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]>
This commit is contained in:
osobh
@@ -0,0 +1,354 @@
|
||||
//! Multi-epoch LoRA training infrastructure.
|
||||
//!
|
||||
//! Provides:
|
||||
//! - [`TrainingExample`] — one `(transcript, audio)` pair, with the audio
|
||||
//! pre-encoded to Mimi codes so each training step doesn't re-encode.
|
||||
//! - [`TrainingDataset`] — scan a directory of `(*.wav, *.txt)` pairs and
|
||||
//! prepare a list of [`TrainingExample`]s.
|
||||
//! - [`Trainer`] — wraps a [`Generator`] (with LoRA already injected) +
|
||||
//! AdamW + LR schedule + gradient clipping. `train_one_epoch` and
|
||||
//! `train_n_epochs` drive the optimizer.
|
||||
//! - [`save_lora_adapter`] / [`load_lora_adapter`] — round-trip the LoRA
|
||||
//! A,B matrices through safetensors.
|
||||
|
||||
use crate::audio_io;
|
||||
use crate::error::{CsmError, Result};
|
||||
use crate::generator::Generator;
|
||||
use crate::prompt::{build_prompt, Segment};
|
||||
use candle_core::Tensor;
|
||||
use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarMap};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// One training example: a transcript paired with the Mimi-encoded audio
|
||||
/// codes for that utterance. Pre-encoding avoids running Mimi on every
|
||||
/// training step.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TrainingExample {
|
||||
pub speaker: u32,
|
||||
pub text: String,
|
||||
/// Per-frame target codes, shape `(num_frames, num_codebooks=32)`.
|
||||
pub frame_codes: Vec<Vec<u32>>,
|
||||
/// Path the audio came from, for logging.
|
||||
pub source: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl TrainingExample {
|
||||
/// Build from raw 24 kHz mono samples by Mimi-encoding the audio.
|
||||
pub fn from_audio(
|
||||
speaker: u32,
|
||||
text: impl Into<String>,
|
||||
audio_24k: &[f32],
|
||||
generator: &mut Generator,
|
||||
) -> Result<Self> {
|
||||
let codes = generator.mimi.encode(audio_24k)?;
|
||||
let (_b, num_codebooks, num_frames) = codes.dims3()?;
|
||||
let mut frame_codes = Vec::with_capacity(num_frames);
|
||||
for frame in 0..num_frames {
|
||||
let row = codes
|
||||
.narrow(2, frame, 1)?
|
||||
.squeeze(2)?
|
||||
.flatten_all()?
|
||||
.to_vec1::<u32>()?;
|
||||
assert_eq!(row.len(), num_codebooks);
|
||||
frame_codes.push(row);
|
||||
}
|
||||
Ok(Self {
|
||||
speaker,
|
||||
text: text.into(),
|
||||
frame_codes,
|
||||
source: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// A dataset of training examples loaded from disk.
|
||||
pub struct TrainingDataset {
|
||||
pub examples: Vec<TrainingExample>,
|
||||
}
|
||||
|
||||
impl TrainingDataset {
|
||||
/// Scan a directory for `*.wav` files. For each `foo.wav` look up the
|
||||
/// matching transcript at `foo.txt`. Skip files that don't have a pair.
|
||||
/// Mimi-encode each audio file once at load time.
|
||||
pub fn load_from_dir<P: AsRef<Path>>(
|
||||
dir: P,
|
||||
speaker: u32,
|
||||
generator: &mut Generator,
|
||||
) -> Result<Self> {
|
||||
let dir = dir.as_ref();
|
||||
let mut examples = Vec::new();
|
||||
for entry in std::fs::read_dir(dir)? {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
let ext = path.extension().and_then(|s| s.to_str()).unwrap_or("");
|
||||
if ext != "wav" {
|
||||
continue;
|
||||
}
|
||||
let txt_path = path.with_extension("txt");
|
||||
if !txt_path.exists() {
|
||||
tracing::warn!(
|
||||
"skipping {}: no matching .txt transcript",
|
||||
path.display()
|
||||
);
|
||||
continue;
|
||||
}
|
||||
let text = std::fs::read_to_string(&txt_path)?.trim().to_string();
|
||||
if text.is_empty() {
|
||||
tracing::warn!("skipping {}: empty transcript", path.display());
|
||||
continue;
|
||||
}
|
||||
let audio = audio_io::load_mono_24k(&path)?;
|
||||
let mut ex = TrainingExample::from_audio(speaker, text, &audio, generator)?;
|
||||
ex.source = Some(path);
|
||||
tracing::info!(
|
||||
"loaded example: {} chars, {} frames",
|
||||
ex.text.len(),
|
||||
ex.frame_codes.len()
|
||||
);
|
||||
examples.push(ex);
|
||||
}
|
||||
if examples.is_empty() {
|
||||
return Err(CsmError::Config(format!(
|
||||
"no (.wav, .txt) pairs found in {}",
|
||||
dir.display()
|
||||
)));
|
||||
}
|
||||
Ok(Self { examples })
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.examples.len()
|
||||
}
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.examples.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// Cosine LR schedule with linear warmup. `step` is 0-indexed.
|
||||
pub fn cosine_lr(step: usize, total_steps: usize, warmup: usize, peak: f64, end: f64) -> f64 {
|
||||
if step < warmup {
|
||||
peak * (step as f64 + 1.0) / (warmup.max(1) as f64)
|
||||
} else {
|
||||
let progress = ((step - warmup) as f64) / ((total_steps - warmup).max(1) as f64);
|
||||
let cos = (1.0 + (progress * std::f64::consts::PI).cos()) * 0.5;
|
||||
end + (peak - end) * cos
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TrainingConfig {
|
||||
pub epochs: usize,
|
||||
pub peak_lr: f64,
|
||||
pub end_lr: f64,
|
||||
pub warmup_steps: usize,
|
||||
pub grad_clip: Option<f64>,
|
||||
/// Number of frames sampled per example per step (random subset of frames).
|
||||
/// Smaller = faster step but noisier; larger = slower but more stable.
|
||||
pub frames_per_step: usize,
|
||||
pub seed: u64,
|
||||
}
|
||||
|
||||
impl Default for TrainingConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
epochs: 5,
|
||||
peak_lr: 5e-4,
|
||||
end_lr: 1e-5,
|
||||
warmup_steps: 16,
|
||||
grad_clip: Some(1.0),
|
||||
frames_per_step: 4,
|
||||
seed: 42,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Trainer<'a> {
|
||||
pub generator: &'a mut Generator,
|
||||
pub vm: &'a VarMap,
|
||||
pub dataset: &'a TrainingDataset,
|
||||
pub config: TrainingConfig,
|
||||
}
|
||||
|
||||
impl<'a> Trainer<'a> {
|
||||
pub fn new(
|
||||
generator: &'a mut Generator,
|
||||
vm: &'a VarMap,
|
||||
dataset: &'a TrainingDataset,
|
||||
config: TrainingConfig,
|
||||
) -> Self {
|
||||
Self {
|
||||
generator,
|
||||
vm,
|
||||
dataset,
|
||||
config,
|
||||
}
|
||||
}
|
||||
|
||||
/// Run a full training schedule: `epochs * dataset_size * frames_per_step`
|
||||
/// gradient steps with the cosine LR schedule. Returns the per-step loss
|
||||
/// trace so callers can plot / log.
|
||||
pub fn train(&mut self) -> Result<Vec<f32>> {
|
||||
use rand::rngs::StdRng;
|
||||
use rand::seq::SliceRandom;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
let total_steps =
|
||||
self.config.epochs * self.dataset.len() * self.config.frames_per_step.max(1);
|
||||
if total_steps == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
tracing::info!(
|
||||
"training: epochs={} dataset={} fps={} → total_steps={} peak_lr={}",
|
||||
self.config.epochs,
|
||||
self.dataset.len(),
|
||||
self.config.frames_per_step,
|
||||
total_steps,
|
||||
self.config.peak_lr
|
||||
);
|
||||
|
||||
let mut rng = StdRng::seed_from_u64(self.config.seed);
|
||||
let mut optim = AdamW::new(
|
||||
self.vm.all_vars(),
|
||||
ParamsAdamW {
|
||||
lr: self.config.peak_lr,
|
||||
..ParamsAdamW::default()
|
||||
},
|
||||
)?;
|
||||
let mut losses = Vec::with_capacity(total_steps);
|
||||
let mut step = 0usize;
|
||||
|
||||
for epoch in 0..self.config.epochs {
|
||||
let mut order: Vec<usize> = (0..self.dataset.len()).collect();
|
||||
order.shuffle(&mut rng);
|
||||
for &idx in order.iter() {
|
||||
let ex = &self.dataset.examples[idx];
|
||||
if ex.frame_codes.is_empty() {
|
||||
continue;
|
||||
}
|
||||
// Build the prompt once for this example.
|
||||
let current = Segment::new_text(ex.speaker, &ex.text);
|
||||
let prompt = build_prompt(
|
||||
&[],
|
||||
¤t,
|
||||
&self.generator.model,
|
||||
&mut self.generator.mimi,
|
||||
&self.generator.tokenizer,
|
||||
)?;
|
||||
let prompt_len = prompt.tokens.dim(1)?;
|
||||
|
||||
for _ in 0..self.config.frames_per_step.max(1) {
|
||||
let lr = cosine_lr(
|
||||
step,
|
||||
total_steps,
|
||||
self.config.warmup_steps,
|
||||
self.config.peak_lr,
|
||||
self.config.end_lr,
|
||||
);
|
||||
optim.set_learning_rate(lr);
|
||||
|
||||
// Pick a random frame from this example as the target.
|
||||
let frame_idx = rng.gen_range(0..ex.frame_codes.len());
|
||||
let target = &ex.frame_codes[frame_idx];
|
||||
|
||||
self.generator.model.clear_kv_cache();
|
||||
let loss = self.generator.model.inner.forward_loss(
|
||||
&prompt.tokens,
|
||||
&prompt.mask,
|
||||
0,
|
||||
target,
|
||||
)?;
|
||||
let loss_val = loss.to_scalar::<f32>()?;
|
||||
losses.push(loss_val);
|
||||
|
||||
let mut grads = loss.backward()?;
|
||||
if let Some(clip) = self.config.grad_clip {
|
||||
clip_grads(self.vm, &mut grads, clip as f32)?;
|
||||
}
|
||||
optim.step(&grads)?;
|
||||
self.generator.model.inner.refresh_lora(self.vm)?;
|
||||
|
||||
step += 1;
|
||||
if step % 10 == 0 || step == total_steps - 1 {
|
||||
tracing::info!(
|
||||
" step {step:>4}/{total_steps} epoch {epoch} ex {idx} frame {frame_idx} \
|
||||
prompt_len={prompt_len} lr={lr:.2e} loss={loss_val:.4}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(losses)
|
||||
}
|
||||
}
|
||||
|
||||
/// Global L2 norm gradient clipping. Walks all Vars in the VarMap, computes
|
||||
/// `||g||_2` across all gradients, and scales every gradient by
|
||||
/// `min(1, max_norm / ||g||_2)`.
|
||||
fn clip_grads(
|
||||
vm: &VarMap,
|
||||
grads: &mut candle_core::backprop::GradStore,
|
||||
max_norm: f32,
|
||||
) -> Result<()> {
|
||||
let mut total_sq: f32 = 0.0;
|
||||
for v in vm.all_vars() {
|
||||
if let Some(g) = grads.get(&v) {
|
||||
let s = g
|
||||
.flatten_all()?
|
||||
.to_dtype(candle_core::DType::F32)?
|
||||
.sqr()?
|
||||
.sum_all()?
|
||||
.to_scalar::<f32>()?;
|
||||
total_sq += s;
|
||||
}
|
||||
}
|
||||
let norm = total_sq.sqrt();
|
||||
if norm > max_norm {
|
||||
let scale = max_norm / norm;
|
||||
for v in vm.all_vars() {
|
||||
if let Some(g) = grads.get(&v) {
|
||||
let scaled = (g * scale as f64)?;
|
||||
grads.insert(v.as_tensor(), scaled);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Save the LoRA adapter parameters (A,B for every layer) as a safetensors
|
||||
/// file. Reload via `load_lora_adapter`.
|
||||
pub fn save_lora_adapter<P: AsRef<Path>>(vm: &VarMap, out: P) -> Result<()> {
|
||||
use std::collections::HashMap;
|
||||
let vars = vm.data().lock().unwrap();
|
||||
let mut tensors: HashMap<String, Tensor> = HashMap::new();
|
||||
for (name, var) in vars.iter() {
|
||||
tensors.insert(name.clone(), var.as_tensor().clone());
|
||||
}
|
||||
drop(vars);
|
||||
candle_core::safetensors::save(&tensors, out.as_ref())
|
||||
.map_err(|e| CsmError::Other(anyhow::anyhow!("safetensors save: {e}")))?;
|
||||
tracing::info!("saved {} LoRA tensors → {}", tensors.len(), out.as_ref().display());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load LoRA adapter weights from a safetensors file, copying into the
|
||||
/// matching Vars in the VarMap. Each Var must already exist (i.e. the model
|
||||
/// must have been constructed with `add_lora_to_backbone` before calling).
|
||||
pub fn load_lora_adapter<P: AsRef<Path>>(
|
||||
vm: &VarMap,
|
||||
path: P,
|
||||
device: &candle_core::Device,
|
||||
) -> Result<()> {
|
||||
let tensors = candle_core::safetensors::load(path.as_ref(), device)?;
|
||||
let vars = vm.data().lock().unwrap();
|
||||
let mut count = 0usize;
|
||||
for (name, t) in tensors.iter() {
|
||||
if let Some(var) = vars.get(name) {
|
||||
var.set(t)?;
|
||||
count += 1;
|
||||
} else {
|
||||
tracing::warn!("safetensors has tensor `{name}` but VarMap has no matching Var");
|
||||
}
|
||||
}
|
||||
tracing::info!("loaded {count} LoRA tensors from {}", path.as_ref().display());
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user