Files
rustytorch/crates/models/rtx-csm/src/watermark.rs
T
osobhandClaude Opus 4.7 9161b32a91 rtx-csm: per-request watermark message in tts_server
Watermarker trait gains embed_with_message(audio, message) with a
default impl forwarding to embed (no-op for watermarkers without a
payload). AudioSealWatermarker overrides to use the requested message
instead of self.message; ResampledWatermarker forwards through the
resample dance.

TtsRequest gains optional watermark_message: Option<String> (decimal or
0xHEX). Useful for clawsample to tag each generation with a unique ID
(e.g. job_id mod 0x10000) for audit trails. When omitted, falls back
to the server-startup --audioseal-message default.

Verified end-to-end: override "0xBEEF" -> detect 0xBEEF (mean_presence
0.9995, 16/16 bits). Default fallback also decodes correctly.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
2026-04-26 04:01:30 -07:00

83 lines
2.7 KiB
Rust

//! Watermarking hook for the TTS pipeline.
//!
//! The trait is sample-rate-agnostic: callers feed PCM in whatever native
//! 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 {
fn embed(&self, audio: &[f32]) -> Result<Vec<f32>>;
/// Embed with a specific 16-bit message payload, overriding the
/// watermarker's default. Default impl ignores `_message` and falls
/// back to [`Self::embed`] — implementations that don't carry a
/// payload (NoopWatermarker, future SilentCipher) inherit this no-op.
/// Implementations that DO carry a payload (AudioSeal) override.
fn embed_with_message(&self, audio: &[f32], _message: u16) -> Result<Vec<f32>> {
self.embed(audio)
}
}
#[derive(Default, Debug, Clone, Copy)]
pub struct NoopWatermarker;
impl Watermarker for NoopWatermarker {
fn embed(&self, audio: &[f32]) -> Result<Vec<f32>> {
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> ResampledWatermarker<W> {
fn run<F>(&self, audio: &[f32], inner_call: F) -> Result<Vec<f32>>
where
F: FnOnce(&[f32]) -> Result<Vec<f32>>,
{
if self.source_rate == self.model_rate {
return inner_call(audio);
}
let n = audio.len();
let down = audio_io::resample(audio, self.source_rate, self.model_rate)?;
let watermarked = inner_call(&down)?;
let mut up = audio_io::resample(&watermarked, self.model_rate, self.source_rate)?;
if up.len() > n {
up.truncate(n);
} else if up.len() < n {
up.resize(n, 0.0);
}
Ok(up)
}
}
impl<W: Watermarker> Watermarker for ResampledWatermarker<W> {
fn embed(&self, audio: &[f32]) -> Result<Vec<f32>> {
self.run(audio, |a| self.inner.embed(a))
}
fn embed_with_message(&self, audio: &[f32], message: u16) -> Result<Vec<f32>> {
self.run(audio, |a| self.inner.embed_with_message(a, message))
}
}