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]>
83 lines
2.7 KiB
Rust
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))
|
|
}
|
|
}
|