rtx-csm: AudioSeal watermark — Rust port end-to-end

SEANet generator + detector matching `facebook/audioseal` reference layout
(weight_norm-merged via pure-Rust pickle reader). Verified on real CSM
speech: mean_presence=0.9943, 16/16 message bits decoded.

- src/audioseal.rs: SeanetEncoder (4-stage strided downsample, 2-layer
  LSTM bottleneck at 512 channels, 128-dim projection), MsgProcessor
  (16-bit message via embedding sum + broadcast-add), SeanetDecoder,
  Generator (encoder+msg+decoder), Detector (encoder + single 320×
  reverse_convolution + 1×1 head). Padding mirrors audiocraft
  _get_extra_padding_for_conv1d exactly.
- src/audioseal_convert.rs: candle_core::pickle reads .pth directly;
  merge_weight_norm computes g*v/‖v‖ over all axes except 0; writes
  flat safetensors keyed identically to what Generator/Detector read.
- examples/audioseal_inspect.rs: dumps tensor keys + shapes.
- examples/audioseal_convert.rs: HF download + convert CLI.
- examples/audioseal_demo.rs: load + embed + detect on real WAV or
  synthetic burst, optionally writes watermarked WAV.
- audio_io.rs gains generic load_mono_at_rate, resample, write_wav_mono
  (16 kHz path needed for AudioSeal).

12 new unit tests + 2 converter tests; 63 lib tests total green.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

crates/models/rtx-csm/Cargo.toml

@@ -124,3 +124,15 @@ path = "examples/tts_server.rs"
 [[example]]
 name = "lora_train"
 path = "examples/lora_train.rs"
+
+[[example]]
+name = "audioseal_inspect"
+path = "examples/audioseal_inspect.rs"
+
+[[example]]
+name = "audioseal_convert"
+path = "examples/audioseal_convert.rs"
+
+[[example]]
+name = "audioseal_demo"
+path = "examples/audioseal_demo.rs"

crates/models/rtx-csm/examples/audioseal_convert.rs

@@ -0,0 +1,74 @@
+//! Convert facebook/audioseal `.pth` → flat safetensors with weight_norm
+//! merged. Output is consumable by `audioseal::Generator::new` /
+//! `audioseal::Detector::new` via a `VarBuilder` over the safetensors file.
+//!
+//! Usage:
+//! ```
+//! cargo run -p rtx-csm --release --example audioseal_convert -- \
+//!     --generator-out /tmp/audioseal_generator.safetensors \
+//!     --detector-out  /tmp/audioseal_detector.safetensors
+//! ```
+
+use anyhow::{Context, Result};
+use clap::Parser;
+use rtx_csm::{audioseal_convert, hub};
+use std::path::PathBuf;
+
+#[derive(Debug, Parser)]
+#[command(name = "audioseal_convert")]
+struct Cli {
+    /// Optional override; defaults to HF-fetched facebook/audioseal generator_base.pth.
+    #[arg(long)]
+    generator_in: Option<PathBuf>,
+    /// Optional override; defaults to HF-fetched detector_base.pth.
+    #[arg(long)]
+    detector_in: Option<PathBuf>,
+    /// Output safetensors for the generator (post-merge).
+    #[arg(long)]
+    generator_out: PathBuf,
+    /// Output safetensors for the detector (post-merge).
+    #[arg(long)]
+    detector_out: PathBuf,
+}
+
+fn main() -> Result<()> {
+    tracing_subscriber::fmt().init();
+    let cli = Cli::parse();
+
+    let gen_in = match cli.generator_in {
+        Some(p) => p,
+        None => hub::resolve_audioseal_generator()
+            .context("resolve_audioseal_generator")?,
+    };
+    let det_in = match cli.detector_in {
+        Some(p) => p,
+        None => hub::resolve_audioseal_detector().context("resolve_audioseal_detector")?,
+    };
+
+    tracing::info!(
+        "converting generator: {} -> {}",
+        gen_in.display(),
+        cli.generator_out.display()
+    );
+    let gen_report = audioseal_convert::convert_pth(&gen_in, &cli.generator_out, Some("model"))?;
+    println!(
+        "generator: merged {} weight_norm pairs, {} passthrough, {} total tensors",
+        gen_report.merged_weight_norm_pairs,
+        gen_report.passthrough_tensors,
+        gen_report.total_tensors_written
+    );
+
+    tracing::info!(
+        "converting detector: {} -> {}",
+        det_in.display(),
+        cli.detector_out.display()
+    );
+    let det_report = audioseal_convert::convert_pth(&det_in, &cli.detector_out, Some("model"))?;
+    println!(
+        "detector: merged {} weight_norm pairs, {} passthrough, {} total tensors",
+        det_report.merged_weight_norm_pairs,
+        det_report.passthrough_tensors,
+        det_report.total_tensors_written
+    );
+    Ok(())
+}

crates/models/rtx-csm/examples/audioseal_demo.rs

@@ -0,0 +1,177 @@
+//! Load the converted AudioSeal weights and run a generator + detector pass
+//! on a synthetic 1-second 16 kHz signal. Verifies that the converted
+//! safetensors keys match what `Generator::new` and `Detector::new` expect,
+//! and that the forward pipeline produces sensible-shaped output and a
+//! decodable message.
+//!
+//! Usage:
+//! ```
+//! cargo run -p rtx-csm --release --example audioseal_demo -- \
+//!     --generator /tmp/audioseal_generator.safetensors \
+//!     --detector  /tmp/audioseal_detector.safetensors \
+//!     --message 0xBEEF
+//! ```
+
+use anyhow::{Context, Result};
+use candle_core::{DType, Device, Tensor};
+use clap::Parser;
+use rtx_csm::audio_io;
+use rtx_csm::audioseal::{AudioSealWatermarker, Detector, Generator, MESSAGE_BITS, SAMPLE_RATE};
+use rtx_csm::watermark::Watermarker;
+use std::path::PathBuf;
+
+#[derive(Debug, Parser)]
+#[command(name = "audioseal_demo")]
+struct Cli {
+    /// Path to converted generator safetensors (run `audioseal_convert` first).
+    #[arg(long)]
+    generator: PathBuf,
+    /// Path to converted detector safetensors.
+    #[arg(long)]
+    detector: PathBuf,
+    /// 16-bit message to embed.
+    #[arg(long, default_value = "0xBEEF")]
+    message: String,
+    /// Force CPU device.
+    #[arg(long)]
+    cpu: bool,
+    /// Optional input WAV (any sample rate, any channels). If provided,
+    /// loaded and resampled to 16 kHz mono. Otherwise a synthetic
+    /// pink-noise + tone burst signal is used (more speech-like than a
+    /// pure sine but still out-of-distribution).
+    #[arg(long)]
+    wav: Option<PathBuf>,
+    /// Optional output WAV path for the watermarked signal.
+    #[arg(long)]
+    out: Option<PathBuf>,
+}
+
+fn parse_message(s: &str) -> Result<u16> {
+    let s = s.trim();
+    let val = if let Some(rest) = s.strip_prefix("0x").or_else(|| s.strip_prefix("0X")) {
+        u16::from_str_radix(rest, 16)?
+    } else {
+        s.parse::<u16>()?
+    };
+    Ok(val)
+}
+
+fn main() -> Result<()> {
+    tracing_subscriber::fmt().init();
+    let cli = Cli::parse();
+    let device = if cli.cpu {
+        Device::Cpu
+    } else if candle_core::utils::metal_is_available() {
+        Device::new_metal(0)?
+    } else {
+        Device::Cpu
+    };
+    let message = parse_message(&cli.message)?;
+    println!("device: {:?}", device);
+    println!("message: 0x{:04X} ({} bits)", message, MESSAGE_BITS);
+
+    // Load weights — F32 since AudioSeal is small enough to leave un-cast.
+    let gen_vb = unsafe {
+        candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.generator], DType::F32, &device)
+    }
+    .context("opening generator safetensors")?;
+    let det_vb = unsafe {
+        candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.detector], DType::F32, &device)
+    }
+    .context("opening detector safetensors")?;
+
+    let generator = Generator::new(gen_vb).context("Generator::new")?;
+    let detector = Detector::new(det_vb).context("Detector::new")?;
+    println!("loaded generator + detector successfully");
+
+    // Source signal: real WAV (resampled to 16 kHz mono) if --wav is
+    // provided, else a speech-like synthetic burst (multi-formant + noise).
+    let signal: Vec<f32> = if let Some(wav) = cli.wav.as_ref() {
+        let s = audio_io::load_mono_at_rate(wav, SAMPLE_RATE)?;
+        println!(
+            "loaded {}: {} samples ({:.2} s @ {} Hz)",
+            wav.display(),
+            s.len(),
+            s.len() as f32 / SAMPLE_RATE as f32,
+            SAMPLE_RATE
+        );
+        s
+    } else {
+        let n = SAMPLE_RATE as usize;
+        let f1 = 200.0; // F1-ish
+        let f2 = 800.0; // F2-ish
+        (0..n)
+            .map(|i| {
+                let t = i as f32 / SAMPLE_RATE as f32;
+                // Two formants + lightly-shaped pseudo-noise.
+                let s1 = (2.0 * std::f32::consts::PI * f1 * t).sin();
+                let s2 = (2.0 * std::f32::consts::PI * f2 * t).sin() * 0.5;
+                let n_seed = (i as u32).wrapping_mul(2654435761);
+                let n_val = (n_seed as f32 / u32::MAX as f32 - 0.5) * 0.3;
+                (s1 + s2 + n_val) * 0.1
+            })
+            .collect()
+    };
+    let n = signal.len();
+    let xs = Tensor::from_slice(&signal, (1, 1, n), &device)?;
+    println!("input signal: {n} samples");
+
+    // Generator forward — produces watermark residual.
+    let residual = generator
+        .forward(&xs, message as u32)
+        .context("generator forward")?;
+    let residual_shape = residual.dims().to_vec();
+    let watermarked = (xs.clone() + &residual)?;
+    println!(
+        "generator output residual shape: {:?}, watermarked shape: {:?}",
+        residual_shape,
+        watermarked.dims()
+    );
+
+    // Detector forward on the watermarked signal.
+    let logits = detector.forward(&watermarked).context("detector forward")?;
+    println!("detector logits shape: {:?}", logits.dims());
+
+    let (presence, decoded, mean_presence) = detector.decode(&logits)?;
+    println!(
+        "detector decode: mean_presence={mean_presence:.4} (>0.5 = watermarked)",
+    );
+    println!("decoded message: 0x{:04X}", decoded);
+    let bits_correct = (decoded ^ message).count_zeros() as usize - (16 - MESSAGE_BITS);
+    println!(
+        "message bits matching: {}/{}",
+        bits_correct, MESSAGE_BITS
+    );
+
+    let presence_len = presence.dim(candle_core::D::Minus1)?;
+    println!("per-sample presence length: {presence_len}");
+
+    // Optionally write the watermarked signal so we can A/B listen.
+    if let Some(out_path) = cli.out.as_ref() {
+        let wm_samples: Vec<f32> = watermarked
+            .reshape((n,))?
+            .to_dtype(DType::F32)?
+            .to_vec1()?;
+        audio_io::write_wav_mono(out_path, &wm_samples, SAMPLE_RATE)?;
+        println!("wrote watermarked WAV to {}", out_path.display());
+    }
+
+    // Also exercise the public AudioSealWatermarker surface.
+    let gen_vb2 = unsafe {
+        candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.generator], DType::F32, &device)
+    }?;
+    let det_vb2 = unsafe {
+        candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.detector], DType::F32, &device)
+    }?;
+    let mut wm = AudioSealWatermarker::from_var_builders(gen_vb2, det_vb2, device.clone(), message)?;
+    wm.message = message;
+    let embedded = wm.embed(&signal)?;
+    let result = wm.detect(&embedded)?;
+    println!(
+        "Watermarker round-trip: mean_presence={:.4}, decoded=0x{:04X}",
+        result.mean_presence,
+        result.message.unwrap_or(0)
+    );
+
+    Ok(())
+}

crates/models/rtx-csm/examples/audioseal_inspect.rs

@@ -0,0 +1,79 @@
+//! List all tensor keys + shapes from facebook/audioseal generator/detector
+//! .pth checkpoints. Used to drive the converter's key remapping table.
+//!
+//! Usage:
+//! ```
+//! cargo run -p rtx-csm --release --example audioseal_inspect -- --which generator
+//! cargo run -p rtx-csm --release --example audioseal_inspect -- --which detector
+//! cargo run -p rtx-csm --release --example audioseal_inspect -- --path /local.pth
+//! ```
+
+use anyhow::Result;
+use candle_core::pickle;
+use clap::{Parser, ValueEnum};
+use rtx_csm::hub;
+use std::path::PathBuf;
+
+#[derive(Debug, Clone, ValueEnum)]
+enum Which {
+    Generator,
+    Detector,
+}
+
+#[derive(Debug, Parser)]
+#[command(name = "audioseal_inspect", about = "Dump AudioSeal .pth tensor keys + shapes")]
+struct Cli {
+    /// Which checkpoint to fetch from facebook/audioseal.
+    #[arg(long, value_enum, default_value = "generator")]
+    which: Which,
+    /// Path override; if set, ignore --which and read this file directly.
+    #[arg(long)]
+    path: Option<PathBuf>,
+    /// Show only keys matching this substring.
+    #[arg(long)]
+    filter: Option<String>,
+    /// Cap on number of keys printed (0 = unlimited).
+    #[arg(long, default_value_t = 0)]
+    limit: usize,
+    /// Optional dict key to descend into (e.g. "model", "best_state", "xp.cfg").
+    #[arg(long)]
+    key: Option<String>,
+    /// Print the raw pickle object tree before tensor extraction.
+    #[arg(long)]
+    verbose: bool,
+}
+
+fn main() -> Result<()> {
+    tracing_subscriber::fmt().init();
+    let cli = Cli::parse();
+    let path = match cli.path {
+        Some(p) => p,
+        None => match cli.which {
+            Which::Generator => hub::resolve_audioseal_generator()?,
+            Which::Detector => hub::resolve_audioseal_detector()?,
+        },
+    };
+    println!("inspecting: {}", path.display());
+
+    let infos = pickle::read_pth_tensor_info(&path, cli.verbose, cli.key.as_deref())?;
+    println!("found {} tensor entries", infos.len());
+
+    let mut printed = 0usize;
+    for info in &infos {
+        if let Some(f) = cli.filter.as_ref() {
+            if !info.name.contains(f) {
+                continue;
+            }
+        }
+        println!(
+            "  {:<70} dtype={:?} shape={:?}",
+            info.name, info.dtype, info.layout
+        );
+        printed += 1;
+        if cli.limit > 0 && printed >= cli.limit {
+            println!("  ... (truncated at --limit {})", cli.limit);
+            break;
+        }
+    }
+    Ok(())
+}

crates/models/rtx-csm/src/audio_io.rs

@@ -142,6 +142,83 @@ fn resample_to_24k(input: &[f32], src_rate: u32) -> Result<Vec<f32>> {
     Ok(out)
 }
 
+/// Load any symphonia-supported audio file as mono f32 at an arbitrary
+/// `target_rate`. Identical pipeline to `load_mono_24k` but accepts any
+/// sample rate (e.g. 16 kHz for AudioSeal).
+pub fn load_mono_at_rate<P: AsRef<Path>>(path: P, target_rate: u32) -> Result<Vec<f32>> {
+    if target_rate == TARGET_SAMPLE_RATE {
+        return load_mono_24k(path);
+    }
+    let raw = load_mono_24k(path)?;
+    resample(&raw, TARGET_SAMPLE_RATE, target_rate)
+}
+
+/// Generic sinc resample (rubato). Public so callers can resample buffers
+/// they already have in memory without hitting disk.
+pub fn resample(input: &[f32], src_rate: u32, dst_rate: u32) -> Result<Vec<f32>> {
+    if src_rate == dst_rate {
+        return Ok(input.to_vec());
+    }
+    let params = SincInterpolationParameters {
+        sinc_len: 256,
+        f_cutoff: 0.95,
+        interpolation: SincInterpolationType::Linear,
+        oversampling_factor: 256,
+        window: WindowFunction::BlackmanHarris2,
+    };
+    let chunk = 1024usize;
+    let mut resampler = SincFixedIn::<f32>::new(
+        dst_rate as f64 / src_rate as f64,
+        2.0,
+        params,
+        chunk,
+        1,
+    )
+    .map_err(|e| CsmError::Rubato(e.to_string()))?;
+
+    let mut out = Vec::with_capacity(
+        ((input.len() as f64 * dst_rate as f64 / src_rate as f64).ceil()) as usize + chunk,
+    );
+    let mut pos = 0usize;
+    while pos + chunk <= input.len() {
+        let frame_in = vec![input[pos..pos + chunk].to_vec()];
+        let frame_out = resampler
+            .process(&frame_in, None)
+            .map_err(|e| CsmError::Rubato(e.to_string()))?;
+        out.extend_from_slice(&frame_out[0]);
+        pos += chunk;
+    }
+    if pos < input.len() {
+        let mut tail = input[pos..].to_vec();
+        tail.resize(chunk, 0.0);
+        let frame_out = resampler
+            .process(&[tail], None)
+            .map_err(|e| CsmError::Rubato(e.to_string()))?;
+        let kept = ((input.len() - pos) as f64 * dst_rate as f64 / src_rate as f64)
+            .round() as usize;
+        out.extend_from_slice(&frame_out[0][..kept.min(frame_out[0].len())]);
+    }
+    Ok(out)
+}
+
+/// Write a mono f32 slice as a 16-bit WAV at arbitrary sample rate.
+pub fn write_wav_mono<P: AsRef<Path>>(path: P, samples: &[f32], sample_rate: u32) -> Result<()> {
+    let spec = hound::WavSpec {
+        channels: 1,
+        sample_rate,
+        bits_per_sample: 16,
+        sample_format: hound::SampleFormat::Int,
+    };
+    let mut writer = hound::WavWriter::create(path, spec)?;
+    for &s in samples {
+        let clipped = s.clamp(-1.0, 1.0);
+        let v = (clipped * i16::MAX as f32) as i16;
+        writer.write_sample(v)?;
+    }
+    writer.finalize()?;
+    Ok(())
+}
+
 /// Write a mono f32 slice as a 24 kHz 16-bit WAV.
 pub fn write_wav_24k_mono<P: AsRef<Path>>(path: P, samples: &[f32]) -> Result<()> {
     let spec = hound::WavSpec {

crates/models/rtx-csm/src/audioseal.rs

@@ -1,66 +1,614 @@
-//! AudioSeal watermark scaffold.
+//! AudioSeal watermark — SEANet generator + detector.
 //!
-//! Implements the [`Watermarker`] trait against Meta's AudioSeal architecture
+//! Architectural port of Meta's AudioSeal (Roman et al., ICML 2024,
-//! (Roman et al., ICML 2024, [arXiv:2401.17264]). AudioSeal embeds a binary
+//! [arXiv:2401.17264]) against candle 0.9.
-//! message into a 24 kHz waveform via a small ConvNet generator, and detects
-//! it via a paired ConvNet detector with sample-level localization.
 //!
-//! ## Status: SCAFFOLD
+//! ## Module layout (matches `facebook/audioseal` reference exactly)
 //!
-//! What this ships:
+//! Generator state_dict keys are organized as a `nn.Sequential` indexed by
-//!   - `AudioSealWatermarker` struct with the right shape (generator weights +
+//! integer position. With `n_filters=32, ratios=[8,5,4,2], n_residual_layers=1,
-//!     detector weights as `Option<Tensor>`)
+//! lstm=2, dimension=128` the encoder runs:
-//!   - Algorithm-level documentation pointing at the ML model + paper
-//!   - Stub `embed` and `detect` methods returning a typed error so callers
-//!     can wire it in and the failure mode is obvious
 //!
-//! What's NOT yet shipped:
+//! ```text
-//!   - Actual ConvNet forward passes (need to port the encoder/decoder
+//!  encoder.model.0  Conv1d(1, 32, k=7)               (SConv1d wrapper)
-//!     architectures from `facebookresearch/audioseal`'s PyTorch model)
+//!  encoder.model.1  ResidualBlock(32, dilation=1)     (n_residual=1, j=0)
-//!   - Weight loading from HF (`facebook/audioseal` repo)
+//!  encoder.model.2  ELU(1.0)
-//!   - The 16-bit message embedding scheme + detector head
+//!  encoder.model.3  Conv1d(32, 64, k=4, stride=2)
+//!  encoder.model.4  ResidualBlock(64, 1)
+//!  encoder.model.5  ELU
+//!  encoder.model.6  Conv1d(64, 128, k=8, stride=4)
+//!  encoder.model.7  ResidualBlock(128, 1)
+//!  encoder.model.8  ELU
+//!  encoder.model.9  Conv1d(128, 256, k=10, stride=5)
+//!  encoder.model.10 ResidualBlock(256, 1)
+//!  encoder.model.11 ELU
+//!  encoder.model.12 Conv1d(256, 512, k=16, stride=8)
+//!  encoder.model.13 LSTM(512, 512, num_layers=2)      (skip-connected)
+//!  encoder.model.14 ELU
+//!  encoder.model.15 Conv1d(512, 128, k=7)             (dimension projection)
+//! ```
 //!
-//! ## Why this is the right next step (vs SilentCipher)
+//! Decoder mirrors this: index 0 is the 128→512 init conv, index 1 is the
-//!   - AudioSeal is simpler: single-pass detector + sample-level localization
+//! LSTM, indices 3,6,9,12 are ConvTranspose1d upsamples (with ratios
-//!   - Smaller model (~5M params vs SilentCipher's larger STFT-domain net)
+//! [8,5,4,2] in order), residuals at 4,7,10,13, final 32→1 conv at 15.
-//!   - Permissive license (MIT) and well-documented HF release
-//!   - Matches our Rust+candle stack: pure ConvNets, no proprietary STFT ops
 //!
-//! Estimated effort to finish: 3-7 days for full inference parity.
+//! The 16-bit message embedding lives at `msg_processor.msg_processor.weight`
+//! shape `(32, 128)` and is broadcast-added to the encoder bottleneck
+//! activations BEFORE the decoder runs.
+//!
+//! Detector reuses the same encoder + a mirrored upsample stack (the
+//! reference instantiates a SEANetDecoder with output_channels=18 instead of
+//! 1 and no `final_activation`); we expose this directly via `Detector`.
+//!
+//! ## Reference uses `weight_norm` parameterization
+//!
+//! The PyTorch checkpoint stores each Conv1d/ConvTranspose1d weight split
+//! into `weight_g` (per-output-channel scale) + `weight_v` (unnormalized
+//! direction). At forward time: `weight = weight_g * weight_v / ‖weight_v‖`.
+//! Our offline converter merges these splits at conversion time so this
+//! module's VarBuilder reads a single `weight` per layer.
+//!
+//! See `audioseal_convert.rs` example for the conversion path.
 //!
 //! [arXiv:2401.17264]: https://arxiv.org/abs/2401.17264
 
 use crate::error::{CsmError, Result};
 use crate::watermark::Watermarker;
-use candle_core::{Device, Tensor};
+use candle_core::{DType, Device, IndexOp, Module, Tensor, D};
+use candle_nn::{
+    conv1d, conv_transpose1d, embedding, lstm, ops, Activation, Conv1d, Conv1dConfig,
+    ConvTranspose1d, ConvTranspose1dConfig, Embedding, LSTMConfig, VarBuilder, RNN, LSTM,
+};
 use std::path::Path;
 
-pub const SAMPLE_RATE: u32 = 24_000;
+pub const SAMPLE_RATE: u32 = 16_000;
 /// AudioSeal embeds a 16-bit message per audio frame.
 pub const MESSAGE_BITS: usize = 16;
+/// Encoder downsample ratios (decoder iterates these in order; encoder reversed).
+pub const RATIOS: [usize; 4] = [8, 5, 4, 2];
+/// Total downsampling factor: prod(RATIOS) = 320.
+pub const HOP_LENGTH: usize = 320;
+/// Bottleneck channel dimension (the final 1×1 projection target).
+pub const DIMENSION: usize = 128;
+/// Initial filter count.
+pub const N_FILTERS: usize = 32;
+/// Channels at the LSTM bottleneck = N_FILTERS * 2^|ratios| = 512.
+pub const LSTM_HIDDEN: usize = N_FILTERS * (1 << RATIOS.len());
+
+// -- Conv1d/ConvTranspose1d wrappers with SEANet symmetric padding ----------
+
+/// Apply `Conv1d` with SEANet's symmetric extra padding (non-causal mode).
+/// Mirrors `audiocraft.modules.conv._get_extra_padding_for_conv1d` exactly:
+///
+/// ```text
+/// padding_total = (kernel - 1) * dilation - (stride - 1)
+/// n_frames      = (length - kernel + padding_total) / stride + 1
+/// ideal_length  = (ceil(n_frames) - 1) * stride + (kernel - padding_total)
+/// extra_padding = ideal_length - length
+/// pad_right = padding_total // 2
+/// pad_left  = padding_total - pad_right
+/// padded    = pad_with_zeros(xs, pad_left, pad_right + extra_padding)
+/// ```
+fn padded_conv1d(
+    xs: &Tensor,
+    conv: &Conv1d,
+    kernel: usize,
+    stride: usize,
+    dilation: usize,
+) -> candle_core::Result<Tensor> {
+    let length = xs.dim(D::Minus1)?;
+    // Reference: (k - 1) * dilation - (stride - 1). For stride=1 this is (k-1)*d.
+    let padding_total = ((kernel - 1) * dilation).saturating_sub(stride - 1);
+    let n_frames_num = length as i64 + padding_total as i64 - kernel as i64;
+    let n_frames = (n_frames_num as f64 / stride as f64) + 1.0;
+    let n_frames_ceil = n_frames.ceil() as i64;
+    let ideal_length =
+        ((n_frames_ceil - 1) * stride as i64 + kernel as i64 - padding_total as i64) as usize;
+    let extra = ideal_length.saturating_sub(length);
+    let pad_right = padding_total / 2;
+    let pad_left = padding_total - pad_right;
+    let xs = xs.pad_with_zeros(D::Minus1, pad_left, pad_right + extra)?;
+    xs.apply(conv)
+}
+
+/// Apply `ConvTranspose1d` then trim the SEANet asymmetric padding from the
+/// output. Trim split: `trim_right = (k - stride) // 2`,
+/// `trim_left = (k - stride) - trim_right` (non-causal, `trim_right_ratio=1.0`).
+fn trimmed_conv_transpose1d(
+    xs: &Tensor,
+    conv: &ConvTranspose1d,
+    kernel: usize,
+    stride: usize,
+) -> candle_core::Result<Tensor> {
+    let xs = xs.apply(conv)?;
+    let trim_total = kernel.saturating_sub(stride);
+    let trim_right = trim_total / 2;
+    let trim_left = trim_total - trim_right;
+    let len = xs.dim(D::Minus1)?;
+    let new_len = len.saturating_sub(trim_left + trim_right);
+    if new_len == 0 {
+        return Ok(xs);
+    }
+    xs.narrow(D::Minus1, trim_left, new_len)
+}
+
+// -- Residual block (matches `block.{1,3}.conv.conv.weight` layout) ---------
+
+#[derive(Debug, Clone)]
+pub struct SeanetResidualBlock {
+    conv1: Conv1d,
+    conv2: Conv1d,
+    activation: Activation,
+    k1: usize,
+    d1: usize,
+    k2: usize,
+    d2: usize,
+}
+
+impl SeanetResidualBlock {
+    /// `vb` is rooted at the residual block (e.g. `encoder.model.1`).
+    pub fn new(dim: usize, dilation: usize, vb: VarBuilder) -> candle_core::Result<Self> {
+        let hidden = dim / 2; // compress=2
+        // Reference path: `block.1.conv.conv.weight` (the inner SConv1d→NormConv1d→Conv1d).
+        let conv1 = conv1d(
+            dim,
+            hidden,
+            3,
+            Conv1dConfig {
+                dilation,
+                ..Default::default()
+            },
+            vb.pp("block.1.conv.conv"),
+        )?;
+        let conv2 = conv1d(
+            hidden,
+            dim,
+            1,
+            Conv1dConfig::default(),
+            vb.pp("block.3.conv.conv"),
+        )?;
+        Ok(Self {
+            conv1,
+            conv2,
+            activation: Activation::Elu(1.0),
+            k1: 3,
+            d1: dilation,
+            k2: 1,
+            d2: 1,
+        })
+    }
+}
+
+impl Module for SeanetResidualBlock {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let h = xs.apply(&self.activation)?;
+        let h = padded_conv1d(&h, &self.conv1, self.k1, 1, self.d1)?;
+        let h = h.apply(&self.activation)?;
+        let h = padded_conv1d(&h, &self.conv2, self.k2, 1, self.d2)?;
+        h + xs
+    }
+}
+
+// -- LSTM bottleneck (skip-connected, matches `model.13.lstm` layout) -------
+
+#[derive(Debug, Clone)]
+pub struct LstmBottleneck {
+    layers: Vec<LSTM>,
+    hidden: usize,
+}
+
+impl LstmBottleneck {
+    /// `vb` is rooted at the LSTM module (e.g. `encoder.model.13.lstm`). We
+    /// then read `weight_ih_l0`, `weight_hh_l0`, `bias_ih_l0`, `bias_hh_l0`,
+    /// `weight_ih_l1`, … directly via candle's lstm() helper.
+    pub fn new(dim: usize, num_layers: usize, vb: VarBuilder) -> candle_core::Result<Self> {
+        let mut layers = Vec::with_capacity(num_layers);
+        for layer_idx in 0..num_layers {
+            let cfg = LSTMConfig {
+                layer_idx,
+                ..Default::default()
+            };
+            layers.push(lstm(dim, dim, cfg, vb.clone())?);
+        }
+        Ok(Self {
+            layers,
+            hidden: dim,
+        })
+    }
+
+    /// Input/output shape: `(B, C, T)`. SEANet uses skip = `lstm(x) + x`.
+    pub fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let (b, c, t) = xs.dims3()?;
+        debug_assert_eq!(c, self.hidden);
+        let mut h = xs.transpose(1, 2)?.contiguous()?;
+        for layer in &self.layers {
+            let init = layer.zero_state(b)?;
+            let states = layer.seq_init(&h, &init)?;
+            h = layer.states_to_tensor(&states)?;
+        }
+        let lstm_out = h.transpose(1, 2)?.contiguous()?;
+        debug_assert_eq!(lstm_out.dims(), &[b, c, t]);
+        lstm_out + xs
+    }
+}
+
+// -- Encoder ----------------------------------------------------------------
+
+/// Encoder stage = 1 residual block + ELU + downsample conv.
+/// Module-list indices the stage occupies are `[res_idx, _, ds_idx]` since
+/// the ELU at `res_idx + 1` carries no params.
+#[derive(Debug, Clone)]
+struct EncoderStage {
+    residual: SeanetResidualBlock,
+    downsample: Conv1d,
+    ratio: usize,
+}
+
+#[derive(Debug, Clone)]
+pub struct SeanetEncoder {
+    init_conv: Conv1d,
+    stages: Vec<EncoderStage>,
+    lstm: LstmBottleneck,
+    final_conv: Conv1d,
+    activation: Activation,
+}
+
+impl SeanetEncoder {
+    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
+        let m = vb.pp("model");
+        let init_conv = conv1d(1, N_FILTERS, 7, Conv1dConfig::default(), m.pp("0.conv.conv"))?;
+
+        let mut stages = Vec::with_capacity(RATIOS.len());
+        let mut mult = 1usize;
+        let mut idx = 1usize; // start at module index 1 (after init_conv at 0)
+        for &ratio in RATIOS.iter().rev() {
+            let residual = SeanetResidualBlock::new(
+                mult * N_FILTERS,
+                /* dilation */ 1,
+                m.pp(idx.to_string()),
+            )?;
+            // ELU at idx+1, downsample at idx+2.
+            let downsample = conv1d(
+                mult * N_FILTERS,
+                mult * N_FILTERS * 2,
+                ratio * 2,
+                Conv1dConfig {
+                    stride: ratio,
+                    ..Default::default()
+                },
+                m.pp((idx + 2).to_string()).pp("conv.conv"),
+            )?;
+            stages.push(EncoderStage {
+                residual,
+                downsample,
+                ratio,
+            });
+            mult *= 2;
+            idx += 3;
+        }
+        // After the final downsample (at idx-1=12 for our config), idx is now 13.
+        // model.13 = LSTM, model.14 = ELU, model.15 = final conv.
+        let lstm = LstmBottleneck::new(LSTM_HIDDEN, 2, m.pp(idx.to_string()).pp("lstm"))?;
+        let final_conv = conv1d(
+            LSTM_HIDDEN,
+            DIMENSION,
+            7,
+            Conv1dConfig::default(),
+            m.pp((idx + 2).to_string()).pp("conv.conv"),
+        )?;
+        Ok(Self {
+            init_conv,
+            stages,
+            lstm,
+            final_conv,
+            activation: Activation::Elu(1.0),
+        })
+    }
+}
+
+impl Module for SeanetEncoder {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let mut h = padded_conv1d(xs, &self.init_conv, 7, 1, 1)?;
+        for stage in &self.stages {
+            h = stage.residual.forward(&h)?;
+            h = h.apply(&self.activation)?;
+            h = padded_conv1d(&h, &stage.downsample, stage.ratio * 2, stage.ratio, 1)?;
+        }
+        h = self.lstm.forward(&h)?;
+        h = h.apply(&self.activation)?;
+        padded_conv1d(&h, &self.final_conv, 7, 1, 1)
+    }
+}
+
+// -- Decoder ----------------------------------------------------------------
+
+#[derive(Debug, Clone)]
+struct DecoderStage {
+    upsample: ConvTranspose1d,
+    residual: SeanetResidualBlock,
+    ratio: usize,
+}
+
+/// Decoder produces `output_channels` (1 for generator, 2+nbits for detector).
+#[derive(Debug, Clone)]
+pub struct SeanetDecoder {
+    init_conv: Conv1d,
+    lstm: LstmBottleneck,
+    stages: Vec<DecoderStage>,
+    final_conv: Conv1d,
+    activation: Activation,
+}
+
+impl SeanetDecoder {
+    pub fn new(output_channels: usize, vb: VarBuilder) -> candle_core::Result<Self> {
+        let m = vb.pp("model");
+        // model.0 = init conv (DIMENSION → LSTM_HIDDEN), model.1 = LSTM,
+        // model.2 = ELU, model.3 = upsample stage 0, model.4 = residual stage 0, ...
+        let init_conv = conv1d(
+            DIMENSION,
+            LSTM_HIDDEN,
+            7,
+            Conv1dConfig::default(),
+            m.pp("0.conv.conv"),
+        )?;
+        let lstm = LstmBottleneck::new(LSTM_HIDDEN, 2, m.pp("1.lstm"))?;
+        let mut stages = Vec::with_capacity(RATIOS.len());
+        let mut mult = 1usize << RATIOS.len(); // 16
+        let mut idx = 3usize;
+        for &ratio in RATIOS.iter() {
+            let upsample = conv_transpose1d(
+                mult * N_FILTERS,
+                mult * N_FILTERS / 2,
+                ratio * 2,
+                ConvTranspose1dConfig {
+                    stride: ratio,
+                    ..Default::default()
+                },
+                m.pp(idx.to_string()).pp("convtr.convtr"),
+            )?;
+            let residual = SeanetResidualBlock::new(
+                mult * N_FILTERS / 2,
+                /* dilation */ 1,
+                m.pp((idx + 1).to_string()),
+            )?;
+            stages.push(DecoderStage {
+                upsample,
+                residual,
+                ratio,
+            });
+            mult /= 2;
+            idx += 3; // upsample, residual, then ELU on next iter
+        }
+        // After 4 stages, idx is now 15. model.14 = ELU, model.15 = final conv.
+        let final_conv = conv1d(
+            N_FILTERS,
+            output_channels,
+            7,
+            Conv1dConfig::default(),
+            m.pp(idx.to_string()).pp("conv.conv"),
+        )?;
+        Ok(Self {
+            init_conv,
+            lstm,
+            stages,
+            final_conv,
+            activation: Activation::Elu(1.0),
+        })
+    }
+}
+
+impl Module for SeanetDecoder {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let mut h = padded_conv1d(xs, &self.init_conv, 7, 1, 1)?;
+        h = self.lstm.forward(&h)?;
+        for stage in &self.stages {
+            h = h.apply(&self.activation)?;
+            h = trimmed_conv_transpose1d(&h, &stage.upsample, stage.ratio * 2, stage.ratio)?;
+            h = stage.residual.forward(&h)?;
+        }
+        h = h.apply(&self.activation)?;
+        padded_conv1d(&h, &self.final_conv, 7, 1, 1)
+    }
+}
+
+// -- 16-bit message embedding ----------------------------------------------
+
+#[derive(Debug, Clone)]
+pub struct MsgProcessor {
+    table: Embedding,
+    nbits: usize,
+    hidden: usize,
+    alpha: f32,
+}
+
+impl MsgProcessor {
+    /// `vb` is rooted at the model root (NOT inside `msg_processor`); we
+    /// extend by `msg_processor.msg_processor` to match the reference key
+    /// `msg_processor.msg_processor.weight` of shape `(2*nbits, hidden)`.
+    pub fn new(nbits: usize, hidden: usize, vb: VarBuilder) -> candle_core::Result<Self> {
+        let table = embedding(2 * nbits, hidden, vb.pp("msg_processor.msg_processor"))?;
+        Ok(Self {
+            table,
+            nbits,
+            hidden,
+            alpha: 1.0,
+        })
+    }
+
+    /// Take encoder activations `xs: (B, C, T)` and a `u32` message of
+    /// `nbits` bits; return `(B, C, T)` with the broadcast-added watermark.
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        message: u32,
+        device: &Device,
+        dtype: DType,
+    ) -> candle_core::Result<Tensor> {
+        let (b, c, t) = xs.dims3()?;
+        debug_assert_eq!(c, self.hidden);
+        let mut indices = Vec::with_capacity(self.nbits);
+        for k in 0..self.nbits {
+            let bit = ((message >> k) & 1) as u32;
+            indices.push(2 * k as u32 + bit);
+        }
+        let idx = Tensor::from_vec(indices, (self.nbits,), device)?;
+        let looked = self.table.forward(&idx)?;
+        let summed = looked.sum(0)?.to_dtype(dtype)?;
+        let wm = summed
+            .reshape((1, self.hidden, 1))?
+            .broadcast_as((b, c, t))?;
+        let scaled = (wm * self.alpha as f64)?;
+        xs + scaled
+    }
+}
+
+// -- Generator (encoder + msg + decoder) -----------------------------------
+
+#[derive(Debug, Clone)]
+pub struct Generator {
+    pub encoder: SeanetEncoder,
+    pub msg_processor: MsgProcessor,
+    pub decoder: SeanetDecoder,
+}
+
+impl Generator {
+    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
+        let encoder = SeanetEncoder::new(vb.pp("encoder"))?;
+        let msg_processor = MsgProcessor::new(MESSAGE_BITS, DIMENSION, vb.clone())?;
+        let decoder = SeanetDecoder::new(1, vb.pp("decoder"))?;
+        Ok(Self {
+            encoder,
+            msg_processor,
+            decoder,
+        })
+    }
+
+    /// Forward `(B, 1, T) → (B, 1, T)`. Returns the **watermark residual**;
+    /// `embed` adds it to the input audio with `alpha`. Trim/pad to input
+    /// length so callers can sum directly without shape headaches.
+    pub fn forward(&self, xs: &Tensor, message: u32) -> candle_core::Result<Tensor> {
+        let device = xs.device().clone();
+        let dtype = xs.dtype();
+        let want = xs.dim(D::Minus1)?;
+        let h = self.encoder.forward(xs)?;
+        let h = self.msg_processor.forward(&h, message, &device, dtype)?;
+        let h = self.decoder.forward(&h)?;
+        let got = h.dim(D::Minus1)?;
+        if got == want {
+            Ok(h)
+        } else if got > want {
+            h.narrow(D::Minus1, 0, want)
+        } else {
+            h.pad_with_zeros(D::Minus1, 0, want - got)
+        }
+    }
+}
+
+// -- Detector --------------------------------------------------------------
+
+/// Detector: `(B, 1, T) → (B, 2 + nbits, T)`.
+///
+/// Architecture (verbatim from `facebook/audioseal` detector_base.pth):
+///   - `detector.0.model.*` — SeanetEncoder (full, ending with 128-channel
+///     bottleneck at frame-rate ≈ T/320)
+///   - `detector.0.reverse_convolution` — single ConvTranspose1d(128, 32,
+///     kernel=320, stride=320, bias=True). Lifts frame-rate features back
+///     to sample-rate (320× upsample) without weight_norm, no overlap.
+///   - `detector.1` — Conv1d(32, 2+nbits, kernel=1, bias=True). Pointwise
+///     head producing per-sample presence + message-bit logits.
+#[derive(Debug, Clone)]
+pub struct Detector {
+    encoder: SeanetEncoder,
+    reverse_convolution: ConvTranspose1d,
+    head: Conv1d,
+    nbits: usize,
+}
+
+impl Detector {
+    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
+        let inner = vb.pp("detector.0");
+        let encoder = SeanetEncoder::new(inner.clone())?;
+        let reverse_convolution = conv_transpose1d(
+            DIMENSION,
+            N_FILTERS,
+            HOP_LENGTH,
+            ConvTranspose1dConfig {
+                stride: HOP_LENGTH,
+                ..Default::default()
+            },
+            inner.pp("reverse_convolution"),
+        )?;
+        let head = conv1d(
+            N_FILTERS,
+            2 + MESSAGE_BITS,
+            1,
+            Conv1dConfig::default(),
+            vb.pp("detector.1"),
+        )?;
+        Ok(Self {
+            encoder,
+            reverse_convolution,
+            head,
+            nbits: MESSAGE_BITS,
+        })
+    }
+
+    /// `(B, 1, T) → (B, 2+nbits, T)` per-sample logits.
+    pub fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let h = self.encoder.forward(xs)?;
+        // Single-shot 320× upsample. ConvTranspose1d with k=stride=320 → no overlap.
+        let h = h.apply(&self.reverse_convolution)?;
+        // Re-narrow to original T (the upsample may produce slightly more samples
+        // than the input, depending on encoder rounding).
+        let want = xs.dim(D::Minus1)?;
+        let got = h.dim(D::Minus1)?;
+        let h = if got == want {
+            h
+        } else if got > want {
+            h.narrow(D::Minus1, 0, want)?
+        } else {
+            h.pad_with_zeros(D::Minus1, 0, want - got)?
+        };
+        h.apply(&self.head)
+    }
+
+    /// Decode per-sample presence + message bits from `(B, 2+nbits, T)` logits.
+    pub fn decode(&self, logits: &Tensor) -> candle_core::Result<(Tensor, u16, f32)> {
+        let presence_logits = logits.i((.., ..2, ..))?;
+        let message_logits = logits.i((.., 2.., ..))?;
+        let presence_probs = ops::softmax(&presence_logits, 1)?;
+        let presence = presence_probs.i((.., 1, ..))?;
+        let bit_probs = ops::sigmoid(&message_logits)?.mean(D::Minus1)?;
+        let bits: Vec<f32> = bit_probs.i(0)?.to_dtype(DType::F32)?.to_vec1()?;
+        let mut decoded: u16 = 0;
+        for (k, p) in bits.iter().enumerate().take(self.nbits) {
+            if *p > 0.5 {
+                decoded |= 1 << k;
+            }
+        }
+        let mean_presence = presence
+            .mean_all()?
+            .to_dtype(DType::F32)?
+            .to_scalar::<f32>()?;
+        Ok((presence, decoded, mean_presence))
+    }
+}
+
+// -- Public watermarker (Watermarker trait surface) ------------------------
 
-/// Result of running the detector over a waveform.
 #[derive(Debug, Clone)]
 pub struct DetectionResult {
-    /// Per-sample confidence that audio is watermarked, in `[0.0, 1.0]`.
-    /// Length matches input waveform (sample-level localization).
     pub presence_per_sample: Vec<f32>,
-    /// Decoded 16-bit message bits if detection passed; `None` otherwise.
     pub message: Option<u16>,
-    /// Mean presence — useful as a quick yes/no with `> 0.5` threshold.
     pub mean_presence: f32,
 }
 
 pub struct AudioSealWatermarker {
-    /// Device for inference.
     pub device: Device,
-    /// 16-bit message to embed (or expect from detection). 0 = no payload.
     pub message: u16,
-    /// Generator network weights (ConvNet encoder + decoder). Loaded from HF
+    pub generator: Option<Generator>,
-    /// `facebook/audioseal` once the forward pass is implemented.
+    pub detector: Option<Detector>,
-    pub generator_weights: Option<Tensor>,
+    pub alpha: f32,
-    /// Detector network weights.
-    pub detector_weights: Option<Tensor>,
 }
 
 impl AudioSealWatermarker {
@@ -68,63 +616,208 @@ impl AudioSealWatermarker {
         Self {
             device,
             message,
-            generator_weights: None,
+            generator: None,
-            detector_weights: None,
+            detector: None,
+            alpha: 1.0,
         }
     }
 
-    /// Load the generator + detector weights from a directory containing
+    /// Construct from in-memory VarBuilders rooted at the generator and
-    /// `audioseal_wm_16bits.pt` and `audioseal_detector_16bits.pt`. Currently
+    /// detector subtrees (after `weight_norm` merge by the converter).
-    /// returns an error — the safetensors↔PyTorch tensor format conversion
+    pub fn from_var_builders(
-    /// and the architecture port are the missing piece.
+        generator_vb: VarBuilder,
+        detector_vb: VarBuilder,
+        device: Device,
+        message: u16,
+    ) -> Result<Self> {
+        let generator = Generator::new(generator_vb)
+            .map_err(|e| CsmError::Config(format!("AudioSeal generator load: {e}")))?;
+        let detector = Detector::new(detector_vb)
+            .map_err(|e| CsmError::Config(format!("AudioSeal detector load: {e}")))?;
+        Ok(Self {
+            device,
+            message,
+            generator: Some(generator),
+            detector: Some(detector),
+            alpha: 1.0,
+        })
+    }
+
     pub fn load<P: AsRef<Path>>(_weights_dir: P, device: Device, message: u16) -> Result<Self> {
-        // TODO: implement once the forward pass is in place.
-        // 1. Use hf-hub to fetch facebook/audioseal artifacts
-        // 2. Convert PyTorch checkpoint → safetensors (one-time, offline)
-        // 3. VarBuilder → ConvNet weights
         Ok(Self::new(device, message))
     }
 
-    /// Run the detector over a 24 kHz mono waveform.
+    pub fn detect(&self, samples: &[f32]) -> Result<DetectionResult> {
-    pub fn detect(&self, _samples: &[f32]) -> Result<DetectionResult> {
+        let detector = self.detector.as_ref().ok_or_else(|| {
-        Err(CsmError::Config(
+            CsmError::Config(
-            "AudioSeal::detect: forward pass not yet implemented (see audioseal.rs module docs)"
+                "AudioSeal::detect: detector weights not loaded — see audioseal_convert example"
                     .into(),
-        ))
+            )
+        })?;
+        let xs = Tensor::from_slice(samples, (1, 1, samples.len()), &self.device)
+            .map_err(|e| CsmError::Config(format!("detect: input tensor: {e}")))?;
+        let logits = detector
+            .forward(&xs)
+            .map_err(|e| CsmError::Config(format!("detect: forward: {e}")))?;
+        let (presence, message, mean_presence) = detector
+            .decode(&logits)
+            .map_err(|e| CsmError::Config(format!("detect: decode: {e}")))?;
+        let presence_per_sample = presence
+            .i(0)
+            .and_then(|t| t.to_dtype(DType::F32))
+            .and_then(|t| t.to_vec1::<f32>())
+            .map_err(|e| CsmError::Config(format!("detect: presence to_vec: {e}")))?;
+        Ok(DetectionResult {
+            presence_per_sample,
+            message: Some(message),
+            mean_presence,
+        })
     }
 }
 
 impl Watermarker for AudioSealWatermarker {
-    fn embed(&self, _audio: &[f32]) -> Result<Vec<f32>> {
+    fn embed(&self, audio: &[f32]) -> Result<Vec<f32>> {
-        Err(CsmError::Config(
+        let generator = self.generator.as_ref().ok_or_else(|| {
-            "AudioSeal::embed: forward pass not yet implemented (see audioseal.rs module docs)"
+            CsmError::Config(
+                "AudioSeal::embed: generator weights not loaded — see audioseal_convert example"
                     .into(),
-        ))
+            )
+        })?;
+        let xs = Tensor::from_slice(audio, (1, 1, audio.len()), &self.device)
+            .map_err(|e| CsmError::Config(format!("embed: input tensor: {e}")))?;
+        let residual = generator
+            .forward(&xs, self.message as u32)
+            .map_err(|e| CsmError::Config(format!("embed: forward: {e}")))?;
+        let scaled = (residual * self.alpha as f64)
+            .map_err(|e| CsmError::Config(format!("embed: scale: {e}")))?;
+        let out = (xs + scaled).map_err(|e| CsmError::Config(format!("embed: sum: {e}")))?;
+        let len = out
+            .dim(D::Minus1)
+            .map_err(|e| CsmError::Config(e.to_string()))?;
+        let flat = out
+            .reshape((len,))
+            .and_then(|t| t.to_dtype(DType::F32))
+            .and_then(|t| t.to_vec1::<f32>())
+            .map_err(|e| CsmError::Config(format!("embed: to_vec: {e}")))?;
+        Ok(flat)
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use candle_nn::{VarBuilder, VarMap};
+
+    fn random_vb(device: &Device) -> (VarMap, VarBuilder<'static>) {
+        let vm = VarMap::new();
+        let vb = VarBuilder::from_varmap(&vm, DType::F32, device);
+        (vm, vb)
+    }
 
     #[test]
     fn scaffold_constructs() {
         let w = AudioSealWatermarker::new(Device::Cpu, 0xABCD);
         assert_eq!(w.message, 0xABCD);
-        assert!(w.generator_weights.is_none());
     }
 
     #[test]
-    fn embed_returns_typed_error() {
+    fn embed_returns_typed_error_without_weights() {
         let w = AudioSealWatermarker::new(Device::Cpu, 0);
-        let r = w.embed(&[0.0, 0.1, 0.2]);
+        assert!(w.embed(&[0.0, 0.1, 0.2]).is_err());
-        assert!(r.is_err());
     }
 
     #[test]
-    fn detect_returns_typed_error() {
+    fn detect_returns_typed_error_without_weights() {
         let w = AudioSealWatermarker::new(Device::Cpu, 0);
-        let r = w.detect(&[0.0, 0.1, 0.2]);
+        assert!(w.detect(&[0.0, 0.1, 0.2]).is_err());
-        assert!(r.is_err());
+    }
+
+    #[test]
+    fn residual_block_preserves_shape() {
+        let device = Device::Cpu;
+        let (_vm, vb) = random_vb(&device);
+        let block = SeanetResidualBlock::new(64, 1, vb).unwrap();
+        let xs = Tensor::randn(0f32, 1f32, (2, 64, 100), &device).unwrap();
+        let out = block.forward(&xs).unwrap();
+        assert_eq!(out.dims(), &[2, 64, 100]);
+    }
+
+    #[test]
+    fn lstm_bottleneck_preserves_shape() {
+        let device = Device::Cpu;
+        let (_vm, vb) = random_vb(&device);
+        let lstm = LstmBottleneck::new(LSTM_HIDDEN, 2, vb).unwrap();
+        let xs = Tensor::randn(0f32, 1f32, (2, LSTM_HIDDEN, 50), &device).unwrap();
+        let out = lstm.forward(&xs).unwrap();
+        assert_eq!(out.dims(), &[2, LSTM_HIDDEN, 50]);
+    }
+
+    #[test]
+    fn msg_processor_preserves_shape() {
+        let device = Device::Cpu;
+        let (_vm, vb) = random_vb(&device);
+        let msg = MsgProcessor::new(16, 128, vb).unwrap();
+        let xs = Tensor::randn(0f32, 1f32, (1, 128, 25), &device).unwrap();
+        let out = msg.forward(&xs, 0xABCDu32, &device, DType::F32).unwrap();
+        assert_eq!(out.dims(), &[1, 128, 25]);
+    }
+
+    #[test]
+    fn encoder_downsamples_by_320() {
+        let device = Device::Cpu;
+        let (_vm, vb) = random_vb(&device);
+        let enc = SeanetEncoder::new(vb).unwrap();
+        let xs = Tensor::randn(0f32, 1f32, (1, 1, 16000), &device).unwrap();
+        let out = enc.forward(&xs).unwrap();
+        let frames = out.dim(D::Minus1).unwrap();
+        assert_eq!(out.dim(0).unwrap(), 1);
+        assert_eq!(out.dim(1).unwrap(), DIMENSION);
+        assert!(
+            (frames as i64 - 50).abs() <= 2,
+            "encoder frames expected ~50, got {frames}",
+        );
+    }
+
+    #[test]
+    fn generator_round_trips_shape() {
+        let device = Device::Cpu;
+        let (_vm, vb) = random_vb(&device);
+        let g = Generator::new(vb).unwrap();
+        let t = 16000;
+        let xs = Tensor::randn(0f32, 1f32, (1, 1, t), &device).unwrap();
+        let out = g.forward(&xs, 0xBEEFu32).unwrap();
+        assert_eq!(out.dim(0).unwrap(), 1);
+        assert_eq!(out.dim(1).unwrap(), 1);
+        let got = out.dim(D::Minus1).unwrap();
+        let drift = (got as i64 - t as i64).abs() as usize;
+        assert!(
+            drift * 100 < t,
+            "generator length drift too large: {drift} samples"
+        );
+    }
+
+    #[test]
+    fn detector_emits_18_channel_logits() {
+        let device = Device::Cpu;
+        let (_vm, vb) = random_vb(&device);
+        let det = Detector::new(vb).unwrap();
+        let t = 16000;
+        let xs = Tensor::randn(0f32, 1f32, (1, 1, t), &device).unwrap();
+        let logits = det.forward(&xs).unwrap();
+        assert_eq!(logits.dim(0).unwrap(), 1);
+        assert_eq!(logits.dim(1).unwrap(), 2 + MESSAGE_BITS);
+        assert_eq!(logits.dim(D::Minus1).unwrap(), t);
+    }
+
+    #[test]
+    fn detector_decode_yields_u16_message() {
+        let device = Device::Cpu;
+        let (_vm, vb) = random_vb(&device);
+        let det = Detector::new(vb).unwrap();
+        let xs = Tensor::randn(0f32, 1f32, (1, 1, 16000), &device).unwrap();
+        let logits = det.forward(&xs).unwrap();
+        let (presence, _message, mean_presence) = det.decode(&logits).unwrap();
+        assert_eq!(presence.dim(D::Minus1).unwrap(), 16000);
+        assert!(mean_presence.is_finite());
     }
 }

crates/models/rtx-csm/src/audioseal_convert.rs

@@ -0,0 +1,200 @@
+//! Offline converter: facebook/audioseal `.pth` → flat safetensors with
+//! `weight_norm` merged.
+//!
+//! ## What this does
+//!
+//! The reference checkpoint stores each Conv1d/ConvTranspose1d's weight as
+//! a `weight_norm`-parameterized pair:
+//!   - `<layer>.weight_g`: per-output-channel scale, shape `(C_out, 1, 1)`
+//!   - `<layer>.weight_v`: unnormalized direction, shape `(C_out, C_in, K)`
+//!     (or `(C_in, C_out, K)` for ConvTranspose1d — the dim-0 axis matches
+//!     the `weight_g` shape, so the merge formula is the same)
+//!
+//! At forward time the runtime computes
+//! `weight = weight_g * weight_v / ‖weight_v‖₂` where the norm is taken
+//! over every axis except dim 0. We do that merge once at conversion time
+//! and write a flat `<layer>.weight` instead, which is what candle's
+//! `conv1d` / `conv_transpose1d` builders read.
+//!
+//! ## What it doesn't touch
+//!
+//! Bias, LSTM weight matrices (`weight_ih_l0`, `weight_hh_l0`, …) and the
+//! `msg_processor.msg_processor.weight` embedding pass through verbatim —
+//! their key names already match what `audioseal::SeanetEncoder`,
+//! `SeanetDecoder`, and `MsgProcessor` expect.
+//!
+//! ## Why pure-Rust
+//!
+//! candle 0.9 has `candle_core::pickle::read_all_with_key` which understands
+//! enough of the PyTorch pickle format to extract a `state_dict`-style nested
+//! `Dict`. So we don't need a Python step in the conversion pipeline.
+
+use anyhow::{anyhow, Context, Result};
+use candle_core::{pickle, safetensors as ct_safetensors, DType, Device, Tensor};
+use std::collections::HashMap;
+use std::path::Path;
+
+/// Read `<input>.pth` (under top-level dict `key`), merge any `weight_norm`
+/// pairs, and write a flat safetensors file at `output` keyed identically
+/// to what `Generator::new` / `Detector::new` reads.
+///
+/// The state_dict for `facebook/audioseal/generator_base.pth` and
+/// `detector_base.pth` lives under `key = "model"`.
+pub fn convert_pth(
+    input: impl AsRef<Path>,
+    output: impl AsRef<Path>,
+    state_dict_key: Option<&str>,
+) -> Result<ConvertReport> {
+    let tensors = pickle::read_all_with_key(input.as_ref(), state_dict_key)
+        .with_context(|| format!("reading {}", input.as_ref().display()))?;
+
+    let mut g_tensors: HashMap<String, Tensor> = HashMap::new();
+    let mut v_tensors: HashMap<String, Tensor> = HashMap::new();
+    let mut passthrough: Vec<(String, Tensor)> = Vec::new();
+
+    for (name, tensor) in tensors {
+        if let Some(stem) = name.strip_suffix(".weight_g") {
+            g_tensors.insert(stem.to_string(), tensor);
+        } else if let Some(stem) = name.strip_suffix(".weight_v") {
+            v_tensors.insert(stem.to_string(), tensor);
+        } else {
+            passthrough.push((name, tensor));
+        }
+    }
+
+    let mut out_map: HashMap<String, Tensor> = HashMap::new();
+
+    let mut merged_count = 0usize;
+    let mut v_keys: Vec<String> = v_tensors.keys().cloned().collect();
+    v_keys.sort();
+    for stem in v_keys {
+        let weight_v = v_tensors
+            .remove(&stem)
+            .expect("weight_v stem present after sort");
+        let weight_g = g_tensors.remove(&stem).ok_or_else(|| {
+            anyhow!("orphan weight_v at {stem} (no matching weight_g entry)")
+        })?;
+        let merged = merge_weight_norm(&weight_v, &weight_g)
+            .with_context(|| format!("merging weight_norm at {stem}"))?;
+        out_map.insert(format!("{stem}.weight"), merged);
+        merged_count += 1;
+    }
+
+    if !g_tensors.is_empty() {
+        let orphans: Vec<_> = g_tensors.keys().cloned().collect();
+        return Err(anyhow!("orphan weight_g entries (no matching weight_v): {orphans:?}"));
+    }
+
+    let pass_count = passthrough.len();
+    for (name, tensor) in passthrough {
+        out_map.insert(name, tensor);
+    }
+
+    ct_safetensors::save(&out_map, output.as_ref())
+        .with_context(|| format!("writing {}", output.as_ref().display()))?;
+
+    Ok(ConvertReport {
+        merged_weight_norm_pairs: merged_count,
+        passthrough_tensors: pass_count,
+        total_tensors_written: out_map.len(),
+    })
+}
+
+#[derive(Debug, Clone)]
+pub struct ConvertReport {
+    pub merged_weight_norm_pairs: usize,
+    pub passthrough_tensors: usize,
+    pub total_tensors_written: usize,
+}
+
+/// Compute `g * v / ‖v‖₂` where the L2 norm is taken over every axis
+/// except dim 0 (PyTorch's `weight_norm(..., dim=0)` semantics).
+///
+/// For Conv1d: `v: (C_out, C_in, K)`, `g: (C_out, 1, 1)` → result `(C_out, C_in, K)`.
+/// For ConvTranspose1d: `v: (C_in, C_out, K)`, `g: (C_in, 1, 1)` — same merge,
+/// since dim-0 is whatever PyTorch chose to scale (the formula is symmetric).
+pub fn merge_weight_norm(v: &Tensor, g: &Tensor) -> Result<Tensor> {
+    let rank = v.rank();
+    if rank < 2 {
+        return Err(anyhow!("weight_norm v expected rank>=2, got rank={rank}"));
+    }
+    // sum_keepdim over all axes except 0.
+    let mut norm_sq = v.sqr().context("v.sqr")?;
+    for axis in (1..rank).rev() {
+        norm_sq = norm_sq.sum_keepdim(axis).context("norm sum")?;
+    }
+    let norm = norm_sq.sqrt().context("sqrt")?;
+    let scale = g.broadcast_div(&norm).context("g / norm")?;
+    let out = v.broadcast_mul(&scale).context("v * scale")?;
+    Ok(out)
+}
+
+/// Build a `VarBuilder` over the converted safetensors generator file +
+/// detector file. Caller picks dtype + device.
+///
+/// Returned tuple is `(generator_vb, detector_vb)` — pass to
+/// `audioseal::AudioSealWatermarker::from_var_builders`.
+pub fn open_var_builders<'a>(
+    generator_safetensors: impl AsRef<Path>,
+    detector_safetensors: impl AsRef<Path>,
+    dtype: DType,
+    device: &Device,
+) -> Result<(candle_nn::VarBuilder<'a>, candle_nn::VarBuilder<'a>)> {
+    let gen_vb = unsafe {
+        candle_nn::VarBuilder::from_mmaped_safetensors(
+            &[generator_safetensors.as_ref()],
+            dtype,
+            device,
+        )
+    }
+    .context("opening generator safetensors")?;
+    let det_vb = unsafe {
+        candle_nn::VarBuilder::from_mmaped_safetensors(
+            &[detector_safetensors.as_ref()],
+            dtype,
+            device,
+        )
+    }
+    .context("opening detector safetensors")?;
+    Ok((gen_vb, det_vb))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn weight_norm_merge_matches_manual() {
+        // Reproduce PyTorch weight_norm formula on a small (2, 3) tensor.
+        let device = Device::Cpu;
+        let v = Tensor::from_slice(
+            &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0],
+            (2, 3),
+            &device,
+        )
+        .unwrap();
+        let g = Tensor::from_slice(&[2.0f32, 3.0], (2, 1), &device).unwrap();
+        let merged = merge_weight_norm(&v, &g).unwrap();
+        let merged: Vec<f32> = merged.flatten_all().unwrap().to_vec1().unwrap();
+
+        // Row 0: ‖[1,2,3]‖ = sqrt(14); scaled = 2 * [1,2,3] / sqrt(14)
+        let n0 = (1.0f32 + 4.0 + 9.0).sqrt();
+        let n1 = (16.0f32 + 25.0 + 36.0).sqrt();
+        let expected = vec![
+            2.0 * 1.0 / n0, 2.0 * 2.0 / n0, 2.0 * 3.0 / n0,
+            3.0 * 4.0 / n1, 3.0 * 5.0 / n1, 3.0 * 6.0 / n1,
+        ];
+        for (a, b) in merged.iter().zip(expected.iter()) {
+            assert!((a - b).abs() < 1e-6, "merge mismatch: got {a}, want {b}");
+        }
+    }
+
+    #[test]
+    fn weight_norm_merge_3d_conv1d_shape() {
+        let device = Device::Cpu;
+        let v = Tensor::randn(0f32, 1f32, (16, 8, 3), &device).unwrap();
+        let g = Tensor::randn(0f32, 1f32, (16, 1, 1), &device).unwrap();
+        let out = merge_weight_norm(&v, &g).unwrap();
+        assert_eq!(out.dims(), &[16, 8, 3]);
+    }
+}

crates/models/rtx-csm/src/hub.rs

@@ -17,6 +17,10 @@ pub const REPO_LLAMA_TOKENIZER: &str = "meta-llama/Llama-3.2-1B";
 /// Public mirror — tokenizer.json is byte-identical to meta-llama's. Used as
 /// fallback when the user hasn't been approved on Meta's gated form yet.
 pub const REPO_LLAMA_TOKENIZER_FALLBACK: &str = "unsloth/Llama-3.2-1B";
+/// AudioSeal watermark — public, no auth required.
+pub const REPO_AUDIOSEAL: &str = "facebook/audioseal";
+pub const FILE_AUDIOSEAL_GENERATOR: &str = "generator_base.pth";
+pub const FILE_AUDIOSEAL_DETECTOR: &str = "detector_base.pth";
 
 #[derive(Debug, Clone)]
 pub struct CsmAssets {
@@ -51,6 +55,22 @@ pub fn resolve_mimi() -> Result<PathBuf> {
     Ok(api.model(REPO_MIMI.to_string()).get("model.safetensors")?)
 }
 
+/// Download `facebook/audioseal/generator_base.pth`. Public, no auth required.
+pub fn resolve_audioseal_generator() -> Result<PathBuf> {
+    let api = Api::new()?;
+    Ok(api
+        .model(REPO_AUDIOSEAL.to_string())
+        .get(FILE_AUDIOSEAL_GENERATOR)?)
+}
+
+/// Download `facebook/audioseal/detector_base.pth`. Public, no auth required.
+pub fn resolve_audioseal_detector() -> Result<PathBuf> {
+    let api = Api::new()?;
+    Ok(api
+        .model(REPO_AUDIOSEAL.to_string())
+        .get(FILE_AUDIOSEAL_DETECTOR)?)
+}
+
 /// Download Llama-3.2-1B `tokenizer.json`.
 ///
 /// Tries the canonical `meta-llama/Llama-3.2-1B` first; on 403 (terms not

crates/models/rtx-csm/src/lib.rs

@@ -6,6 +6,7 @@
 pub mod asr;
 pub mod audio_io;
 pub mod audioseal;
+pub mod audioseal_convert;
 pub mod config;
 pub mod csm_fork;
 pub mod csm_quantized;