osobhandClaude Opus 4.7
f4aaee2f1b
Single CLI ties together every capability shipped this session: text -> CSM-1B (with optional LoRA) -> post-process (HPF/declick/LUFS) -> AudioSeal watermark embed -> AudioSeal detect verify -> WavLM-SV speaker embedding + optional reference scoring. Verified on Metal: 4s speech generated + watermarked + detected (mean_presence=0.9999, 16/16 bits decoded) + 512-d speaker embedding extracted in ~30s. Cross-content same-speaker cosine sits around 0.49 vs 0.998 for same-content same-speaker — suggests the WavLM-SV port may leak content into the speaker embedding more than the HF reference. Phase 5d numerical parity work (Python sidecar comparison) would tighten this. This is the canonical usage example for downstream callers (clawsample-csm etc.) — copy the structure. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
247 lines
8.8 KiB
Rust
247 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::{
|
||
audio_io,
|
||
audioseal::AudioSealWatermarker,
|
||
speaker_sim::{SpeakerSimilarity, WavLmSimilarity},
|
||
watermark::ResampledWatermarker,
|
||
GenerateOptions, Generator, PostProcess, Segment,
|
||
};
|
||
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(())
|
||
}
|