Files
rustytorch/crates/models/rtx-csm/examples/audioseal_inspect.rs
T
osobhandClaude Opus 4.7 d1dba7a05c 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]>
2026-04-25 19:57:52 -07:00

82 lines
2.5 KiB
Rust

//! 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,
WavlmSv,
}
#[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()?,
Which::WavlmSv => hub::resolve_wavlm_sv()?,
},
};
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(())
}