rtx-csm: WavLM-SV converter + end-to-end speaker similarity
Phase 5c — pure-Rust converter for microsoft/wavlm-base-plus-sv with auto-detecting weight_norm merger; verified on real 100M-param weights: load + embed + cosine-similarity round-trip works on Metal. - wavlm_sv_convert.rs: candle_core::pickle reads pytorch_model.bin directly. merge_weight_norm_auto picks the kept dim from g's shape (dim=0 for AudioSeal SEANet, dim=2 for WavLM pos_conv_embed). Skips classifier.*/objective.* (train-only AMSoftmax head). - examples/wavlm_sv_convert: HF download + convert CLI. Verified output: 1 weight_norm pair merged + 261 passthrough + 3 skipped = 262 tensors. - examples/wavlm_sv_demo: load + embed pair of WAVs + cosine similarity. - examples/audioseal_inspect: gains --which wavlm-sv variant for key discovery. - hub.rs: REPO_WAVLM_SV + resolve_wavlm_sv() helper. - wavlm_sv::XVectorHead bug fix: layer_weights is top-level, not nested under prefix. Verified end-to-end on Metal: cosine sim 0.9985 on same-speaker pair (CSM vs CSM-watermarked, 10s @ 24 kHz resampled to 16 kHz). Numerical parity vs HF reference is Phase 5d. 3 converter tests + 12 wavlm_sv tests; 78 lib tests total green. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -0,0 +1,208 @@
|
||||
//! 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::{anyhow, Context, Result};
|
||||
use candle_core::{pickle, safetensors as ct_safetensors, Tensor};
|
||||
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();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user