1use 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#[derive(Clone, Copy, PartialEq)]
20pub enum Flavor {
21 Sample,
23 Produce,
26}
27
28pub 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 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
167pub 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}