osobhandClaude Opus 4.7
938b54a2b0
SEANet generator + detector matching `facebook/audioseal` reference layout (weight_norm-merged via pure-Rust pickle reader). Verified on real CSM speech: mean_presence=0.9943, 16/16 message bits decoded. - src/audioseal.rs: SeanetEncoder (4-stage strided downsample, 2-layer LSTM bottleneck at 512 channels, 128-dim projection), MsgProcessor (16-bit message via embedding sum + broadcast-add), SeanetDecoder, Generator (encoder+msg+decoder), Detector (encoder + single 320× reverse_convolution + 1×1 head). Padding mirrors audiocraft _get_extra_padding_for_conv1d exactly. - src/audioseal_convert.rs: candle_core::pickle reads .pth directly; merge_weight_norm computes g*v/‖v‖ over all axes except 0; writes flat safetensors keyed identically to what Generator/Detector read. - examples/audioseal_inspect.rs: dumps tensor keys + shapes. - examples/audioseal_convert.rs: HF download + convert CLI. - examples/audioseal_demo.rs: load + embed + detect on real WAV or synthetic burst, optionally writes watermarked WAV. - audio_io.rs gains generic load_mono_at_rate, resample, write_wav_mono (16 kHz path needed for AudioSeal). 12 new unit tests + 2 converter tests; 63 lib tests total green. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
80 lines
2.4 KiB
Rust
80 lines
2.4 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,
|
|
}
|
|
|
|
#[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()?,
|
|
},
|
|
};
|
|
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(())
|
|
}
|