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
@@ -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]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user