rustytorch/crates/models/rtx-csm/src/audioseal_convert.rs

//! 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::{Context, Result, anyhow};
use candle_core::{DType, Device, Tensor, pickle, safetensors as ct_safetensors};
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]);
    }
}