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,172 @@
+//! Frame-level repetition detection ("loop-escape").
+//!
+//! CSM's known pathological mode is **syllable looping** — the AR decoder
+//! gets stuck in a low-entropy attractor and emits the same 32-codebook
+//! frame over and over. Users hear it as a syllable repeating forever.
+//!
+//! The ideal fix is per-codebook logit blocking on recent frame
+//! fingerprints, but the actual sampling happens *inside* candle's
+//! `csm::Model::generate_frame` and we can't hook between codebook
+//! samples without forking the upstream module. So instead we detect the
+//! loop *after* the frame is emitted and short-circuit generation.
+//!
+//! This catches both:
+//!   - period-1 loops (frame_k == frame_{k-1}, ...)
+//!   - higher-period loops where a small set of frames cycles
+//!
+//! Detection rule: within the last `window` emitted frames, if any single
+//! fingerprint occurs `max_repeats` or more times, we declare a loop.
+
+use std::collections::{HashMap, VecDeque};
+
+#[derive(Debug, Clone, Copy)]
+pub struct RepetitionConfig {
+    /// How many recent frames to track. 16 frames = 1.28s at 12.5 Hz.
+    pub window: usize,
+    /// A single fingerprint occurring this many times in the window triggers
+    /// a loop break. 4 is a safe default — natural speech can have 2-3
+    /// repeated frames at a sustained vowel without being pathological.
+    pub max_repeats: usize,
+}
+
+impl Default for RepetitionConfig {
+    fn default() -> Self {
+        Self {
+            window: 16,
+            max_repeats: 4,
+        }
+    }
+}
+
+/// Tracks recent frame fingerprints. Cheap (a small ring buffer of u64).
+pub struct RepetitionGuard {
+    cfg: RepetitionConfig,
+    recent: VecDeque<u64>,
+}
+
+impl RepetitionGuard {
+    pub fn new(cfg: RepetitionConfig) -> Self {
+        Self {
+            cfg,
+            recent: VecDeque::with_capacity(cfg.window),
+        }
+    }
+
+    pub fn reset(&mut self) {
+        self.recent.clear();
+    }
+
+    /// Record a frame and return true if a loop is detected. Frame is the
+    /// `Vec<u32>` of length num_codebooks returned by `generate_frame`.
+    pub fn observe(&mut self, frame: &[u32]) -> bool {
+        let fp = fingerprint(frame);
+        if self.recent.len() == self.cfg.window {
+            self.recent.pop_front();
+        }
+        self.recent.push_back(fp);
+        // Count occurrences only when the window has enough frames to make
+        // the threshold meaningful — avoids spurious early triggers.
+        if self.recent.len() < self.cfg.max_repeats {
+            return false;
+        }
+        let mut counts: HashMap<u64, usize> = HashMap::new();
+        let mut max_count = 0usize;
+        for h in self.recent.iter() {
+            let c = counts.entry(*h).or_insert(0);
+            *c += 1;
+            if *c > max_count {
+                max_count = *c;
+            }
+        }
+        max_count >= self.cfg.max_repeats
+    }
+}
+
+/// FNV-1a 64-bit hash over the codebook ids. Fast, avoids extra deps.
+fn fingerprint(frame: &[u32]) -> u64 {
+    const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
+    const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
+    let mut h = FNV_OFFSET;
+    for v in frame {
+        let bytes = v.to_le_bytes();
+        for b in bytes {
+            h ^= b as u64;
+            h = h.wrapping_mul(FNV_PRIME);
+        }
+    }
+    h
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn no_repeats_no_trigger() {
+        let mut g = RepetitionGuard::new(RepetitionConfig::default());
+        for i in 0..20u32 {
+            let frame: Vec<u32> = (0..32).map(|c| i.wrapping_add(c)).collect();
+            assert!(!g.observe(&frame), "false trigger at frame {i}");
+        }
+    }
+
+    #[test]
+    fn period_1_loop_trips_after_max_repeats() {
+        let mut g = RepetitionGuard::new(RepetitionConfig {
+            window: 16,
+            max_repeats: 4,
+        });
+        let frame = vec![42u32; 32];
+        // Three identical frames must NOT trigger.
+        assert!(!g.observe(&frame));
+        assert!(!g.observe(&frame));
+        assert!(!g.observe(&frame));
+        // Fourth one trips the guard.
+        assert!(g.observe(&frame));
+    }
+
+    #[test]
+    fn period_2_alternating_loop_trips() {
+        let mut g = RepetitionGuard::new(RepetitionConfig {
+            window: 16,
+            max_repeats: 4,
+        });
+        let a = vec![1u32; 32];
+        let b = vec![2u32; 32];
+        let mut tripped_at = None;
+        for i in 0..16 {
+            let f = if i % 2 == 0 { &a } else { &b };
+            if g.observe(f) {
+                tripped_at = Some(i);
+                break;
+            }
+        }
+        assert!(tripped_at.is_some(), "alternating loop never tripped");
+    }
+
+    #[test]
+    fn fingerprint_distinguishes_neighbors() {
+        let a: Vec<u32> = (0..32).collect();
+        let mut b = a.clone();
+        b[5] = 999;
+        assert_ne!(fingerprint(&a), fingerprint(&b));
+    }
+
+    #[test]
+    fn reset_clears_history() {
+        let mut g = RepetitionGuard::new(RepetitionConfig {
+            window: 16,
+            max_repeats: 4,
+        });
+        let frame = vec![42u32; 32];
+        for _ in 0..5 {
+            g.observe(&frame);
+        }
+        g.reset();
+        // After reset, we need 4 more observations to trip again.
+        assert!(!g.observe(&frame));
+        assert!(!g.observe(&frame));
+        assert!(!g.observe(&frame));
+        assert!(g.observe(&frame));
+    }
+}