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]>
94 lines
3.4 KiB
Rust
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(())
|
|
}
|