Files
rustytorch/crates/models/rtx-csm/tests/unit.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

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]);
}