rustytorch/crates/models/rtx-csm/examples/diarize.rs

//! Diarize a wav file: produce JSON of (start_s, end_s, speaker) segments.
//!
//! Composed entirely from in-crate parts (Silero V5 VAD + WavLM-SV +
//! agglomerative clustering — see `src/diarize.rs`). No external runtime
//! deps. Suitable as a data-prep CLI for the Phase 12 LoRA training stack.
//!
//! Usage:
//!   cargo run -p rtx-csm --release --features metal --example diarize -- \
//!       --in /path/to/audio.wav --out /tmp/segments.json
//!
//! Optional `--n-speakers k` forces exactly k clusters (skip the auto
//! threshold). Pipe the JSON into `jq` to slice up the audio with
//! `ffmpeg`/`sox` for per-speaker training.

use anyhow::Context;
use anyhow::Result;
use candle_core::DType;
use clap::Parser;
use rtx_csm::Generator;
use rtx_csm::audio_io;
use rtx_csm::diarize::{DiarizationConfig, Diarizer, SAMPLE_RATE};
use rtx_csm::silero_vad::SileroVad;
use rtx_csm::wavlm_sv::WavLmSv;
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Cli {
    /// Input wav file (any sample rate; resampled internally to 16 kHz).
    #[arg(long = "in")]
    input: PathBuf,

    /// Output JSON path.
    #[arg(long)]
    out: PathBuf,

    /// Converted WavLM-SV safetensors (output of `examples/wavlm_sv_convert`).
    #[arg(long)]
    wavlm_sv_weights: PathBuf,

    /// Force the clustering algorithm to produce exactly this many speakers.
    /// When omitted, auto-detect via the cosine-distance threshold.
    #[arg(long)]
    n_speakers: Option<usize>,

    /// VAD probability threshold (0–1). 0.5 is the Silero V5 default.
    #[arg(long, default_value_t = 0.5)]
    vad_threshold: f32,

    /// Cosine-distance threshold above which two clusters do NOT merge.
    /// Lower = more clusters. 0.5 is a literature default for WavLM-SV.
    #[arg(long, default_value_t = 0.5)]
    cluster_threshold: f32,

    /// WavLM-SV embedding window length in seconds.
    #[arg(long, default_value_t = 2.0)]
    window_s: f32,

    /// Hop between consecutive embedding windows in seconds.
    #[arg(long, default_value_t = 1.0)]
    hop_s: f32,

    /// Drop output segments shorter than this (seconds) after the merge step.
    #[arg(long, default_value_t = 0.5)]
    min_segment_s: f32,

    #[arg(long)]
    cpu: bool,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();

    let device = if cli.cpu {
        candle_core::Device::Cpu
    } else {
        Generator::default_device()?
    };
    println!("device: {device:?}");

    // Load 16 kHz audio (audio_io resamples internally if needed).
    let audio = audio_io::load_mono_at_rate(&cli.input, SAMPLE_RATE)?;
    println!(
        "loaded {} samples ({:.2}s at {} Hz)",
        audio.len(),
        audio.len() as f32 / SAMPLE_RATE as f32,
        SAMPLE_RATE,
    );

    // Build the two backbones.
    let vad_path = rtx_csm::silero_vad::ensure_default_weights()?;
    let vad = SileroVad::load_from_file(&vad_path, &device)?;
    println!("loaded Silero V5 VAD");

    let vb = unsafe {
        candle_nn::VarBuilder::from_mmaped_safetensors(
            &[&cli.wavlm_sv_weights],
            DType::F32,
            &device,
        )
    }
    .context("opening WavLM-SV safetensors")?;
    let sv = WavLmSv::new(vb).context("WavLmSv::new")?;
    println!("loaded WavLM-SV from {}", cli.wavlm_sv_weights.display());

    let mut diarizer = Diarizer::new(vad, sv, device);
    let cfg = DiarizationConfig {
        vad_threshold: cli.vad_threshold,
        cluster_threshold: cli.cluster_threshold,
        window_s: cli.window_s,
        hop_s: cli.hop_s,
        n_speakers: cli.n_speakers,
        min_segment_s: cli.min_segment_s,
        ..DiarizationConfig::default()
    };

    let t0 = std::time::Instant::now();
    let segments = diarizer.diarize(&audio, &cfg)?;
    let elapsed_ms = t0.elapsed().as_millis();
    let realtime_factor =
        (elapsed_ms as f32 / 1000.0) / (audio.len() as f32 / SAMPLE_RATE as f32).max(0.001);
    println!(
        "diarized {} segments in {} ms ({:.3}× realtime)",
        segments.len(),
        elapsed_ms,
        realtime_factor,
    );
    let n_speakers: std::collections::HashSet<_> = segments.iter().map(|s| s.speaker).collect();
    println!("detected {} unique speaker(s)", n_speakers.len());
    for s in &segments {
        println!(
            "  speaker {} : {:.2}s → {:.2}s",
            s.speaker, s.start_s, s.end_s
        );
    }

    let json = serde_json::to_string_pretty(&segments)?;
    std::fs::write(&cli.out, json)?;
    println!("\n✓ wrote {}", cli.out.display());
    Ok(())
}