Files
rustytorch/crates/models/rtx-csm/tests/mimi_roundtrip.rs
T
osobhandClaude Opus 4.7 15dd3575d4 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

94 lines
3.4 KiB
Rust

//! Step A: Mimi codec round-trip parity check.
//!
//! Synthesizes a pseudo-speech signal at 24 kHz, encodes through Mimi to
//! discrete codes (32 codebooks), decodes back to waveform, asserts the RMS
//! reconstruction error is within bound.
//!
//! Marked `#[ignore]` because it requires ~300 MB of weights from
//! `kyutai/mimi`. Run manually:
//!
//! ```
//! HF_HUB_ENABLE_HF_TRANSFER=0 cargo test -p rtx-csm --test mimi_roundtrip --release -- --ignored --nocapture
//! ```
use rtx_csm::{
audio_io::{write_wav_24k_mono, TARGET_SAMPLE_RATE},
error::Result,
hub,
mimi::{Mimi, NUM_CODEBOOKS},
util::{pick_device, rms},
};
use std::f32::consts::TAU;
/// Build a speech-shaped 24 kHz mono signal: voiced fundamental + first three
/// formants, amplitude-modulated to simulate syllable rate. Roughly 3 seconds.
fn synth_speech_like(seconds: f32) -> Vec<f32> {
let n = (seconds * TARGET_SAMPLE_RATE as f32) as usize;
let mut out = Vec::with_capacity(n);
let f0 = 130.0_f32; // male-ish fundamental
let formants = [700.0_f32, 1220.0, 2600.0];
let formant_gains = [1.0_f32, 0.6, 0.3];
let syllable_rate = 4.0_f32; // Hz
for i in 0..n {
let t = i as f32 / TARGET_SAMPLE_RATE as f32;
// amplitude envelope (avoid 0 for stability)
let env = 0.5 * (1.0 + (TAU * syllable_rate * t).sin()).max(0.05);
let voiced = (TAU * f0 * t).sin();
let mut formant_sum = 0.0;
for (f, g) in formants.iter().zip(formant_gains.iter()) {
formant_sum += g * (TAU * f * t).sin();
}
let s = 0.5 * env * (0.6 * voiced + 0.4 * formant_sum / formant_gains.iter().sum::<f32>());
out.push(s);
}
out
}
#[test]
#[ignore]
fn mimi_roundtrip_rms_bound() -> Result<()> {
let _ = tracing_subscriber::fmt::try_init();
let device = pick_device()?;
println!("device: {device:?}");
let mimi_weights = hub::resolve_mimi()?;
println!("mimi weights: {}", mimi_weights.display());
let mut mimi = Mimi::load(&mimi_weights, &device)?;
let original = synth_speech_like(3.0);
println!("input samples: {} (~{:.1}s)", original.len(), original.len() as f32 / TARGET_SAMPLE_RATE as f32);
let codes = mimi.encode(&original)?;
let dims = codes.dims();
println!("codes shape: {dims:?}");
assert_eq!(dims.len(), 3, "codes must be (B, C, T)");
assert_eq!(dims[0], 1, "batch=1");
assert_eq!(dims[1], NUM_CODEBOOKS, "codebooks=32");
let decoded = mimi.decode(&codes)?;
println!("decoded samples: {}", decoded.len());
// Truncate to common length — Mimi may pad to a multiple of 1920 (24kHz/12.5Hz).
let n = original.len().min(decoded.len());
let err = rms(&original[..n], &decoded[..n]);
println!("rms error: {err:.5}");
// Save artifacts for ear-test debugging when -- --nocapture is used.
if let Ok(dir) = std::env::var("CSM_TEST_OUT") {
let _ = std::fs::create_dir_all(&dir);
write_wav_24k_mono(format!("{dir}/mimi_in.wav"), &original)?;
write_wav_24k_mono(format!("{dir}/mimi_out.wav"), &decoded[..n])?;
println!("artifacts written to {dir}");
}
// Synthesized non-speech: looser bound than the 0.05 target for real speech.
// Real-speech bound will be re-checked when a fixture WAV is added in Step B.
assert!(
err < 0.20,
"Mimi round-trip RMS {err} exceeds 0.20 bound on synthesized signal"
);
Ok(())
}