Skip to main content

xil/
tts.rs

1//! Dialogue TTS backends shared by `sample` and `produce`: the Chatterbox
2//! Turbo worker bridge and the gTTS draft backend. Ports of `_ChatterboxClient`
3//! and `_gtts_generate` (the two copies in the Python are identical).
4
5use std::fs;
6use std::path::{Path, PathBuf};
7
8use anyhow::bail;
9use regex::Regex;
10use serde_json::{Map, Value};
11use xil_audio::fx::py_repr;
12use xil_core::fsutil::basename;
13use xil_core::log;
14use xil_workers::Worker;
15
16const LABEL: &str = "Chatterbox Turbo";
17
18/// The two Python copies of `_ChatterboxClient` differ in two details.
19#[derive(Clone, Copy, PartialEq)]
20pub enum Flavor {
21    /// `XILU004`: reads exactly one line per response.
22    Sample,
23    /// `XILP002`: skips non-JSON noise between responses, and names the
24    /// venv in its startup failure.
25    Produce,
26}
27
28/// `_ChatterboxClient`: starts `chatterbox_turbo_worker.py` on first use.
29pub struct Chatterbox {
30    python: String,
31    voice_refs_dir: String,
32    device: String,
33    flavor: Flavor,
34    worker: Option<Worker>,
35}
36
37impl Chatterbox {
38    pub fn new(python: &str, voice_refs_dir: &str, device: &str, flavor: Flavor) -> Chatterbox {
39        Chatterbox {
40            python: python.to_string(),
41            voice_refs_dir: voice_refs_dir.to_string(),
42            device: device.to_string(),
43            flavor,
44            worker: None,
45        }
46    }
47
48    fn start(&mut self) -> anyhow::Result<()> {
49        let script = crate::workers::python_package_dir()
50            .map(|d| xil_workers::worker_script(&d, "chatterbox_turbo_worker.py"))
51            .ok_or_else(|| {
52                anyhow::anyhow!(
53                    "cannot locate the xil_pipeline package that ships chatterbox_turbo_worker.py"
54                )
55            })?;
56        log::info(&format!(
57            "Starting {LABEL} worker ({}, {})…",
58            self.python, self.device
59        ));
60        let mut w = Worker::spawn(
61            Path::new(&self.python),
62            &script,
63            std::slice::from_ref(&self.device),
64        )?;
65        let Some(msg) =
66            w.wait_ready(|line| log::debug(&format!("{LABEL} worker startup: {line}")))?
67        else {
68            if self.flavor == Flavor::Produce {
69                bail!("RuntimeError: {LABEL} worker exited before sending ready signal. Check that venv-chatterbox is correctly set up and the model is downloaded.");
70            }
71            bail!("RuntimeError: {LABEL} worker exited before sending ready signal.");
72        };
73        let actual = match msg.get("device") {
74            Some(Value::String(s)) => s.clone(),
75            Some(other) => other.to_string(),
76            None => self.device.clone(),
77        };
78        if actual != self.device {
79            log::warning(&format!(
80                "{LABEL}: requested device {} unavailable, running on {}",
81                py_repr(&self.device),
82                py_repr(&actual)
83            ));
84        }
85        let Some(sr) = msg.get("sr").and_then(Value::as_i64) else {
86            bail!("KeyError: 'sr'");
87        };
88        log::info(&format!(
89            "{LABEL} worker ready (sample_rate={sr}, device={actual})"
90        ));
91        self.worker = Some(w);
92        Ok(())
93    }
94
95    fn ref_for(&self, speaker: &str) -> Option<PathBuf> {
96        [".wav", ".mp3"]
97            .iter()
98            .map(|ext| Path::new(&self.voice_refs_dir).join(format!("{speaker}{ext}")))
99            .find(|p| p.exists())
100    }
101
102    /// `generate(text, out_path, speaker_key)`.
103    pub fn generate(&mut self, text: &str, out_path: &Path, speaker: &str) -> anyhow::Result<()> {
104        if self.worker.is_none() {
105            self.start()?;
106        }
107        let r = self.ref_for(speaker);
108        if let Some(r) = &r {
109            log::info(&format!("   ref: {}", basename(r)));
110        }
111        let mut req = Map::new();
112        req.insert("text".into(), text.into());
113        req.insert(
114            "out_path".into(),
115            out_path.to_string_lossy().into_owned().into(),
116        );
117        req.insert(
118            "ref_audio".into(),
119            r.map(|p| Value::from(p.to_string_lossy().into_owned()))
120                .unwrap_or(Value::Null),
121        );
122        req.insert(
123            "cond_path".into(),
124            Path::new(&self.voice_refs_dir)
125                .join(format!("{speaker}.turbo.conds.pt"))
126                .to_string_lossy()
127                .into_owned()
128                .into(),
129        );
130        let flavor = self.flavor;
131        let w = self.worker.as_mut().expect("started");
132        w.send(&Value::Object(req))?;
133        let resp: Value = loop {
134            let raw = w.read_line()?;
135            if raw.is_empty() {
136                bail!("RuntimeError: {LABEL} worker closed pipe unexpectedly.");
137            }
138            if flavor == Flavor::Sample {
139                break serde_json::from_str(&raw)?;
140            }
141            let line = raw.trim();
142            if line.is_empty() {
143                continue;
144            }
145            match serde_json::from_str(line) {
146                Ok(v) => break v,
147                Err(_) => log::debug(&format!("{LABEL} worker: {line}")),
148            }
149        };
150        if let Some(err) = resp.get("error") {
151            let msg = match err {
152                Value::String(s) => s.clone(),
153                other => other.to_string(),
154            };
155            bail!("RuntimeError: {LABEL}: {msg}");
156        }
157        Ok(())
158    }
159
160    pub fn close(&mut self) {
161        if let Some(w) = self.worker.take() {
162            let _ = w.close();
163        }
164    }
165}
166
167/// `_gtts_generate(text, out_path)`: strip `[tags]`, synthesize, atomic write.
168/// The producer's copy also logs the request.
169pub fn gtts_generate(text: &str, out_path: &Path, flavor: Flavor) -> anyhow::Result<()> {
170    static TAGS: std::sync::LazyLock<Regex> =
171        std::sync::LazyLock::new(|| Regex::new(r"\[[^\]]*\]").unwrap());
172    let cleaned = TAGS.replace_all(text, "").trim().to_string();
173    if cleaned.is_empty() {
174        return Ok(());
175    }
176    let dir = out_path
177        .parent()
178        .filter(|p| !p.as_os_str().is_empty())
179        .unwrap_or(Path::new("."));
180    let tmp = tempfile::Builder::new()
181        .prefix("tmp")
182        .rand_bytes(8)
183        .suffix(".tmp")
184        .tempfile_in(dir)?;
185    if flavor == Flavor::Produce {
186        log::info(&format!(
187            "   > gTTS ({} chars) → {}",
188            cleaned.chars().count(),
189            basename(out_path)
190        ));
191    }
192    let audio = xil_api::gtts::synthesize(&cleaned, "en")
193        .map_err(|e| anyhow::anyhow!("gtts.tts.gTTSError: {e}"))?;
194    fs::write(tmp.path(), audio)?;
195    tmp.persist(out_path)
196        .map_err(|e| anyhow::anyhow!(e.error))?;
197    Ok(())
198}