From 8e868894cbcaf421c263f3555af4d29bef653420 Mon Sep 17 00:00:00 2001 From: salmonik Date: Sat, 5 Sep 2026 19:35:21 +0200 Subject: feat: piper now not a subprocess, cache system for TTS --- src/tts.rs | 384 ++++++++++++++++++++++++++++++++++++++++++------------------- 1 file changed, 264 insertions(+), 120 deletions(-) (limited to 'src/tts.rs') diff --git a/src/tts.rs b/src/tts.rs index ded76ad..3ae5140 100644 --- a/src/tts.rs +++ b/src/tts.rs @@ -1,147 +1,291 @@ +use std::io::Cursor; use std::path::PathBuf; -use std::process::Stdio; -use std::time::Instant; +use std::sync::OnceLock; +use std::time::{Duration, Instant}; -use tokio::io::AsyncWriteExt; +use dashmap::DashMap; +use piper_rs::Piper; +use tokio::sync::Mutex; use tracing::{debug, info, warn}; -use crate::config; -use crate::config::TEMP_PATH; use crate::error::{Error, Result}; -/// Synthesize speech for `text`. -/// -/// Primary engine is local Piper TTS (via the `piper` CLI). If Piper is not -/// installed, fails, or `TTS_MODEL_PATH` is unset, we fall back to Google -/// Translate TTS over HTTP. +// ── Constants ────────────────────────────────────────────────────────── + +/// Target sample rate for songbird playback. +const TARGET_SAMPLE_RATE: u32 = 48_000; + +/// Cache entries older than this are evicted on access. +const CACHE_TTL: Duration = Duration::from_secs(3600); + +/// All voice-state TTS event types. +const EVENT_TYPES: &[&str] = &[ + "joined", + "left", + "deafened", + "undeafened", + "muted", + "unmuted", +]; + +// ── TtsState ─────────────────────────────────────────────────────────── + +/// Shared TTS state: piper model instance + audio cache. /// -/// Returns the encoded audio bytes plus the file extension to use when saving -/// (`"wav"` for Piper, `"mp3"` for Google). -pub async fn synthesize( - config: &config::Config, - http_client: &reqwest::Client, - text: &str, -) -> Result<(Vec, &'static str)> { - match config.piper_model_path() { - Some(model_path) => { - let started = Instant::now(); - match piper_synthesize(&model_path, text).await { - Ok(bytes) => { - info!( - model = %model_path.display(), - text_len = text.len(), - bytes = bytes.len(), - elapsed_ms = started.elapsed().as_millis() as u64, - "Piper TTS synthesized" - ); - return Ok((bytes, "wav")); - } - Err(e) => { - warn!( - model = %model_path.display(), - text_len = text.len(), - elapsed_ms = started.elapsed().as_millis() as u64, - "Piper TTS failed, falling back to Google TTS: {e}" - ); - } - } +/// Held in poise `Data`. Piper is loaded lazily on first synthesis call. +/// Cache maps full text string → resampled 48 kHz stereo PCM bytes. +pub struct TtsState { + piper: OnceLock>, + cache: DashMap, Instant)>, + model_path: Option, +} + +impl TtsState { + pub fn new(model_path: Option) -> Self { + Self { + piper: OnceLock::new(), + cache: DashMap::new(), + model_path, } - None => { - debug!("TTS_MODEL_PATH is not set; using Google TTS fallback"); + } + + /// Get or initialize the Piper model. Returns error if model path is + /// unset or model fails to load. + async fn get_piper(&self) -> Result<&Mutex> { + if let Some(piper) = self.piper.get() { + return Ok(piper); } + + let path = self + .model_path + .as_ref() + .ok_or_else(|| Error::Custom("TTS_MODEL_PATH is not set".into()))?; + + // Resolve .onnx.json config path from model path + let onnx_path = path.clone(); + let config_path = { + let s = path.to_string_lossy(); + if s.ends_with(".onnx") { + PathBuf::from(format!("{}.json", s)) + } else { + path.clone() + } + }; + + debug!( + onnx = %onnx_path.display(), + config = %config_path.display(), + "loading piper model" + ); + let piper = Piper::new(&onnx_path, &config_path).map_err(|e| { + Error::Custom(format!("Failed to load piper model: {e}")) + })?; + info!("piper model loaded successfully"); + + // set() may fail if another task initialized first — that's fine + let _ = self.piper.set(Mutex::new(piper)); + Ok(self.piper.get().unwrap()) } - let started = Instant::now(); - let bytes = google_synthesize(http_client, text).await?; - info!( - text_len = text.len(), - bytes = bytes.len(), - elapsed_ms = started.elapsed().as_millis() as u64, - "Google TTS synthesized" - ); - Ok((bytes, "mp3")) -} + /// Look up cached audio for a text string. Returns `None` if uncached + /// or expired. + fn cache_lookup(&self, text: &str) -> Option> { + let entry = self.cache.get(text)?; + if entry.1.elapsed() > CACHE_TTL { + drop(entry); + self.cache.remove(text); + None + } else { + Some(entry.0.clone()) + } + } + + /// Store audio bytes in the cache. + fn cache_store(&self, text: String, audio: Vec) { + self.cache.insert(text, (audio, Instant::now())); + } -/// Run `piper` as a subprocess, feeding `text` on stdin and collecting the -/// generated WAV from `--output_file`. -async fn piper_synthesize(model_path: &PathBuf, text: &str) -> Result> { - tokio::fs::create_dir_all(TEMP_PATH).await.ok(); - - let out_path = format!("{}/tts_piper_{}.wav", TEMP_PATH, uuid::Uuid::new_v4()); - let model = model_path.to_string_lossy().to_string(); - - debug!(model = %model, "spawning piper subprocess"); - let started = Instant::now(); - - let mut child = tokio::process::Command::new("piper") - .arg("--model") - .arg(&model) - .arg("--output_file") - .arg(&out_path) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .map_err(|e| Error::Custom(crate::lang::fmt_tts_piper_launch(&e.to_string())))?; - - { - let mut stdin = child - .stdin - .take() - .ok_or_else(|| Error::Custom("Piper stdin unavailable".into()))?; - stdin.write_all(text.as_bytes()).await?; - stdin.write_all(b"\n").await?; + /// Evict all expired cache entries. + fn cache_evict(&self) { + self.cache + .retain(|_, (_, inserted_at)| inserted_at.elapsed() <= CACHE_TTL); } - let output = child.wait_with_output().await?; - let elapsed_ms = started.elapsed().as_millis() as u64; + /// Synthesize text → resampled 48 kHz stereo PCM bytes. + /// + /// Checks cache first. On miss, runs piper inference, resamples, caches, + /// and returns. + pub async fn synthesize(&self, text: &str) -> Result> { + // Cache hit + if let Some(cached) = self.cache_lookup(text) { + debug!(text_len = text.len(), "tts cache hit"); + return Ok(cached); + } + + // Cache miss — synthesize + let piper = self.get_piper().await?; + let started = Instant::now(); - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - warn!( - model = %model, + let (samples, sample_rate) = { + let mut lock = piper.lock().await; + lock.create(text, false, None, None, None, None) + .map_err(|e| Error::Custom(format!("piper synthesis failed: {e}")))? + }; + + let elapsed_ms = started.elapsed().as_millis() as u64; + debug!( + text_len = text.len(), + samples = samples.len(), + sample_rate, elapsed_ms, - "piper exited with error: {}", - stderr.trim_end() + "piper synthesized" ); - let _ = tokio::fs::remove_file(&out_path).await; - return Err(Error::Custom(crate::lang::fmt_tts_piper_failed( - stderr.trim_end(), - ))); + + // Resample mono → stereo at target sample rate + let pcm_bytes = resample_to_stereo_bytes(&samples, sample_rate, TARGET_SAMPLE_RATE); + + self.cache_store(text.to_string(), pcm_bytes.clone()); + Ok(pcm_bytes) } - let bytes = tokio::fs::read(&out_path).await?; - debug!( - model = %model, - bytes = bytes.len(), - elapsed_ms, - "piper succeeded" - ); - let _ = tokio::fs::remove_file(&out_path).await; - Ok(bytes) + /// Pre-generate all 6 event phrases for a display name. + /// + /// Called on first voice event for a user. Synthesizes all phrases and + /// caches them so subsequent events are instant. + pub async fn pre_generate(&self, display_name: &str) { + let phrases: Vec = EVENT_TYPES + .iter() + .map(|evt| make_event_phrase(display_name, evt)) + .collect(); + + for phrase in &phrases { + if self.cache_lookup(phrase).is_some() { + continue; + } + if let Err(e) = self.synthesize(phrase).await { + warn!( + display_name, + phrase, "pre-generation failed: {e}" + ); + } + } + + // Evict expired entries while we're at it + self.cache_evict(); + } } -/// Fetch Google Translate TTS (Polish) over HTTP. -async fn google_synthesize(client: &reqwest::Client, text: &str) -> Result> { - let tts_url = format!( - "https://translate.google.com/translate_tts?ie=UTF-8&q={}&tl=pl&client=tw-ob", - urlencoding(text) - ); - debug!(url_len = tts_url.len(), "fetching Google TTS"); - let bytes = client.get(&tts_url).send().await?.bytes().await?; - Ok(bytes.to_vec()) +// ── Phrase construction ──────────────────────────────────────────────── + +/// Build the TTS phrase for a voice event + display name. +pub fn make_event_phrase(display_name: &str, event: &str) -> String { + let suffix = match event { + "joined" => "sie zespawnił", + "left" => "jest chujem", + "deafened" => "jest bity przez ojca", + "undeafened" => "uciekł od ojca", + "muted" => "idzie robić jajecznice", + "unmuted" => "zrobił jajecznice", + _ => return format!("{display_name} {event}"), + }; + format!("{display_name} {suffix}") } -pub fn urlencoding(s: &str) -> String { - let mut out = String::with_capacity(s.len()); - for &byte in s.as_bytes() { - match byte { - b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'_' | b'-' | b'.' | b'~' => { - out.push(byte as char) - } - b' ' => out.push('+'), - _ => out.push_str(&format!("%{:02X}", byte)), +// ── Resampling ───────────────────────────────────────────────────────── + +/// Resample mono f32 PCM from `src_rate` to `dst_rate`, duplicate to stereo, +/// and return as little-endian bytes. +fn resample_to_stereo_bytes(samples: &[f32], src_rate: u32, dst_rate: u32) -> Vec { + if src_rate == dst_rate { + // Just duplicate mono → stereo, no resampling needed + let mut out = Vec::with_capacity(samples.len() * 2 * 4); + for &s in samples { + out.extend_from_slice(&s.to_le_bytes()); + out.extend_from_slice(&s.to_le_bytes()); } + return out; } + + let ratio = src_rate as f64 / dst_rate as f64; + let out_len = (samples.len() as f64 / ratio).ceil() as usize; + let mut out = Vec::with_capacity(out_len * 2 * 4); + + for i in 0..out_len { + let src_idx = i as f64 * ratio; + let idx = src_idx as usize; + let frac = src_idx - idx as f64; + + let sample = if idx + 1 < samples.len() { + // Linear interpolation + samples[idx] + (samples[idx + 1] - samples[idx]) * frac as f32 + } else if idx < samples.len() { + samples[idx] + } else { + 0.0 + }; + + // Mono → stereo duplication + out.extend_from_slice(&sample.to_le_bytes()); + out.extend_from_slice(&sample.to_le_bytes()); + } + out } + +// ── WAV encoding ─────────────────────────────────────────────────────── + +/// Encode raw f32 stereo PCM bytes into a WAV file (16-bit PCM). +/// +/// `pcm_f32_bytes` is the raw little-endian f32 interleaved stereo data +/// (e.g. from `resample_to_stereo_bytes`). +pub fn encode_wav(pcm_f32_bytes: &[u8], sample_rate: u32, channels: u16) -> Vec { + let bits_per_sample: u16 = 16; + let byte_rate = sample_rate * channels as u32 * bits_per_sample as u32 / 8; + let block_align = channels * bits_per_sample / 8; + + // Convert f32 samples to i16 + let sample_count = pcm_f32_bytes.len() / 4; // 4 bytes per f32 + let mut pcm_i16 = Vec::with_capacity(sample_count * 2); + for chunk in pcm_f32_bytes.chunks_exact(4) { + let f = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); + let i = (f * 32767.0).clamp(-32768.0, 32767.0) as i16; + pcm_i16.extend_from_slice(&i.to_le_bytes()); + } + + let data_size = pcm_i16.len() as u32; + let file_size = 36 + data_size; + + let mut wav = Vec::with_capacity(44 + pcm_i16.len()); + + // RIFF header + wav.extend_from_slice(b"RIFF"); + wav.extend_from_slice(&file_size.to_le_bytes()); + wav.extend_from_slice(b"WAVE"); + + // fmt chunk + wav.extend_from_slice(b"fmt "); + wav.extend_from_slice(&16u32.to_le_bytes()); // chunk size + wav.extend_from_slice(&1u16.to_le_bytes()); // PCM format + wav.extend_from_slice(&channels.to_le_bytes()); + wav.extend_from_slice(&sample_rate.to_le_bytes()); + wav.extend_from_slice(&byte_rate.to_le_bytes()); + wav.extend_from_slice(&block_align.to_le_bytes()); + wav.extend_from_slice(&bits_per_sample.to_le_bytes()); + + // data chunk + wav.extend_from_slice(b"data"); + wav.extend_from_slice(&data_size.to_le_bytes()); + wav.extend_from_slice(&pcm_i16); + + wav +} + +// ── Songbird Input helper ────────────────────────────────────────────── + +/// Create a songbird `Input` from resampled f32 stereo PCM bytes. +pub fn pcm_bytes_to_input(pcm_f32_bytes: Vec) -> songbird::input::Input { + use songbird::input::RawAdapter; + let cursor = Cursor::new(pcm_f32_bytes); + // 2 channels (stereo), 48 kHz — matches TARGET_SAMPLE_RATE + RawAdapter::new(cursor, TARGET_SAMPLE_RATE, 2).into() +} -- cgit v1.2.3