rtx-csm: Generator inline watermarker + ResampledWatermarker adapter

A single \`generate\` invocation now produces a watermarked WAV when
AudioSeal weights are passed via CLI. End-to-end verified on real CSM
speech: mean_presence=1.0000, 16/16 message bits decoded.

- Generator gains \`watermarker: Option<Box<dyn Watermarker>>\` slot;
  \`generate_to_wav\` runs \`wm.embed(&pcm)\` after post-process, before
  WAV write. Field is Send+Sync so the existing Arc<Mutex<Generator>>
  tts_server pattern still works.
- watermark.rs ships ResampledWatermarker<W> adapter for handling rate
  mismatches (CSM 24 kHz ↔ AudioSeal 16 kHz). Output length is normalized
  to input length so it's a transparent drop-in.
- examples/generate.rs gains --watermark-generator/--watermark-detector/
  --watermark-message flags. Loads AudioSeal, wraps in resampler, installs.

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

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

@@ -98,6 +98,24 @@ struct Cli {
     /// LUFS target for loudness normalization. Ignored if --raw.
     #[arg(long, default_value_t = -16.0)]
     lufs: f32,
+
+    /// Path to converted AudioSeal generator safetensors (run
+    /// `audioseal_convert` first). When set together with
+    /// `--watermark-detector`, the watermarker is wired into
+    /// `generate_to_wav` so output is automatically watermarked.
+    #[arg(long)]
+    watermark_generator: Option<std::path::PathBuf>,
+
+    /// Path to converted AudioSeal detector safetensors. Required for the
+    /// watermarker even if you only want to embed (the detector is part of
+    /// AudioSealWatermarker construction; future builds may make it
+    /// optional).
+    #[arg(long)]
+    watermark_detector: Option<std::path::PathBuf>,
+
+    /// 16-bit watermark message (decimal or 0xHEX).
+    #[arg(long, default_value = "0")]
+    watermark_message: String,
 }
 
 fn main() -> Result<()> {
@@ -152,6 +170,50 @@ fn main() -> Result<()> {
         context.push(Segment::new(cli.context_speaker, txt, audio));
     }
 
+    // Optional watermarker wiring (AudioSeal + 24k↔16k resample adapter).
+    if let (Some(gen_path), Some(det_path)) = (
+        cli.watermark_generator.as_ref(),
+        cli.watermark_detector.as_ref(),
+    ) {
+        let msg_str = cli.watermark_message.trim();
+        let message: u16 = if let Some(rest) = msg_str
+            .strip_prefix("0x")
+            .or_else(|| msg_str.strip_prefix("0X"))
+        {
+            u16::from_str_radix(rest, 16)?
+        } else {
+            msg_str.parse::<u16>()?
+        };
+        let gen_vb = unsafe {
+            candle_nn::VarBuilder::from_mmaped_safetensors(
+                &[gen_path],
+                candle_core::DType::F32,
+                &device,
+            )
+        }?;
+        let det_vb = unsafe {
+            candle_nn::VarBuilder::from_mmaped_safetensors(
+                &[det_path],
+                candle_core::DType::F32,
+                &device,
+            )
+        }?;
+        let inner = rtx_csm::AudioSealWatermarker::from_var_builders(
+            gen_vb,
+            det_vb,
+            device.clone(),
+            message,
+        )?;
+        // CSM produces 24 kHz; AudioSeal native is 16 kHz.
+        let wm = rtx_csm::ResampledWatermarker::new(inner, generator.config.sample_rate, 16_000);
+        generator.set_watermarker(Box::new(wm));
+        tracing::info!(
+            "watermarker installed (message=0x{:04X}, model 16 kHz, output {} Hz)",
+            message,
+            generator.config.sample_rate
+        );
+    }
+
     let opts = GenerateOptions {
         max_audio_ms: cli.max_audio_ms,
         temperature: cli.temperature,

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

@@ -56,6 +56,11 @@ pub struct Generator {
     /// Applied to every text input before tokenization. Set to
     /// [`TextNormalize::passthrough`] if you've pre-normalized upstream.
     pub text_normalize: TextNormalize,
+    /// Optional watermarker applied inside [`Self::generate_to_wav`] AFTER
+    /// post-processing and BEFORE WAV write. Set via
+    /// [`Self::set_watermarker`]. `None` = no-op (Sesame's reference TTS
+    /// also ships unwatermarked by default; this is the integration hook).
+    pub watermarker: Option<Box<dyn crate::watermark::Watermarker>>,
 }
 
 impl Generator {
@@ -68,9 +73,22 @@ impl Generator {
             config,
             device,
             text_normalize: TextNormalize::default(),
+            watermarker: None,
         }
     }
 
+    /// Install a watermarker that runs inside `generate_to_wav` after
+    /// post-processing. Take ownership of the watermarker so the generator
+    /// can be moved into worker threads (Watermarker is `Send + Sync`).
+    pub fn set_watermarker(&mut self, wm: Box<dyn crate::watermark::Watermarker>) {
+        self.watermarker = Some(wm);
+    }
+
+    /// Drop any installed watermarker.
+    pub fn clear_watermarker(&mut self) {
+        self.watermarker = None;
+    }
+
     /// Download (cached) CSM-1B + Mimi + Llama tokenizer from HuggingFace and
     /// build a ready-to-generate `Generator`.
     pub fn load_csm_1b(device: &Device) -> Result<Self> {
@@ -439,8 +457,15 @@ impl Generator {
         self.generate(text, profile.id, &ctx, opts)
     }
 
-    /// Convenience: `generate` + apply default post-processing + write WAV.
+    /// Convenience: `generate` + apply default post-processing + watermark
-    /// Pass [`PostProcess::disabled`] to skip post-processing.
+    /// (if installed) + write WAV. Pass [`PostProcess::disabled`] to skip
+    /// post-processing. Pass [`Self::clear_watermarker`] (or never install
+    /// one) to skip watermarking.
+    ///
+    /// Order: model → post-process (HPF + declick + LUFS) → watermark.
+    /// Watermarking comes last so the loudness target the user sees on disk
+    /// is the loudness target the user requested (the watermark residual
+    /// is at most a few dB and well below LUFS measurement floor).
     pub fn generate_to_wav(
         &mut self,
         text: &str,
@@ -452,6 +477,9 @@ impl Generator {
     ) -> Result<()> {
         let mut pcm = self.generate(text, speaker, context, opts)?;
         post.apply(&mut pcm, self.config.sample_rate)?;
+        if let Some(wm) = self.watermarker.as_ref() {
+            pcm = wm.embed(&pcm)?;
+        }
         crate::audio_io::write_wav_24k_mono(out_path, &pcm)?;
         Ok(())
     }

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

@@ -47,4 +47,4 @@ pub use speaker_sim::{
 };
 pub use text_norm::TextNormalize;
 pub use wer::{wer as compute_wer, WerResult};
-pub use watermark::{NoopWatermarker, Watermarker};
+pub use watermark::{NoopWatermarker, ResampledWatermarker, Watermarker};

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

@@ -1,8 +1,11 @@
-//! Watermarking hook (SilentCipher in Sesame's Python; no Rust port yet).
+//! Watermarking hook for the TTS pipeline.
 //!
-//! Stage 1 ships `NoopWatermarker`. Stage 3 will implement a real port or a
+//! The trait is sample-rate-agnostic: callers feed PCM in whatever native
-//! PyO3 bridge.
+//! rate they have and receive PCM at the same rate. Wrap a model-bound
+//! watermarker in [`ResampledWatermarker`] when the model expects a
+//! different rate (e.g. AudioSeal trained at 16 kHz on a 24 kHz CSM stream).
 
+use crate::audio_io;
 use crate::error::Result;
 
 pub trait Watermarker: Send + Sync {
@@ -17,3 +20,43 @@ impl Watermarker for NoopWatermarker {
         Ok(audio.to_vec())
     }
 }
+
+/// Wraps a `Watermarker` that operates at `model_rate` so callers can use
+/// it on audio at `source_rate`. Resamples in/out via the existing rubato
+/// pipeline. Output length is normalized to the input length to keep this
+/// as a drop-in equivalent of a same-rate watermarker.
+pub struct ResampledWatermarker<W: Watermarker> {
+    pub inner: W,
+    pub source_rate: u32,
+    pub model_rate: u32,
+}
+
+impl<W: Watermarker> ResampledWatermarker<W> {
+    pub fn new(inner: W, source_rate: u32, model_rate: u32) -> Self {
+        Self {
+            inner,
+            source_rate,
+            model_rate,
+        }
+    }
+}
+
+impl<W: Watermarker> Watermarker for ResampledWatermarker<W> {
+    fn embed(&self, audio: &[f32]) -> Result<Vec<f32>> {
+        if self.source_rate == self.model_rate {
+            return self.inner.embed(audio);
+        }
+        let n = audio.len();
+        let down = audio_io::resample(audio, self.source_rate, self.model_rate)?;
+        let watermarked = self.inner.embed(&down)?;
+        let mut up = audio_io::resample(&watermarked, self.model_rate, self.source_rate)?;
+        // Sinc resample isn't bidirectionally length-preserving — clamp/pad
+        // to match the input so callers can drop this in transparently.
+        if up.len() > n {
+            up.truncate(n);
+        } else if up.len() < n {
+            up.resize(n, 0.0);
+        }
+        Ok(up)
+    }
+}