Add rtx-csm: Rust-native port of Sesame CSM-1B with LoRA voice cloning

A new model crate at crates/models/rtx-csm implementing end-to-end inference, quantization, and fine-tuning for Sesame's Conversational Speech Model (CSM-1B). Built on candle 0.9 + Kyutai Mimi codec. Key capabilities: - Inference (FP F16 on Metal, F32 on CPU, BF16 on CUDA) - Quantized inference (Q8_0 / Q4_K_M GGUF, ~3x speedup, ~50% memory) - Streaming Mimi decode with proper StreamTensor state machine - In-context voice cloning via SpeakerProfile - Classifier-Free Guidance (Koel-TTS recipe) - Long-form chunked generation with rolling context - Audio post-processing (HPF + declick + EBU R128 LUFS) - Text input normalization (brackets, times, unicode, length caps) - Frame-level repetition guard (loop-escape) - Top-k + top-p sampling - LoRA fine-tuning end-to-end (training + inference, on FP and Q8 bases) - In-process Whisper ASR via whisper-rs (under --features asr) - Standalone TTS HTTP server (Axum) - Bench harness with manifest export + per-prompt WER Phases delivered: quantization, ASR/WER eval, LoRA voice cloning, HTTP service. AudioSeal/WavLM/Unmute remain as documented future work. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
2026-04-25 18:33:57 -07:00
parent 15b62a8f6e
commit 15dd3575d4
42 changed files with 8213 additions and 0 deletions
@@ -0,0 +1,141 @@
+//! In-process Whisper ASR via whisper.cpp bindings (`whisper-rs`).
+//!
+//! Provides word-level transcription of generated audio for in-process WER
+//! scoring inside the bench harness. Requires the `asr` (or `asr-metal`,
+//! `asr-cuda`) feature flag — pulls a C++ build (cmake + clang).
+//!
+//! Default model is `ggerganov/whisper.cpp/ggml-tiny.en.bin` — small (~75 MB)
+//! and English-only, matching CSM-1B's primary language. Larger models can be
+//! selected via `WhisperAsr::load_with_repo`.
+
+#![cfg(feature = "asr")]
+
+use crate::audio_io::TARGET_SAMPLE_RATE;
+use crate::error::{CsmError, Result};
+use rubato::{Resampler, SincFixedIn, SincInterpolationParameters, SincInterpolationType, WindowFunction};
+use std::path::PathBuf;
+use std::sync::Mutex;
+use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};
+
+pub const DEFAULT_REPO: &str = "ggerganov/whisper.cpp";
+pub const DEFAULT_FILE: &str = "ggml-tiny.en.bin";
+pub const WHISPER_SAMPLE_RATE: u32 = 16_000;
+
+pub struct WhisperAsr {
+    ctx: WhisperContext,
+    /// Lock around the mutable state to make `transcribe` thread-safe.
+    state_lock: Mutex<()>,
+    language: Option<String>,
+}
+
+impl WhisperAsr {
+    /// Resolve and load the default tiny.en model from HF, caching via hf-hub.
+    pub fn load_default() -> Result<Self> {
+        Self::load_with_repo(DEFAULT_REPO, DEFAULT_FILE, Some("en"))
+    }
+
+    pub fn load_with_repo(repo: &str, file: &str, language: Option<&str>) -> Result<Self> {
+        let api = hf_hub::api::sync::Api::new()
+            .map_err(|e| CsmError::Other(anyhow::anyhow!("hf-hub init: {e}")))?;
+        let path = api
+            .model(repo.to_string())
+            .get(file)
+            .map_err(|e| CsmError::Other(anyhow::anyhow!("hf-hub get {repo}/{file}: {e}")))?;
+        Self::load(&path, language)
+    }
+
+    pub fn load<P: AsRef<std::path::Path>>(model_path: P, language: Option<&str>) -> Result<Self> {
+        let path = model_path.as_ref();
+        let path_str = path
+            .to_str()
+            .ok_or_else(|| CsmError::Config(format!("non-utf8 model path: {path:?}")))?;
+        tracing::info!("loading whisper model from {}", path_str);
+        let ctx = WhisperContext::new_with_params(path_str, WhisperContextParameters::default())
+            .map_err(|e| CsmError::Other(anyhow::anyhow!("WhisperContext::new: {e}")))?;
+        Ok(Self {
+            ctx,
+            state_lock: Mutex::new(()),
+            language: language.map(|s| s.to_string()),
+        })
+    }
+
+    /// Transcribe 24 kHz mono f32 (CSM's native rate). Resamples internally to
+    /// 16 kHz before invoking Whisper.
+    pub fn transcribe_24k(&self, samples_24k: &[f32]) -> Result<String> {
+        if samples_24k.is_empty() {
+            return Ok(String::new());
+        }
+        let samples_16k = resample_to_16k(samples_24k)?;
+        self.transcribe_16k(&samples_16k)
+    }
+
+    /// Transcribe samples that are already 16 kHz mono f32.
+    pub fn transcribe_16k(&self, samples_16k: &[f32]) -> Result<String> {
+        let _g = self.state_lock.lock().expect("state lock poisoned");
+        let mut state = self
+            .ctx
+            .create_state()
+            .map_err(|e| CsmError::Other(anyhow::anyhow!("create_state: {e}")))?;
+        let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
+        if let Some(lang) = self.language.as_deref() {
+            params.set_language(Some(lang));
+        }
+        params.set_print_progress(false);
+        params.set_print_realtime(false);
+        params.set_print_special(false);
+        params.set_print_timestamps(false);
+        state
+            .full(params, samples_16k)
+            .map_err(|e| CsmError::Other(anyhow::anyhow!("whisper full: {e}")))?;
+        let mut out = String::new();
+        for segment in state.as_iter() {
+            // The iterator yields segments whose Display impl returns the text
+            // with invalid-UTF8 mapped to U+FFFD. That's what we want.
+            out.push_str(&segment.to_string());
+        }
+        Ok(out.trim().to_string())
+    }
+}
+
+fn resample_to_16k(samples_24k: &[f32]) -> Result<Vec<f32>> {
+    let src_rate = TARGET_SAMPLE_RATE as f64;
+    let dst_rate = WHISPER_SAMPLE_RATE as f64;
+    let chunk = 1024usize;
+    let params = SincInterpolationParameters {
+        sinc_len: 256,
+        f_cutoff: 0.95,
+        interpolation: SincInterpolationType::Linear,
+        oversampling_factor: 256,
+        window: WindowFunction::BlackmanHarris2,
+    };
+    let mut resampler = SincFixedIn::<f32>::new(dst_rate / src_rate, 2.0, params, chunk, 1)
+        .map_err(|e| CsmError::Rubato(e.to_string()))?;
+    let mut out = Vec::with_capacity(
+        ((samples_24k.len() as f64 * dst_rate / src_rate).ceil()) as usize + chunk,
+    );
+    let mut pos = 0usize;
+    while pos + chunk <= samples_24k.len() {
+        let frame_in = vec![samples_24k[pos..pos + chunk].to_vec()];
+        let frame_out = resampler
+            .process(&frame_in, None)
+            .map_err(|e| CsmError::Rubato(e.to_string()))?;
+        out.extend_from_slice(&frame_out[0]);
+        pos += chunk;
+    }
+    if pos < samples_24k.len() {
+        let mut tail = samples_24k[pos..].to_vec();
+        tail.resize(chunk, 0.0);
+        let frame_out = resampler
+            .process(&[tail], None)
+            .map_err(|e| CsmError::Rubato(e.to_string()))?;
+        let kept = ((samples_24k.len() - pos) as f64 * dst_rate / src_rate).round() as usize;
+        out.extend_from_slice(&frame_out[0][..kept.min(frame_out[0].len())]);
+    }
+    Ok(out)
+}
+
+// PathBuf re-export so docs can reference it without conditional imports.
+#[allow(dead_code)]
+fn _path_export() -> PathBuf {
+    PathBuf::new()
+}