osobhandClaude Opus 4.7
a5cedfb46a
8-gen bench (4 emotions × 2 corpora) at seed=42 against firdhokk Whisper-LV3: target RAVDESS CREMA-D happy happy (0.999) ✓ happy (0.999) ✓ angry neutral (0.92) sad (0.99) fearful happy (0.998) fearful (0.984) ✓ sad angry (0.99) fearful (0.99) CREMA-D 2/4 vs RAVDESS 1/4. Larger / more naturalistic corpus produces more class-pure fearful direction. Neither corpus solves angry or sad — recipe shifts into 'vague expressivity' rather than class-specific corners. Practical: prefer CREMA-D when available; A/B both per emotion if class precision matters. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
242 lines
8.8 KiB
Rust
242 lines
8.8 KiB
Rust
//! End-to-end production pipeline showcase.
|
||
//!
|
||
//! Single command demonstrates the full rtx-csm stack:
|
||
//! 1. CSM-1B generates 24 kHz speech from text (with optional LoRA voice).
|
||
//! 2. Post-process: HPF + declick + EBU R128 −16 LUFS.
|
||
//! 3. AudioSeal watermark embedded transparently via the resampling
|
||
//! adapter (24 kHz ↔ 16 kHz round-trip).
|
||
//! 4. AudioSeal detector verifies the watermark on the written WAV.
|
||
//! 5. WavLM-SV computes a 512-d speaker embedding for the output and,
|
||
//! if a reference WAV is provided, scores cosine similarity against it.
|
||
//!
|
||
//! Usage:
|
||
//! ```
|
||
//! cargo run -p rtx-csm --release --features metal --example pipeline -- \
|
||
//! --text "Hello from the full Rust pipeline." \
|
||
//! --speaker 0 \
|
||
//! --audioseal-generator /tmp/audioseal_generator.safetensors \
|
||
//! --audioseal-detector /tmp/audioseal_detector.safetensors \
|
||
//! --audioseal-message 0xCAFE \
|
||
//! --wavlm-sv /tmp/wavlm_sv.safetensors \
|
||
//! --reference /tmp/csm_24k.wav \
|
||
//! --out /tmp/pipeline_out.wav
|
||
//! ```
|
||
//!
|
||
//! The reference WAV is expected to be a known-speaker sample (any rate;
|
||
//! resampled to 16 kHz internally). Cosine sim ≈ 1.0 confirms the output
|
||
//! sounds like the same speaker as the reference; lower values flag drift.
|
||
|
||
use anyhow::Result;
|
||
use candle_core::{DType, Device};
|
||
use clap::Parser;
|
||
use rtx_csm::{
|
||
GenerateOptions, Generator, PostProcess, Segment, audio_io,
|
||
audioseal::AudioSealWatermarker,
|
||
speaker_sim::{SpeakerSimilarity, WavLmSimilarity},
|
||
watermark::ResampledWatermarker,
|
||
};
|
||
use std::path::PathBuf;
|
||
|
||
const AUDIOSEAL_RATE: u32 = 16_000;
|
||
|
||
#[derive(Debug, Parser)]
|
||
#[command(name = "pipeline", about = "rtx-csm end-to-end showcase")]
|
||
struct Cli {
|
||
/// Text to synthesize.
|
||
#[arg(long)]
|
||
text: String,
|
||
/// Speaker id (0 or 1).
|
||
#[arg(long, default_value_t = 0)]
|
||
speaker: u32,
|
||
/// Output WAV path.
|
||
#[arg(long)]
|
||
out: PathBuf,
|
||
/// Max audio length in milliseconds.
|
||
#[arg(long, default_value_t = 8_000)]
|
||
max_audio_ms: u32,
|
||
/// LUFS target for loudness normalization.
|
||
#[arg(long, default_value_t = -16.0)]
|
||
lufs: f32,
|
||
/// Optional LoRA adapter (run examples/lora_train to produce one).
|
||
#[arg(long)]
|
||
lora: Option<PathBuf>,
|
||
#[arg(long, default_value_t = 8)]
|
||
lora_rank: usize,
|
||
#[arg(long, default_value_t = 16.0)]
|
||
lora_alpha: f32,
|
||
|
||
/// AudioSeal generator safetensors (produced by audioseal_convert).
|
||
/// If both --audioseal-generator and --audioseal-detector are set,
|
||
/// output is watermarked + detection round-trip is run.
|
||
#[arg(long)]
|
||
audioseal_generator: Option<PathBuf>,
|
||
#[arg(long)]
|
||
audioseal_detector: Option<PathBuf>,
|
||
/// 16-bit watermark message (decimal or 0xHEX).
|
||
#[arg(long, default_value = "0xCAFE")]
|
||
audioseal_message: String,
|
||
|
||
/// WavLM-SV safetensors (from wavlm_sv_convert). When set, the
|
||
/// pipeline computes a 512-d speaker embedding for the output WAV.
|
||
#[arg(long)]
|
||
wavlm_sv: Option<PathBuf>,
|
||
/// Optional reference WAV to score the output's speaker against.
|
||
#[arg(long)]
|
||
reference: Option<PathBuf>,
|
||
|
||
/// Force CPU device.
|
||
#[arg(long)]
|
||
cpu: bool,
|
||
}
|
||
|
||
fn parse_message(s: &str) -> Result<u16> {
|
||
let s = s.trim();
|
||
if let Some(rest) = s.strip_prefix("0x").or_else(|| s.strip_prefix("0X")) {
|
||
Ok(u16::from_str_radix(rest, 16)?)
|
||
} else {
|
||
Ok(s.parse::<u16>()?)
|
||
}
|
||
}
|
||
|
||
fn main() -> Result<()> {
|
||
tracing_subscriber::fmt().init();
|
||
let cli = Cli::parse();
|
||
let device = if cli.cpu {
|
||
Device::Cpu
|
||
} else {
|
||
Generator::default_device()?
|
||
};
|
||
println!("== rtx-csm pipeline ==");
|
||
println!("device: {device:?}");
|
||
|
||
// -- 1. Build the Generator --------------------------------------------
|
||
let t0 = std::time::Instant::now();
|
||
let mut generator = Generator::load_csm_1b(&device)?;
|
||
println!(
|
||
"[1/5] CSM-1B loaded in {:.2}s (sr={} Hz)",
|
||
t0.elapsed().as_secs_f32(),
|
||
generator.config.sample_rate
|
||
);
|
||
|
||
// Optional LoRA voice clone.
|
||
if let Some(lora_path) = cli.lora.as_ref() {
|
||
let lora_cfg = rtx_csm::lora::LoraConfig {
|
||
rank: cli.lora_rank,
|
||
alpha: cli.lora_alpha,
|
||
..rtx_csm::lora::LoraConfig::default()
|
||
};
|
||
let vm = candle_nn::VarMap::new();
|
||
generator.model.inner.add_lora_to_backbone(&lora_cfg, &vm)?;
|
||
rtx_csm::training::load_lora_adapter(&vm, lora_path, &device)?;
|
||
generator.model.inner.refresh_lora(&vm)?;
|
||
println!(
|
||
" LoRA adapter loaded: {} (rank={} alpha={})",
|
||
lora_path.display(),
|
||
cli.lora_rank,
|
||
cli.lora_alpha
|
||
);
|
||
}
|
||
|
||
// -- 2. Optional inline watermarker -----------------------------------
|
||
let watermark_message = parse_message(&cli.audioseal_message)?;
|
||
let mut audioseal_for_detect: Option<AudioSealWatermarker> = None;
|
||
if let (Some(g), Some(d)) = (
|
||
cli.audioseal_generator.as_ref(),
|
||
cli.audioseal_detector.as_ref(),
|
||
) {
|
||
let g_vb =
|
||
unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[g], DType::F32, &device) }?;
|
||
let d_vb =
|
||
unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[d], DType::F32, &device) }?;
|
||
let inner =
|
||
AudioSealWatermarker::from_var_builders(g_vb, d_vb, device.clone(), watermark_message)?;
|
||
let wm = ResampledWatermarker::new(inner, generator.config.sample_rate, AUDIOSEAL_RATE);
|
||
generator.set_watermarker(Box::new(wm));
|
||
// Build a second instance for detect-only (the one above moves into the
|
||
// generator). This is cheap: weights are mmap'd, only metadata is duplicated.
|
||
let g_vb2 =
|
||
unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[g], DType::F32, &device) }?;
|
||
let d_vb2 =
|
||
unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[d], DType::F32, &device) }?;
|
||
audioseal_for_detect = Some(AudioSealWatermarker::from_var_builders(
|
||
g_vb2,
|
||
d_vb2,
|
||
device.clone(),
|
||
watermark_message,
|
||
)?);
|
||
println!(
|
||
"[2/5] AudioSeal watermarker installed (message=0x{:04X})",
|
||
watermark_message
|
||
);
|
||
} else {
|
||
println!(
|
||
"[2/5] AudioSeal: skipped (pass both --audioseal-generator and --audioseal-detector to enable)"
|
||
);
|
||
}
|
||
|
||
// -- 3. Generate --------------------------------------------------------
|
||
let opts = GenerateOptions {
|
||
max_audio_ms: cli.max_audio_ms,
|
||
..GenerateOptions::default()
|
||
};
|
||
let post = PostProcess {
|
||
lufs_target: Some(cli.lufs),
|
||
..PostProcess::default()
|
||
};
|
||
let context: Vec<Segment> = Vec::new();
|
||
let t_gen = std::time::Instant::now();
|
||
generator.generate_to_wav(&cli.text, cli.speaker, &context, opts, &post, &cli.out)?;
|
||
let gen_secs = t_gen.elapsed().as_secs_f32();
|
||
println!(
|
||
"[3/5] generated + post-processed{} in {:.2}s -> {}",
|
||
if audioseal_for_detect.is_some() {
|
||
" + watermarked"
|
||
} else {
|
||
""
|
||
},
|
||
gen_secs,
|
||
cli.out.display()
|
||
);
|
||
|
||
// -- 4. Watermark verification round-trip ------------------------------
|
||
if let Some(detector) = audioseal_for_detect.as_ref() {
|
||
let raw = audio_io::load_mono_at_rate(&cli.out, AUDIOSEAL_RATE)?;
|
||
let result = detector.detect(&raw)?;
|
||
let decoded = result.message.unwrap_or(0);
|
||
let bits_match = 16 - (decoded ^ watermark_message).count_ones() as usize;
|
||
println!(
|
||
"[4/5] AudioSeal detect: mean_presence={:.4}, decoded=0x{:04X}, bits={}/16",
|
||
result.mean_presence, decoded, bits_match
|
||
);
|
||
} else {
|
||
println!("[4/5] AudioSeal detect: skipped");
|
||
}
|
||
|
||
// -- 5. WavLM-SV speaker embedding + optional reference scoring -------
|
||
if let Some(wavlm_path) = cli.wavlm_sv.as_ref() {
|
||
let scorer = WavLmSimilarity::load(wavlm_path, &device)?;
|
||
let out_samples = audio_io::load_mono_at_rate(&cli.out, AUDIOSEAL_RATE)?;
|
||
let out_emb = scorer.embed(&out_samples)?;
|
||
println!("[5/5] WavLM-SV embedded output (len={})", out_emb.len());
|
||
if let Some(ref_path) = cli.reference.as_ref() {
|
||
let ref_samples = audio_io::load_mono_at_rate(ref_path, AUDIOSEAL_RATE)?;
|
||
let sim = scorer.score(&out_samples, &ref_samples)?;
|
||
println!(
|
||
" cosine vs reference {}: {:.4} ({})",
|
||
ref_path.display(),
|
||
sim,
|
||
if sim > 0.5 {
|
||
"likely same speaker"
|
||
} else {
|
||
"likely different speakers"
|
||
}
|
||
);
|
||
}
|
||
} else {
|
||
println!("[5/5] WavLM-SV: skipped (pass --wavlm-sv to enable)");
|
||
}
|
||
|
||
println!("== pipeline complete: {} ==", cli.out.display());
|
||
Ok(())
|
||
}
|