osobhandClaude Opus 4.7
a5cedfb46a
8-gen bench (4 emotions × 2 corpora) at seed=42 against firdhokk Whisper-LV3: target RAVDESS CREMA-D happy happy (0.999) ✓ happy (0.999) ✓ angry neutral (0.92) sad (0.99) fearful happy (0.998) fearful (0.984) ✓ sad angry (0.99) fearful (0.99) CREMA-D 2/4 vs RAVDESS 1/4. Larger / more naturalistic corpus produces more class-pure fearful direction. Neither corpus solves angry or sad — recipe shifts into 'vague expressivity' rather than class-specific corners. Practical: prefer CREMA-D when available; A/B both per emotion if class precision matters. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
842 lines
29 KiB
Rust
842 lines
29 KiB
Rust
//! AudioSeal watermark — SEANet generator + detector.
|
||
//!
|
||
//! Architectural port of Meta's AudioSeal (Roman et al., ICML 2024,
|
||
//! [arXiv:2401.17264]) against candle 0.9.
|
||
//!
|
||
//! ## Module layout (matches `facebook/audioseal` reference exactly)
|
||
//!
|
||
//! Generator state_dict keys are organized as a `nn.Sequential` indexed by
|
||
//! integer position. With `n_filters=32, ratios=[8,5,4,2], n_residual_layers=1,
|
||
//! lstm=2, dimension=128` the encoder runs:
|
||
//!
|
||
//! ```text
|
||
//! encoder.model.0 Conv1d(1, 32, k=7) (SConv1d wrapper)
|
||
//! encoder.model.1 ResidualBlock(32, dilation=1) (n_residual=1, j=0)
|
||
//! encoder.model.2 ELU(1.0)
|
||
//! 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)
|
||
//! ```
|
||
//!
|
||
//! Decoder mirrors this: index 0 is the 128→512 init conv, index 1 is the
|
||
//! LSTM, indices 3,6,9,12 are ConvTranspose1d upsamples (with ratios
|
||
//! [8,5,4,2] in order), residuals at 4,7,10,13, final 32→1 conv at 15.
|
||
//!
|
||
//! 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::{D, DType, Device, IndexOp, Module, Tensor};
|
||
use candle_nn::{
|
||
Activation, Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, Embedding, LSTM,
|
||
LSTMConfig, RNN, VarBuilder, conv_transpose1d, conv1d, embedding, lstm, ops,
|
||
};
|
||
use std::path::Path;
|
||
|
||
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;
|
||
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) ------------------------
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct DetectionResult {
|
||
pub presence_per_sample: Vec<f32>,
|
||
pub message: Option<u16>,
|
||
pub mean_presence: f32,
|
||
}
|
||
|
||
pub struct AudioSealWatermarker {
|
||
pub device: Device,
|
||
pub message: u16,
|
||
pub generator: Option<Generator>,
|
||
pub detector: Option<Detector>,
|
||
pub alpha: f32,
|
||
}
|
||
|
||
impl AudioSealWatermarker {
|
||
pub fn new(device: Device, message: u16) -> Self {
|
||
Self {
|
||
device,
|
||
message,
|
||
generator: None,
|
||
detector: None,
|
||
alpha: 1.0,
|
||
}
|
||
}
|
||
|
||
/// Construct from in-memory VarBuilders rooted at the generator and
|
||
/// detector subtrees (after `weight_norm` merge by the converter).
|
||
pub fn from_var_builders(
|
||
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> {
|
||
Ok(Self::new(device, message))
|
||
}
|
||
|
||
pub fn detect(&self, samples: &[f32]) -> Result<DetectionResult> {
|
||
let detector = self.detector.as_ref().ok_or_else(|| {
|
||
CsmError::Config(
|
||
"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 AudioSealWatermarker {
|
||
/// Internal implementation used by both `embed` (default message) and
|
||
/// `embed_with_message` (per-call override).
|
||
fn embed_inner(&self, audio: &[f32], message: u16) -> Result<Vec<f32>> {
|
||
let generator = self.generator.as_ref().ok_or_else(|| {
|
||
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, 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)
|
||
}
|
||
}
|
||
|
||
impl Watermarker for AudioSealWatermarker {
|
||
fn embed(&self, audio: &[f32]) -> Result<Vec<f32>> {
|
||
self.embed_inner(audio, self.message)
|
||
}
|
||
|
||
fn embed_with_message(&self, audio: &[f32], message: u16) -> Result<Vec<f32>> {
|
||
self.embed_inner(audio, message)
|
||
}
|
||
}
|
||
|
||
#[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);
|
||
}
|
||
|
||
#[test]
|
||
fn embed_returns_typed_error_without_weights() {
|
||
let w = AudioSealWatermarker::new(Device::Cpu, 0);
|
||
assert!(w.embed(&[0.0, 0.1, 0.2]).is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn detect_returns_typed_error_without_weights() {
|
||
let w = AudioSealWatermarker::new(Device::Cpu, 0);
|
||
assert!(w.detect(&[0.0, 0.1, 0.2]).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());
|
||
}
|
||
}
|