piperust/src/tts/fish_audio.rs

//! Fish Audio / SGLang Drama 3 TTS plugin.
//!
//! POSTs text to `/v1/audio/speech` (OpenAI-compatible) requesting PCM
//! streaming. Fish Audio Drama 3 outputs at 44.1 kHz, which matches the
//! mixer's output rate, so samples are pushed directly into the TTS queue
//! without resampling. A [`CancellationToken`] allows immediate interruption
//! (barge-in).
//!
//! Drama 3 can sound overly theatrical on some references. Two mitigations are
//! applied here, both configurable from the environment: optional sampling
//! knobs (`temperature` / `top_p` / `top_k`) that trim the tails of the
//! acoustic token distribution, and a neutral style tag prepended to untagged
//! requests so plain LLM prose is read in a conversational register.

use crate::audio::SampleRingBuffer;
use crate::config::Config;
use crate::tts::TtsEngine;
use anyhow::{Context, Result};
use async_trait::async_trait;
use futures_util::StreamExt;
use reqwest::Client;
use serde::Serialize;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info};

/// A voice-cloning reference: path to a WAV file + its transcript.
#[derive(Debug, Clone, Serialize)]
struct Reference {
    audio_path: String,
    text: String,
}

/// Request body for the SGLang `/v1/audio/speech` endpoint.
///
/// Sampling knobs are optional and only serialized when set. Drama 3 accepts
/// `temperature`, `top_p` and `top_k` at this level; `repetition_penalty` is
/// deliberately never sent because this server rejects any value other than
/// `1.0` (`Drama 3 request: repetition_penalty must be 1.0`).
#[derive(Debug, Serialize)]
struct SpeechRequest<'a> {
    model: &'a str,
    input: &'a str,
    voice: &'a str,
    response_format: &'a str,
    stream: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    references: Option<Vec<Reference>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_p: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_k: Option<u32>,
}

/// Streaming TTS client backed by Fish Audio / SGLang-Omni Drama 3.
#[derive(Clone)]
pub struct FishAudioTts {
    http: Client,
    endpoint: String,
    model: String,
    voice: String,
    /// Current voice-cloning reference (config default, overridable at runtime).
    reference: Arc<parking_lot::RwLock<Option<Reference>>>,
    /// Optional sampling overrides (de-emphasis of the recitative register).
    temperature: Option<f32>,
    top_p: Option<f32>,
    top_k: Option<u32>,
    /// Neutral style tag prepended to untagged synthesis requests.
    style_tag: String,
}

impl FishAudioTts {
    pub fn new(config: &Config) -> Self {
        let reference = config
            .tts_reference()
            .map(|(audio_path, text)| Reference { audio_path, text });
        if let Some(r) = &reference {
            info!(
                audio_path = %r.audio_path,
                text_len = r.text.len(),
                "TTS voice-cloning reference configured"
            );
        }
        info!(
            temperature = ?config.tts_temperature,
            top_p = ?config.tts_top_p,
            top_k = ?config.tts_top_k,
            style_tag = %config.tts_style_tag,
            "TTS de-emphasis settings"
        );
        Self {
            http: Client::builder()
                .build()
                .expect("failed to build reqwest client"),
            endpoint: config.speech_endpoint(),
            model: config.tts_model.clone(),
            voice: config.tts_voice.clone(),
            reference: Arc::new(parking_lot::RwLock::new(reference)),
            temperature: config.tts_temperature,
            top_p: config.tts_top_p,
            top_k: config.tts_top_k,
            style_tag: config.tts_style_tag.clone(),
        }
    }

    /// The text actually sent to the model: the configured neutral style tag is
    /// prepended unless the text already opens with a style tag or the tag is
    /// disabled. LLM replies carry no tags, so they get de-emphasized; clip
    /// banks that start with their own tag are left untouched.
    fn styled<'a>(&self, text: &'a str) -> std::borrow::Cow<'a, str> {
        let tag = self.style_tag.trim();
        if tag.is_empty() || text.trim_start().starts_with('[') {
            std::borrow::Cow::Borrowed(text)
        } else {
            std::borrow::Cow::Owned(format!("{tag} {text}"))
        }
    }

    /// Override the voice-cloning reference at runtime. Pass `None` to clear.
    pub fn set_reference(&self, audio_path: Option<String>, text: Option<String>) {
        let mut guard = self.reference.write();
        match (audio_path, text) {
            (Some(audio_path), text) => {
                let text = text.unwrap_or_default();
                info!(audio_path = %audio_path, text_len = text.len(), "TTS reference updated");
                *guard = Some(Reference { audio_path, text });
            }
            _ => {
                info!("TTS reference cleared");
                *guard = None;
            }
        }
    }
}

#[async_trait]
impl TtsEngine for FishAudioTts {
    fn name(&self) -> &str {
        "fish-audio"
    }

    fn sample_rate(&self) -> u32 {
        // Fish Audio Drama 3 outputs natively at 44.1 kHz, matching the mixer.
        44_100
    }

    async fn synthesize(
        &self,
        text: &str,
        tts_queue: &SampleRingBuffer,
        cancel: &CancellationToken,
    ) -> Result<()> {
        let references = self.reference.read().clone().map(|r| vec![r]);
        let input = self.styled(text);

        let body = SpeechRequest {
            model: &self.model,
            input: &input,
            voice: &self.voice,
            response_format: "pcm",
            stream: true,
            references,
            temperature: self.temperature,
            top_p: self.top_p,
            top_k: self.top_k,
        };

        info!(model = %self.model, text_len = text.len(), "SGLang TTS request");
        debug!(
            endpoint = %self.endpoint,
            model = %self.model,
            voice = %self.voice,
            response_format = "pcm",
            stream = true,
            has_reference = body.references.is_some(),
            temperature = ?self.temperature,
            top_p = ?self.top_p,
            top_k = ?self.top_k,
            styled = input.len() != text.len(),
            text = %input,
            "TTS request dispatched"
        );

        let req = self
            .http
            .post(&self.endpoint)
            .json(&body)
            .build()
            .context("failed to build SGLang request")?;

        let response = self
            .http
            .execute(req)
            .await
            .context("SGLang request failed")?;

        if !response.status().is_success() {
            let status = response.status();
            let err_text = response.text().await.unwrap_or_default();
            error!(%status, body = %err_text, "SGLang returned an error");
            anyhow::bail!("SGLang error: {} {}", status, err_text);
        }

        // Read the actual sample rate from the response header for logging.
        let tts_sr = response
            .headers()
            .get("x-sample-rate")
            .and_then(|v| v.to_str().ok())
            .and_then(|s| s.parse::<u32>().ok())
            .unwrap_or(44100);
        info!(tts_sample_rate = tts_sr, "TTS response sample rate");

        let mut byte_stream = response.bytes_stream();
        let mut carry: Option<u8> = None;
        let mut chunk_count = 0usize;
        let mut total_samples = 0usize;
        let t0 = std::time::Instant::now();
        let mut last_log = t0;

        while let Some(chunk) = tokio::select! {
            c = byte_stream.next() => c,
            _ = cancel.cancelled() => {
                info!("TTS stream cancelled (barge-in)");
                return Ok(());
            }
        } {
            let chunk = match chunk {
                Ok(b) => b,
                Err(e) => {
                    error!(error = %e, "error reading TTS stream chunk");
                    return Err(e.into());
                }
            };
            let n = ingest_pcm_chunk(&chunk, &mut carry, tts_queue);
            chunk_count += 1;
            total_samples += n;
            if chunk_count == 1 {
                debug!(
                    ttfb_ms = t0.elapsed().as_millis(),
                    bytes = chunk.len(),
                    "TTS time-to-first-byte"
                );
            }
            debug!(
                chunk = chunk_count,
                bytes = chunk.len(),
                samples = n,
                queue_len = tts_queue.len(),
                elapsed_ms = t0.elapsed().as_millis(),
                "TTS chunk received"
            );
            if last_log.elapsed() > std::time::Duration::from_millis(500) {
                last_log = std::time::Instant::now();
                info!(
                    chunk = chunk_count,
                    total_samples,
                    queue_len = tts_queue.len(),
                    elapsed_ms = t0.elapsed().as_millis(),
                    "TTS stream progress"
                );
            }
        }

        info!(
            chunk_count,
            total_samples,
            duration_s = format!("{:.2}", t0.elapsed().as_secs_f32()),
            "TTS stream complete"
        );
        debug!(
            chunk_count,
            total_samples,
            avg_chunk_bytes = if chunk_count > 0 {
                (total_samples * 2) / chunk_count
            } else {
                0
            },
            audio_s = format!("{:.2}", total_samples as f32 / tts_sr as f32),
            elapsed_ms = t0.elapsed().as_millis(),
            queue_len = tts_queue.len(),
            "TTS stream summary"
        );
        Ok(())
    }

    fn set_reference(&self, audio_path: Option<String>, text: Option<String>) {
        FishAudioTts::set_reference(self, audio_path, text);
    }
}

/// Parse 16-bit LE PCM samples from a chunk and push them directly to the
/// queue (Fish Audio's 44.1 kHz output matches the mixer rate).
/// Returns the number of samples pushed.
fn ingest_pcm_chunk(chunk: &[u8], carry: &mut Option<u8>, tts_queue: &SampleRingBuffer) -> usize {
    let mut samples: Vec<i16> = Vec::with_capacity(chunk.len() / 2 + 1);
    let mut iter = chunk.iter().copied();

    if let Some(lo) = carry.take() {
        if let Some(hi) = iter.next() {
            samples.push(i16::from_le_bytes([lo, hi]));
        } else {
            *carry = Some(lo);
        }
    }

    loop {
        let lo = match iter.next() {
            Some(b) => b,
            None => break,
        };
        let hi = match iter.next() {
            Some(b) => b,
            None => {
                *carry = Some(lo);
                break;
            }
        };
        samples.push(i16::from_le_bytes([lo, hi]));
    }

    if !samples.is_empty() {
        tts_queue.push(&samples);
    }
    samples.len()
}