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

//! Speech synthesis on top of qwentts.cpp.
//!
//! It talks to `tts-server` over HTTP instead of invoking the `qwen-tts`
//! binary once per sentence. The difference is not a matter of style: the
//! process loads 2.1 GB of talker model plus 291 MB of codec, and paying that
//! load on every sentence would put several seconds in front of every answer.
//! The server keeps them resident on the GPU and answers in `pcm` format,
//! which arrives in chunks as it is generated.

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};

/// Rate the model synthesizes at.
pub const SAMPLE_RATE: u32 = 24_000;

/// How a synthesis ended.
#[derive(Debug, Clone, Default)]
pub struct SpeechOutcome {
    /// Time until the first audio block.
    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
    }

    /// Multiple of real time. Above 1.0 the synthesizer is slower than
    /// speech and the playback queue will eventually run dry.
    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!(
            "{} did not answer /health within {} s",
            self.http.authority,
            timeout.as_secs()
        )))
    }

    /// Registers a cloned voice from the `qwen-codec` latents.
    ///
    /// The already extracted `.spk` and `.rvq` are sent instead of the original
    /// WAV: the server takes them as they are and does not have to analyze the
    /// reference again on every start.
    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!("could not read {what} ({}): {e}", path.display())))
        };
        let b64 = base64::engine::general_purpose::STANDARD;
        let speaker = b64.encode(read(&voice.speaker, "the speaker embedding")?);
        let codes = b64.encode(read(&voice.codes, "the reference codes")?);
        let transcript = String::from_utf8_lossy(&read(&voice.transcript, "the transcript")?)
            .trim()
            .to_string();

        if transcript.is_empty() {
            return Err(Error::Tts(format!(
                "the reference transcript {} is empty; without it cloning is not enabled",
                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", speech = %voice.name, "voz clonada registrada");
        Ok(())
    }

    /// Voices the server knows right now.
    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())
    }

    /// Synthesizes `text` and delivers the audio in pieces as it is generated.
    ///
    /// `on_audio` gets mono f32 samples at 24 kHz and returns `false` to stop;
    /// `cancel` does the same from another thread. That cut is what allows going
    /// silent at once when the user interrupts, without waiting for a sentence
    /// that will no longer be heard to finish generating.
    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();
        // The server does not guarantee each block carries an even number of
        // bytes, so a stray byte from one block waits for the next instead of
        // shifting every later sample.
        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",
            chars = text.chars().count(),
            ttfb_ms = outcome.ttfb.map(|d| d.as_millis()).unwrap_or(0),
            audio_s = outcome.audio_secs(),
            rtf = outcome.rtf(),
            cancelled = outcome.cancelled,
            "synthesis"
        );
        Ok(outcome)
    }

    /// Synthesizes a short sentence to pay upfront for building the graphs:
    /// the first request costs about 3.5 s more than the following ones, and
    /// they are better not spent on the first real turn.
    pub fn warmup(&self) -> Result<()> {
        if !self.config.warmup {
            return Ok(());
        }
        let started = Instant::now();
        let cancel = Cancel::new();
        // The audio is thrown away: only the side effect matters.
        self.speak("Listo.", &cancel, |_| true)?;
        tracing::info!(target: "tts", ms = started.elapsed().as_millis(), "precalentado");
        Ok(())
    }
}

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

    #[test]
    fn rtf_relates_compute_time_and_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 without_audio_rtf_is_infinite_instead_of_dividing_by_zero() {
        assert!(SpeechOutcome::default().rtf().is_infinite());
    }
}