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]>
This commit is contained in:
osobh
2026-04-25 19:26:13 -07:00
co-authored by Claude Opus 4.7
parent 15dd3575d4
commit 938b54a2b0
9 changed files with 1398 additions and 65 deletions
+12
View File
@@ -124,3 +124,15 @@ path = "examples/tts_server.rs"
[[example]] [[example]]
name = "lora_train" name = "lora_train"
path = "examples/lora_train.rs" 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"
@@ -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(())
}
@@ -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(())
}
@@ -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(())
}
+77
View File
@@ -142,6 +142,83 @@ fn resample_to_24k(input: &[f32], src_rate: u32) -> Result<Vec<f32>> {
Ok(out) 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. /// 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<()> { pub fn write_wav_24k_mono<P: AsRef<Path>>(path: P, samples: &[f32]) -> Result<()> {
let spec = hound::WavSpec { let spec = hound::WavSpec {
+758 -65
View File
@@ -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 //! [arXiv:2401.17264]: https://arxiv.org/abs/2401.17264
use crate::error::{CsmError, Result}; use crate::error::{CsmError, Result};
use crate::watermark::Watermarker; 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; 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. /// AudioSeal embeds a 16-bit message per audio frame.
pub const MESSAGE_BITS: usize = 16; 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)] #[derive(Debug, Clone)]
pub struct DetectionResult { 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>, pub presence_per_sample: Vec<f32>,
/// Decoded 16-bit message bits if detection passed; `None` otherwise.
pub message: Option<u16>, pub message: Option<u16>,
/// Mean presence — useful as a quick yes/no with `> 0.5` threshold.
pub mean_presence: f32, pub mean_presence: f32,
} }
pub struct AudioSealWatermarker { pub struct AudioSealWatermarker {
/// Device for inference.
pub device: Device, pub device: Device,
/// 16-bit message to embed (or expect from detection). 0 = no payload.
pub message: u16, 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 { impl AudioSealWatermarker {
@@ -68,63 +616,208 @@ impl AudioSealWatermarker {
Self { Self {
device, device,
message, 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> { 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)) 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(), .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 { 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(
.into(), "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)] #[cfg(test)]
mod tests { mod tests {
use super::*; 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] #[test]
fn scaffold_constructs() { fn scaffold_constructs() {
let w = AudioSealWatermarker::new(Device::Cpu, 0xABCD); let w = AudioSealWatermarker::new(Device::Cpu, 0xABCD);
assert_eq!(w.message, 0xABCD); assert_eq!(w.message, 0xABCD);
assert!(w.generator_weights.is_none());
} }
#[test] #[test]
fn embed_returns_typed_error() { fn embed_returns_typed_error_without_weights() {
let w = AudioSealWatermarker::new(Device::Cpu, 0); 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] #[test]
fn detect_returns_typed_error() { fn detect_returns_typed_error_without_weights() {
let w = AudioSealWatermarker::new(Device::Cpu, 0); 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());
} }
} }
@@ -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]);
}
}
+20
View File
@@ -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 /// 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. /// 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"; 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)] #[derive(Debug, Clone)]
pub struct CsmAssets { pub struct CsmAssets {
@@ -51,6 +55,22 @@ pub fn resolve_mimi() -> Result<PathBuf> {
Ok(api.model(REPO_MIMI.to_string()).get("model.safetensors")?) 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`. /// Download Llama-3.2-1B `tokenizer.json`.
/// ///
/// Tries the canonical `meta-llama/Llama-3.2-1B` first; on 403 (terms not /// Tries the canonical `meta-llama/Llama-3.2-1B` first; on 403 (terms not
+1
View File
@@ -6,6 +6,7 @@
pub mod asr; pub mod asr;
pub mod audio_io; pub mod audio_io;
pub mod audioseal; pub mod audioseal;
pub mod audioseal_convert;
pub mod config; pub mod config;
pub mod csm_fork; pub mod csm_fork;
pub mod csm_quantized; pub mod csm_quantized;