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]>
82 lines
2.5 KiB
Rust
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(())
|
|
}
|