car-voice 0.51.0

//! In-process Kokoro-82M TTS provider.
//!
//! Wraps `car_inference::backend::mlx_kokoro::KokoroBackend` so voice
//! synthesis runs entirely on-device via MLX/Metal — no HTTP server, no
//! subprocess, no Python. Replaces the `LocalTtsSpeaker` path for users
//! who don't want to keep an `mlx-audio` (or Piper-shim) Python server
//! running alongside CAR.
//!
//! Only compiled on Apple Silicon macOS because that's the only platform
//! where the MLX backend builds. Other targets should use
//! `TtsProvider::Elevenlabs` or `TtsProvider::Local`.

use crate::tts::{AudioFormat, Speaker, SynthesizedAudio};
use crate::{Result, VoiceConfig, VoiceError};
use async_trait::async_trait;
use car_inference::backend::mlx_kokoro::KokoroBackend;
use car_inference::backend_cache::{estimate_model_size, BackendRetention, CachedBackend};
use std::path::{Path, PathBuf};
use std::sync::Arc;

/// In-process Kokoro TTS.
///
/// Holds an `Arc<Mutex<KokoroBackend>>` handle out of the shared
/// process-wide cache. Concurrent `synth` calls on the same `voice`
/// serialize on the mutex (MLX ops are not `Sync`); two different
/// voices would hit the same backend anyway since Kokoro voices are
/// style vectors, not separate models — only one backend gets loaded.
pub struct KokoroSpeaker {
    backend: CachedBackend<KokoroBackend>,
    transient_reservation: Option<Arc<car_inference::resource_policy::LocalLoadReservation>>,
    _cache_scope: car_inference::ScopedKokoroBackendCache,
    voice: String,
    tmp_dir: PathBuf,
}

impl std::fmt::Debug for KokoroSpeaker {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("KokoroSpeaker")
            .field("voice", &self.voice)
            .field("tmp_dir", &self.tmp_dir)
            .finish()
    }
}

impl KokoroSpeaker {
    /// Construct from a [`VoiceConfig`]. Loads the Kokoro backend via
    /// the shared cache — first construction pays the ~1 s load cost,
    /// subsequent constructions / calls are cheap.
    pub fn from_config(config: &VoiceConfig) -> Result<Self> {
        let state_root = car_home::root_or_relative();
        Self::from_config_in_scope(config, &state_root)
    }

    /// Construct within an explicit CAR state root. Admission, cache identity,
    /// policy, and later maintenance all use this same physical scope.
    pub fn from_config_in_scope(config: &VoiceConfig, state_root: &Path) -> Result<Self> {
        let model_dir = resolve_kokoro_model_dir(&config.local_tts_model).ok_or_else(|| {
            VoiceError::Config(format!(
                "Kokoro TTS: model '{}' not found in HuggingFace cache — \
                 run `car models pull {}` first",
                config.local_tts_model, config.local_tts_model,
            ))
        })?;
        let key = config.local_tts_model.clone();
        let size = estimate_model_size(&model_dir);
        let cache_scope = car_inference::scoped_kokoro_backend_cache(state_root);
        let admission = cache_scope.admission().clone();
        let mut aliases = crate::local_admission::catalog_aliases(&key, Some(&model_dir));
        aliases.push(format!("voice/kokoro/{key}"));
        admission.register_model_aliases(&key, aliases);
        let cache_key = admission.canonical_model_id(&key);
        let mut reservation = admission
            .reserve_measured_host(&cache_key, size, 128)
            .map_err(|error| VoiceError::Config(error.to_string()))?;
        // LOCAL_ADMISSION_BOUNDARY:voice-kokoro
        let (backend, retention) = cache_scope
            .cache()
            .get_or_load_admitted(&cache_key, size, &mut reservation, || {
                KokoroBackend::load(&model_dir)
            })
            .map_err(|e: car_inference::InferenceError| {
                VoiceError::Config(format!("load Kokoro model: {e}"))
            })?;
        Ok(Self {
            backend,
            transient_reservation: (retention == BackendRetention::Transient)
                .then(|| Arc::new(reservation)),
            _cache_scope: cache_scope,
            voice: config.local_tts_voice.clone(),
            tmp_dir: std::env::temp_dir(),
        })
    }
}

#[async_trait]
impl Speaker for KokoroSpeaker {
    async fn synth(&self, text: &str) -> Result<SynthesizedAudio> {
        // Write to a per-call temp WAV, read it back into memory, and
        // return the bytes. The backend API is path-based; the disk
        // round-trip is negligible vs synthesis time.
        let wav_path = self
            .tmp_dir
            .join(format!("car-kokoro-{}.wav", uuid::Uuid::new_v4()));
        let voice = self.voice.clone();
        let text = text.to_string();
        let backend = Arc::clone(&self.backend);
        // A cancelled caller may drop the speaker while spawn_blocking is
        // still executing. Clone the transient allocation lease into the task
        // so memory remains charged until native synthesis actually exits.
        let transient_reservation = self.transient_reservation.clone();
        let wav_path_clone = wav_path.clone();

        // `synthesize` is synchronous and CPU/GPU-bound. Hold the
        // sync Mutex only on the blocking worker thread to keep the
        // tokio runtime moving.
        let wav_path = tokio::task::spawn_blocking(move || -> Result<PathBuf> {
            let _transient_reservation = transient_reservation;
            let mut guard = backend
                .lock()
                .map_err(|_| VoiceError::Config("Kokoro backend mutex poisoned".into()))?;
            guard
                .synthesize(&text, Some(&voice), &wav_path_clone)
                .map_err(|e| VoiceError::Config(format!("Kokoro synth: {e}")))
        })
        .await
        .map_err(|e| VoiceError::Config(format!("Kokoro synth task: {e}")))??;

        let bytes = tokio::fs::read(&wav_path)
            .await
            .map_err(|e| VoiceError::Config(format!("read Kokoro wav: {e}")))?;
        let _ = tokio::fs::remove_file(&wav_path).await;

        Ok(SynthesizedAudio {
            bytes,
            format: AudioFormat::Wav,
        })
    }
}

/// Locate the Kokoro model dir in the HuggingFace cache.
///
/// Matches `~/.cache/huggingface/hub/models--{org}--{repo}/snapshots/*/`
/// — the same layout `hf_hub::api` writes and that our
/// `car_inference` loaders read.
fn resolve_kokoro_model_dir(repo: &str) -> Option<PathBuf> {
    let home = dirs::home_dir()?;
    let hub_dir = home.join(".cache/huggingface/hub");
    let normalized = repo.replace('/', "--");
    let candidate = hub_dir
        .join(format!("models--{normalized}"))
        .join("snapshots");
    let entries = std::fs::read_dir(&candidate).ok()?;
    // Pick the newest snapshot dir (there's typically one; take the most
    // recent mtime to be resilient to partial pulls).
    entries
        .flatten()
        .filter_map(|e| {
            let p = e.path();
            let meta = e.metadata().ok()?;
            if !meta.is_dir() {
                return None;
            }
            let mtime = meta.modified().ok()?;
            Some((mtime, p))
        })
        .max_by_key(|(m, _)| *m)
        .map(|(_, p)| p)
}