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]>
88 lines
3.0 KiB
Rust
88 lines
3.0 KiB
Rust
//! Unit tests for the weight-free code paths: pure math, formatting, and the
|
|
//! audio I/O WAV round-trip. These run on every `cargo test` — no HF auth or
|
|
//! model downloads required.
|
|
|
|
use approx::assert_relative_eq;
|
|
use rtx_csm::{
|
|
audio_io::{load_mono_24k, write_wav_24k_mono, TARGET_SAMPLE_RATE},
|
|
config::ModelConfig,
|
|
util::{all_zero_codebooks, rms},
|
|
};
|
|
|
|
#[test]
|
|
fn rms_zero_on_identical_signals() {
|
|
let a = vec![0.1, -0.2, 0.3, -0.4, 0.5];
|
|
assert_relative_eq!(rms(&a, &a), 0.0, epsilon = 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn rms_matches_hand_calc() {
|
|
let a = vec![0.0, 1.0, 2.0, 3.0];
|
|
let b = vec![0.0, 0.0, 0.0, 0.0];
|
|
// sum_sq = 0 + 1 + 4 + 9 = 14, n=4, rms = sqrt(14/4) ≈ 1.8708286
|
|
assert_relative_eq!(rms(&a, &b), (14.0f32 / 4.0).sqrt(), epsilon = 1e-5);
|
|
}
|
|
|
|
#[test]
|
|
fn all_zero_codebooks_true_only_on_all_zero() {
|
|
use candle_core::{Device, Tensor};
|
|
let dev = Device::Cpu;
|
|
let zeros = Tensor::zeros((1, 32), candle_core::DType::I64, &dev).unwrap();
|
|
assert!(all_zero_codebooks(&zeros).unwrap());
|
|
|
|
// One non-zero → false.
|
|
let mut v = vec![0i64; 32];
|
|
v[7] = 42;
|
|
let nonzero = Tensor::from_vec(v, (1, 32), &dev).unwrap();
|
|
assert!(!all_zero_codebooks(&nonzero).unwrap());
|
|
}
|
|
|
|
#[test]
|
|
fn speaker_formatting_brackets_id() {
|
|
use rtx_csm::tokenizer::CsmTokenizer;
|
|
assert_eq!(CsmTokenizer::format_segment(0, "hello"), "[0]hello");
|
|
assert_eq!(CsmTokenizer::format_segment(1, "world"), "[1]world");
|
|
assert_eq!(CsmTokenizer::format_segment(0, ""), "[0]");
|
|
}
|
|
|
|
#[test]
|
|
fn config_csm_1b_defaults() {
|
|
let cfg = ModelConfig::csm_1b();
|
|
assert_eq!(cfg.text_vocab_size, 128_256);
|
|
assert_eq!(cfg.audio_vocab_size, 2051);
|
|
assert_eq!(cfg.audio_num_codebooks, 32);
|
|
assert_eq!(cfg.sample_rate, 24_000);
|
|
assert_eq!(cfg.frame_rate_hz, 12.5);
|
|
assert_relative_eq!(cfg.frame_duration_ms(), 80.0, epsilon = 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn wav_write_read_roundtrip() {
|
|
let tmp = tempfile::NamedTempFile::with_suffix(".wav").unwrap();
|
|
// Synth a 0.1-second sine wave at 440 Hz so the test file is ~4.8 KB.
|
|
let n = (TARGET_SAMPLE_RATE as f32 * 0.1) as usize;
|
|
let input: Vec<f32> = (0..n)
|
|
.map(|i| 0.5 * (2.0 * std::f32::consts::PI * 440.0 * i as f32 / TARGET_SAMPLE_RATE as f32).sin())
|
|
.collect();
|
|
|
|
write_wav_24k_mono(tmp.path(), &input).unwrap();
|
|
let read_back = load_mono_24k(tmp.path()).unwrap();
|
|
|
|
// hound writes 16-bit PCM so we lose some precision; expect same length and
|
|
// small per-sample error.
|
|
assert_eq!(read_back.len(), input.len());
|
|
let err = rms(&input, &read_back);
|
|
assert!(err < 1e-3, "wav round-trip rms {err} too high");
|
|
}
|
|
|
|
#[test]
|
|
fn empty_prompt_has_shape_one_zero_cb_plus_one() {
|
|
use candle_core::{DType, Device};
|
|
// Directly validate that a zero-length audio tensor has the right shape —
|
|
// catches regressions in the empty-audio branch of prompt.rs.
|
|
let dev = Device::Cpu;
|
|
let cb = 32;
|
|
let t = candle_core::Tensor::zeros((1, 0, cb + 1), DType::U32, &dev).unwrap();
|
|
assert_eq!(t.dims(), &[1, 0, 33]);
|
|
}
|