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

//! Offline converter: `microsoft/wavlm-base-plus-sv/pytorch_model.bin`
//! → flat safetensors with `weight_norm` merged for the positional conv.
//!
//! ## What this does
//!
//! The HF state_dict is mostly a direct passthrough — every key matches
//! what `wavlm_sv::WavLmSv::new` expects. The only structural change is
//! merging the `weight_norm` parametrization on
//! `wavlm.encoder.pos_conv_embed.conv`:
//!
//!   - `weight_g` shape `(1, 1, 128)` — per-position scale (norm over dims 0,1)
//!   - `weight_v` shape `(768, 48, 128)` — unnormalized direction
//!
//! At forward time PyTorch computes
//! `weight = weight_g * weight_v / ‖weight_v‖₂` where the L2 norm runs
//! over every axis EXCEPT dim=2 (i.e. axes 0,1). We do the merge once at
//! conversion time and write a flat `weight` instead.
//!
//! Also dropped (training-only / no inference value):
//!   - `classifier.weight`, `classifier.bias` (logits head, not the embedding)
//!   - `objective.weight` (AMSoftmax train-only)
//!
//! ## Why pure-Rust
//!
//! `candle_core::pickle::read_all` reads the `.bin` directly. No Python
//! step needed in the conversion pipeline.

use anyhow::{Context, Result, anyhow};
use candle_core::{Tensor, pickle, safetensors as ct_safetensors};
use std::collections::HashMap;
use std::path::Path;

const SKIP_PREFIXES: &[&str] = &["classifier.", "objective."];

/// Read `pytorch_model.bin`, merge `weight_norm` on `pos_conv_embed.conv`,
/// drop classifier/objective tensors, and write a flat safetensors keyed
/// identically to what `wavlm_sv::WavLmSv::new` reads.
pub fn convert_pth(input: impl AsRef<Path>, output: impl AsRef<Path>) -> Result<ConvertReport> {
    let tensors = pickle::read_all(input.as_ref())
        .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();
    let mut skipped = 0usize;

    for (name, tensor) in tensors {
        if SKIP_PREFIXES.iter().any(|p| name.starts_with(p)) {
            skipped += 1;
            continue;
        }
        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("present");
        let weight_g = g_tensors
            .remove(&stem)
            .ok_or_else(|| anyhow!("orphan weight_v at {stem}"))?;
        let merged = merge_weight_norm_auto(&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: {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,
        skipped_tensors: skipped,
        total_tensors_written: out_map.len(),
    })
}

#[derive(Debug, Clone)]
pub struct ConvertReport {
    pub merged_weight_norm_pairs: usize,
    pub passthrough_tensors: usize,
    pub skipped_tensors: usize,
    pub total_tensors_written: usize,
}

/// Compute `g * v / ‖v‖₂` with the norm taken over every axis EXCEPT
/// the one corresponding to `g`'s non-singleton dimension. Auto-detects
/// the kept axis from `g`'s shape:
///   - g shape `(C, 1, 1)` → kept dim = 0 (audiocraft / SEANet style)
///   - g shape `(1, 1, K)` → kept dim = 2 (WavLM positional conv)
pub fn merge_weight_norm_auto(v: &Tensor, g: &Tensor) -> Result<Tensor> {
    let g_shape = g.dims();
    let kept_dim = g_shape
        .iter()
        .position(|&d| d != 1)
        .ok_or_else(|| anyhow!("weight_g has no non-singleton dim: shape {g_shape:?}"))?;
    merge_weight_norm_dim(v, g, kept_dim)
}

/// Generic `weight_norm` merger: `weight = g * v / ‖v‖₂` with the L2 norm
/// taken over every axis except `kept_dim`.
pub fn merge_weight_norm_dim(v: &Tensor, g: &Tensor, kept_dim: usize) -> Result<Tensor> {
    let rank = v.rank();
    if kept_dim >= rank {
        return Err(anyhow!("kept_dim {kept_dim} out of range for rank {rank}"));
    }
    let mut norm_sq = v.sqr().context("v.sqr")?;
    for axis in (0..rank).rev() {
        if axis == kept_dim {
            continue;
        }
        norm_sq = norm_sq.sum_keepdim(axis).context("sum_keepdim")?;
    }
    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)
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::Device;

    #[test]
    fn merge_dim_0_matches_audioseal_path() {
        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_dim(&v, &g, 0).unwrap();
        let m: Vec<f32> = merged.flatten_all().unwrap().to_vec1().unwrap();
        let n0 = (1.0f32 + 4.0 + 9.0).sqrt();
        let n1 = (16.0f32 + 25.0 + 36.0).sqrt();
        let expected = [
            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 m.iter().zip(expected.iter()) {
            assert!((a - b).abs() < 1e-6);
        }
    }

    #[test]
    fn merge_dim_2_matches_pytorch_pos_conv() {
        // Simulate a 2-output 2-in 3-kernel weight with weight_norm dim=2
        // (norm taken over axes 0,1). Per-kernel-position scale.
        let device = Device::Cpu;
        let v = Tensor::from_slice(
            &[
                1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
            ],
            (2, 2, 3),
            &device,
        )
        .unwrap();
        let g = Tensor::from_slice(&[1.0f32, 0.5, 2.0], (1, 1, 3), &device).unwrap();
        let merged = merge_weight_norm_dim(&v, &g, 2).unwrap();
        // For each kernel position k, norm = ||v[:,:,k]||, weight = g[k] * v[:,:,k] / norm.
        let v_flat: Vec<f32> = v.flatten_all().unwrap().to_vec1().unwrap();
        let g_flat: Vec<f32> = g.flatten_all().unwrap().to_vec1().unwrap();
        let m_flat: Vec<f32> = merged.flatten_all().unwrap().to_vec1().unwrap();
        for k in 0..3 {
            let mut sumsq = 0.0f32;
            for i in 0..2 {
                for j in 0..2 {
                    sumsq += v_flat[i * 6 + j * 3 + k].powi(2);
                }
            }
            let norm = sumsq.sqrt();
            for i in 0..2 {
                for j in 0..2 {
                    let idx = i * 6 + j * 3 + k;
                    let want = g_flat[k] * v_flat[idx] / norm;
                    assert!(
                        (m_flat[idx] - want).abs() < 1e-5,
                        "k={k} i={i} j={j}: got {} want {want}",
                        m_flat[idx]
                    );
                }
            }
        }
    }

    #[test]
    fn auto_detect_picks_correct_dim() {
        let device = Device::Cpu;
        let v0 = Tensor::randn(0f32, 1f32, (4, 8, 3), &device).unwrap();
        let g0 = Tensor::randn(0f32, 1f32, (4, 1, 1), &device).unwrap();
        let _ = merge_weight_norm_auto(&v0, &g0).unwrap();
        let v2 = Tensor::randn(0f32, 1f32, (4, 8, 3), &device).unwrap();
        let g2 = Tensor::randn(0f32, 1f32, (1, 1, 3), &device).unwrap();
        let _ = merge_weight_norm_auto(&v2, &g2).unwrap();
    }
}