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

//! Mimi neural audio codec wrapper.
//!
//! Uses `candle_transformers::models::mimi` (the HF-compatible implementation)
//! rather than the `moshi` crate — the two use different weight-key naming
//! conventions and the HF `kyutai/mimi/model.safetensors` is laid out for the
//! former.
//!
//! 24 kHz, 12.5 Hz frames (80 ms), 32 codebooks, per-codebook vocab 2048.
//! CSM uses vocab 2051 (3 reserved special tokens).

use crate::error::{CsmError, Result};
use candle_core::{Device, StreamTensor, Tensor};
use candle_transformers::models::mimi as cmimi;
use std::path::Path;

pub const SAMPLE_RATE: u32 = 24_000;
pub const NUM_CODEBOOKS: usize = 32;

pub struct Mimi {
    inner: cmimi::Model,
    device: Device,
    /// Path the weights were loaded from. Cached so [`Self::reload`] can
    /// rebuild a fresh `inner` when `reset_state()` proves insufficient
    /// (Mimi's transformer carries a position counter that the upstream
    /// reset doesn't fully clear; surfaces as a `narrow [8192, 32]`
    /// panic after ~80 encoded clips even with reset_state between them).
    weights_path: std::path::PathBuf,
}

impl Mimi {
    /// Load the HF-hosted Mimi weights. `path` must point at the safetensors
    /// downloaded from `kyutai/mimi`.
    pub fn load<P: AsRef<Path>>(path: P, device: &Device) -> Result<Self> {
        let path_str = path
            .as_ref()
            .to_str()
            .ok_or_else(|| CsmError::Config("non-utf8 path".into()))?;
        let inner = cmimi::load(path_str, Some(NUM_CODEBOOKS), device)?;
        Ok(Self {
            inner,
            device: device.clone(),
            weights_path: path.as_ref().to_path_buf(),
        })
    }

    /// Drop and rebuild the inner Mimi model from the cached weights
    /// path. Use when `reset_state()` is known to be insufficient
    /// (e.g. data-prep loops that encode hundreds of clips back-to-back).
    /// Cost: ~200 ms on Metal — call sparingly, e.g. every 30-50 clips.
    pub fn reload(&mut self) -> Result<()> {
        let path_str = self
            .weights_path
            .to_str()
            .ok_or_else(|| CsmError::Config("non-utf8 path".into()))?;
        self.inner = cmimi::load(path_str, Some(NUM_CODEBOOKS), &self.device)?;
        Ok(())
    }

    /// Encode 24 kHz mono samples to discrete codes.
    /// Returns shape `(1, num_codebooks=32, num_frames)` i64.
    pub fn encode(&mut self, samples: &[f32]) -> Result<Tensor> {
        let pcm = Tensor::from_slice(samples, (1, 1, samples.len()), &self.device)?;
        let codes = self.inner.encode(&pcm)?;
        Ok(codes)
    }

    /// Decode discrete codes back to 24 kHz mono samples.
    /// `codes` must have shape `(1, num_codebooks=32, num_frames)`.
    pub fn decode(&mut self, codes: &Tensor) -> Result<Vec<f32>> {
        let pcm = self.inner.decode(codes)?;
        let samples = pcm.flatten_all()?.to_vec1::<f32>()?;
        Ok(samples)
    }

    pub fn device(&self) -> &Device {
        &self.device
    }

    /// Reset all streaming state. Call before starting a new stream — leftover
    /// state from the prior stream will produce wrong audio at the boundary.
    pub fn reset_state(&mut self) {
        self.inner.reset_state();
    }

    /// Streaming decode: feed a new chunk of codes (or `None` to flush internal
    /// state) and receive the output samples produced for that chunk. Mimi
    /// maintains its decoder/upsampler state across calls, so the cost is O(n)
    /// total rather than the O(n²) of repeatedly calling `decode` on a growing
    /// prefix.
    ///
    /// `codes` shape: `(1, num_codebooks=32, k)` for a chunk of `k` frames.
    /// Returns `Some(samples)` when the streaming chain has emitted output for
    /// this step, `None` when it's still buffering.
    pub fn decode_step(&mut self, codes: Option<&Tensor>) -> Result<Option<Vec<f32>>> {
        let st = match codes {
            Some(t) => StreamTensor::from_tensor(t.clone()),
            None => StreamTensor::empty(),
        };
        let out = self.inner.decode_step(&st)?;
        match out.as_option() {
            None => Ok(None),
            Some(t) => {
                let samples = t.flatten_all()?.to_vec1::<f32>()?;
                Ok(Some(samples))
            }
        }
    }
}