rustytorch/crates/models/rtx-csm/src/speaker.rs

//! `SpeakerProfile` — first-class voice-cloning API.
//!
//! CSM's strongest perceptual lever is **in-context audio conditioning**:
//! given 30 s – 3 min of reference audio (with transcripts) for a target
//! speaker, the model produces speech that closely resembles them. This
//! module wraps that pattern in an ergonomic API with budget management
//! and optional input loudness-matching.
//!
//! Token budget math (CSM-1B):
//!   - 2048 backbone max sequence
//!   - Each audio frame (80 ms) costs 1 sequence position (33 slots wide)
//!   - Each text token costs 1 position
//!   - Reserve ~500 positions for the new utterance + safety margin
//!   - → ~1500 positions of *context* available
//!   - At 12.5 fps that's ~120 s of audio history before tokens are needed
//!     for the utterance text
//!
//! The 2048 ceiling is the hard wall. `fit_within_budget` evicts oldest
//! segments until total ≤ budget.

use crate::audio_io::TARGET_SAMPLE_RATE;
use crate::error::{CsmError, Result};
use crate::prompt::Segment;

/// Reasonable budget leaving headroom for the new utterance and EOT.
pub const DEFAULT_PROFILE_BUDGET_TOKENS: usize = 1500;

#[derive(Debug, Clone)]
pub struct SpeakerProfile {
    pub id: u32,
    pub segments: Vec<Segment>,
}

impl SpeakerProfile {
    pub fn new(id: u32) -> Self {
        Self {
            id,
            segments: Vec::new(),
        }
    }

    /// Add a reference utterance: a transcript and the audio (24 kHz mono f32).
    /// Audio shorter than 200 ms is rejected — too short to convey speaker
    /// identity and risks degrading the prompt.
    pub fn add_reference(mut self, text: impl Into<String>, mut audio: Vec<f32>) -> Result<Self> {
        let min_samples = (TARGET_SAMPLE_RATE as f32 * 0.2) as usize;
        if audio.len() < min_samples {
            return Err(CsmError::Config(format!(
                "reference audio is {} samples (~{:.0}ms); need at least 200ms",
                audio.len(),
                audio.len() as f32 / TARGET_SAMPLE_RATE as f32 * 1000.0
            )));
        }
        // Loudness-match: rough peak normalize to -6 dBFS so all references
        // sit at comparable level. Avoid LUFS here since it's slow per-call;
        // peak norm is good enough as an input-side equalizer.
        peak_normalize(&mut audio, 0.5012);
        self.segments.push(Segment::new(self.id, text, audio));
        Ok(self)
    }

    /// Returns an estimate of the token cost of this profile when used as
    /// generation context. Estimate is conservative (audio frames, text token
    /// per 4 ASCII characters as a rule-of-thumb).
    pub fn estimated_tokens(&self) -> usize {
        let mut total = 0usize;
        for seg in &self.segments {
            // Audio frames at 12.5 fps. 1920 samples per frame at 24 kHz.
            if let Some(audio) = &seg.audio {
                total += audio.len().div_ceil(1920);
            }
            // Text: ~4 chars per BPE token, rounded up. Conservative estimate.
            total += seg.text.chars().count().div_ceil(4);
        }
        total
    }

    /// Drop oldest segments until the estimated token cost is ≤ `budget`.
    /// Always retains at least one segment if any are present (a partial
    /// reference is better than no reference).
    pub fn fit_within_budget(&mut self, budget: usize) {
        while self.segments.len() > 1 && self.estimated_tokens() > budget {
            self.segments.remove(0);
        }
    }

    pub fn is_empty(&self) -> bool {
        self.segments.is_empty()
    }

    pub fn segments(&self) -> &[Segment] {
        &self.segments
    }
}

fn peak_normalize(samples: &mut [f32], target_peak: f32) {
    if samples.is_empty() {
        return;
    }
    let peak = samples.iter().fold(0.0f32, |a, &b| a.max(b.abs()));
    if peak < 1e-6 {
        return;
    }
    let gain = target_peak / peak;
    for s in samples.iter_mut() {
        *s *= gain;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn dummy_audio(seconds: f32) -> Vec<f32> {
        let n = (seconds * TARGET_SAMPLE_RATE as f32) as usize;
        (0..n).map(|i| 0.1 * (i as f32 * 0.001).sin()).collect()
    }

    #[test]
    fn rejects_too_short_audio() {
        let r = SpeakerProfile::new(0).add_reference("hi", vec![0.1; 100]);
        assert!(r.is_err());
    }

    #[test]
    fn accepts_short_reference() {
        let p = SpeakerProfile::new(0)
            .add_reference("hello", dummy_audio(1.0))
            .unwrap();
        assert_eq!(p.segments.len(), 1);
        assert_eq!(p.id, 0);
    }

    #[test]
    fn estimated_tokens_rough_math() {
        let p = SpeakerProfile::new(0)
            .add_reference("hello world", dummy_audio(2.0))
            .unwrap();
        // 2 seconds = 25 frames, "hello world" = 11 chars ~= 3 tokens
        let est = p.estimated_tokens();
        assert!(
            est >= 25 && est <= 40,
            "estimated {est} outside reasonable bounds"
        );
    }

    #[test]
    fn fit_within_budget_evicts_oldest() {
        let p = SpeakerProfile::new(0)
            .add_reference("first", dummy_audio(5.0))
            .unwrap()
            .add_reference("second", dummy_audio(5.0))
            .unwrap()
            .add_reference("third", dummy_audio(5.0))
            .unwrap();
        // 3 × ~63 frames = ~189 tokens. Squeeze into 100.
        let mut p = p;
        p.fit_within_budget(100);
        assert!(p.estimated_tokens() <= 100);
        // Last segment ("third") preserved
        assert_eq!(p.segments.last().unwrap().text, "third");
    }

    #[test]
    fn fit_within_budget_keeps_at_least_one() {
        let mut p = SpeakerProfile::new(0)
            .add_reference("only", dummy_audio(3.0))
            .unwrap();
        p.fit_within_budget(1);
        // Budget is impossibly small but we keep the single segment.
        assert_eq!(p.segments.len(), 1);
    }

    #[test]
    fn peak_normalize_scales_to_target() {
        let mut x = vec![0.0, 0.1, -0.2, 0.05];
        peak_normalize(&mut x, 0.5);
        let new_peak = x.iter().fold(0.0f32, |a, &b| a.max(b.abs()));
        assert!((new_peak - 0.5).abs() < 1e-5);
    }
}