path: root/crates/asist-tts/src/lib.rs

//! Síntesis de voz sobre qwentts.cpp.
//!
//! Se habla con `tts-server` por HTTP en lugar de invocar el binario
//! `qwen-tts` una vez por frase. La diferencia no es de estilo: el proceso
//! carga 2,1 GB de modelo hablante más 291 MB de codec, y pagar esa carga en
//! cada frase pondría varios segundos delante de cada respuesta. El servidor
//! los mantiene residentes en la GPU y responde en formato `pcm`, que llega
//! troceado según se genera.

use std::time::{Duration, Instant};

use base64::Engine;
use serde_json::json;

use asist_core::config::{ReferenceVoice, TtsConfig};
use asist_core::error::{Error, Result};
use asist_core::http::{Cancel, HttpClient};

/// Frecuencia a la que sintetiza el modelo.
pub const SAMPLE_RATE: u32 = 24_000;

/// Cómo terminó una síntesis.
#[derive(Debug, Clone, Default)]
pub struct SpeechOutcome {
    /// Tiempo hasta el primer bloque de audio.
    pub ttfb: Option<Duration>,
    pub total: Duration,
    pub samples: usize,
    pub cancelled: bool,
}

impl SpeechOutcome {
    pub fn audio_secs(&self) -> f32 {
        self.samples as f32 / SAMPLE_RATE as f32
    }

    /// Múltiplo del tiempo real. Por encima de 1,0 el sintetizador va más
    /// lento que el habla y la cola de reproducción acabará vaciándose.
    pub fn rtf(&self) -> f32 {
        let audio = self.audio_secs();
        if audio <= 0.0 {
            return f32::INFINITY;
        }
        self.total.as_secs_f32() / audio
    }
}

pub struct TtsClient {
    http: HttpClient,
    config: TtsConfig,
}

impl TtsClient {
    pub fn new(authority: String, config: &TtsConfig) -> Self {
        Self {
            http: HttpClient::new(authority)
                .with_read_timeout(Duration::from_secs(config.request_timeout_secs)),
            config: config.clone(),
        }
    }

    pub fn healthy(&self) -> bool {
        self.http.healthy("/health")
    }

    pub fn wait_ready(&self, timeout: Duration) -> Result<()> {
        let deadline = Instant::now() + timeout;
        while Instant::now() < deadline {
            if self.healthy() {
                return Ok(());
            }
            std::thread::sleep(Duration::from_millis(500));
        }
        Err(Error::Tts(format!(
            "{} no respondió a /health en {} s",
            self.http.authority,
            timeout.as_secs()
        )))
    }

    /// Registra una voz clonada a partir de los latentes de `qwen-codec`.
    ///
    /// Se mandan los `.spk` y `.rvq` ya extraídos en lugar del WAV original:
    /// el servidor los toma tal cual y se ahorra volver a analizar la
    /// referencia en cada arranque.
    pub fn register_voice(&self, voice: &ReferenceVoice) -> Result<()> {
        let read = |path: &std::path::Path, what: &str| -> Result<Vec<u8>> {
            std::fs::read(path).map_err(|e| {
                Error::Tts(format!("no se pudo leer {what} ({}): {e}", path.display()))
            })
        };
        let b64 = base64::engine::general_purpose::STANDARD;
        let speaker = b64.encode(read(&voice.speaker, "el embedding del hablante")?);
        let codes = b64.encode(read(&voice.codes, "los códigos de referencia")?);
        let transcript = String::from_utf8_lossy(&read(&voice.transcript, "la transcripción")?)
            .trim()
            .to_string();

        if transcript.is_empty() {
            return Err(Error::Tts(format!(
                "la transcripción de referencia {} está vacía; sin ella no se activa el clonado",
                voice.transcript.display()
            )));
        }

        let body = json!({
            "name": voice.name,
            "ref_text": transcript,
            "spk_b64": speaker,
            "rvq_b64": codes,
        });
        self.http.post_json("/v1/audio/voices", &body)?;
        tracing::info!(target: "tts", voz = %voice.name, "voz clonada registrada");
        Ok(())
    }

    /// Voces que el servidor conoce ahora mismo.
    pub fn voices(&self) -> Result<Vec<String>> {
        let body = self.http.get("/v1/audio/voices")?;
        let parsed: serde_json::Value = serde_json::from_slice(&body)?;
        Ok(parsed
            .get("voices")
            .and_then(|v| v.as_array())
            .map(|voices| {
                voices
                    .iter()
                    .filter_map(|v| {
                        v.get("name")
                            .and_then(|n| n.as_str())
                            .or_else(|| v.as_str())
                            .map(String::from)
                    })
                    .collect()
            })
            .unwrap_or_default())
    }

    /// Sintetiza `text` y entrega el audio a trozos según se genera.
    ///
    /// `on_audio` recibe muestras mono f32 a 24 kHz y devuelve `false` para
    /// cortar; `cancel` hace lo propio desde otro hilo. Ese corte es lo que
    /// permite callar de golpe cuando el usuario interrumpe, sin esperar a que
    /// termine de generarse una frase que ya no se va a oír.
    pub fn speak(
        &self,
        text: &str,
        cancel: &Cancel,
        mut on_audio: impl FnMut(&[f32]) -> bool,
    ) -> Result<SpeechOutcome> {
        let body = json!({
            "input": text,
            "voice": self.config.voice,
            "response_format": "pcm",
            "temperature": self.config.temperature,
            "top_k": self.config.top_k,
            "top_p": self.config.top_p,
            "repetition_penalty": self.config.repetition_penalty,
            "max_new_tokens": self.config.max_new_tokens,
        });

        let started = Instant::now();
        let mut outcome = SpeechOutcome::default();
        // El servidor no garantiza que cada bloque traiga un número par de
        // bytes, así que el byte suelto de un bloque espera al siguiente en
        // vez de desplazar todas las muestras posteriores.
        let mut odd: Option<u8> = None;
        let mut samples = Vec::with_capacity(8192);
        let mut stopped = false;

        let result = self
            .http
            .post_json_streaming("/v1/audio/speech", &body, cancel, |chunk| {
                if outcome.ttfb.is_none() {
                    outcome.ttfb = Some(started.elapsed());
                }
                samples.clear();
                let mut bytes = chunk;
                if let Some(pending) = odd.take() {
                    if let Some((first, rest)) = bytes.split_first() {
                        samples.push(i16::from_le_bytes([pending, *first]) as f32 / 32768.0);
                        bytes = rest;
                    } else {
                        odd = Some(pending);
                    }
                }
                if bytes.len() % 2 == 1 {
                    odd = Some(bytes[bytes.len() - 1]);
                    bytes = &bytes[..bytes.len() - 1];
                }
                for pair in bytes.chunks_exact(2) {
                    samples.push(i16::from_le_bytes([pair[0], pair[1]]) as f32 / 32768.0);
                }
                outcome.samples += samples.len();
                if !on_audio(&samples) {
                    stopped = true;
                    return false;
                }
                true
            });

        outcome.total = started.elapsed();
        match result {
            Ok(()) => {}
            Err(Error::Cancelled) => outcome.cancelled = true,
            Err(err) => return Err(err),
        }
        outcome.cancelled |= stopped;

        tracing::debug!(
            target: "tts",
            caracteres = text.chars().count(),
            ttfb_ms = outcome.ttfb.map(|d| d.as_millis()).unwrap_or(0),
            audio_s = outcome.audio_secs(),
            rtf = outcome.rtf(),
            cancelada = outcome.cancelled,
            "síntesis"
        );
        Ok(outcome)
    }

    /// Sintetiza una frase corta para pagar por adelantado la construcción de
    /// los grafos: la primera petición cuesta unos 3,5 s más que las
    /// siguientes, y no conviene gastarlos en el primer turno de verdad.
    pub fn warmup(&self) -> Result<()> {
        if !self.config.warmup {
            return Ok(());
        }
        let started = Instant::now();
        let cancel = Cancel::new();
        // El audio se tira: sólo interesa el efecto secundario.
        self.speak("Listo.", &cancel, |_| true)?;
        tracing::info!(target: "tts", ms = started.elapsed().as_millis(), "precalentado");
        Ok(())
    }
}

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

    #[test]
    fn el_rtf_relaciona_tiempo_de_calculo_y_audio() {
        let outcome = SpeechOutcome {
            total: Duration::from_secs(2),
            samples: SAMPLE_RATE as usize * 4,
            ..Default::default()
        };
        assert_eq!(outcome.audio_secs(), 4.0);
        assert_eq!(outcome.rtf(), 0.5);
    }

    #[test]
    fn sin_audio_el_rtf_es_infinito_en_vez_de_dividir_por_cero() {
        assert!(SpeechOutcome::default().rtf().is_infinite());
    }
}