1use std::path::PathBuf;
6
7use forge_dictate::{ModelFacts, ModelSpec};
8use serde::Deserialize;
9
10use crate::dictate::{DictateModel, DictateModelState, DictateRole};
11
12pub const SPEAKING_MODEL: &str = "hexgrad/Kokoro-82M";
15
16pub const HUB_REPO: &str = "onnx-community/Kokoro-82M-v1.0-ONNX";
19
20pub const PINNED_VARIANT: &str = "kokoro-82m-int8";
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct SpeakingVariant {
26 pub variant: &'static str,
28 pub path: &'static str,
30 pub file: &'static str,
32 pub quant: &'static str,
33 pub size: u64,
35 pub sha256: &'static str,
37}
38
39pub const VARIANTS: &[SpeakingVariant] = &[
42 SpeakingVariant {
43 variant: "kokoro-82m-fp32",
44 path: "onnx/model.onnx",
45 file: "model.onnx",
46 quant: "fp32",
47 size: 325_532_232,
48 sha256: "8fbea51ea711f2af382e88c833d9e288c6dc82ce5e98421ea61c058ce21a34cb",
49 },
50 SpeakingVariant {
51 variant: "kokoro-82m-fp16",
52 path: "onnx/model_fp16.onnx",
53 file: "model_fp16.onnx",
54 quant: "fp16",
55 size: 163_234_740,
56 sha256: "ba4527a874b42b21e35f468c10d326fdff3c7fc8cac1f85e9eb6c0dfc35c334a",
57 },
58 SpeakingVariant {
59 variant: "kokoro-82m-int8",
60 path: "onnx/model_quantized.onnx",
61 file: "model_quantized.onnx",
62 quant: "int8",
63 size: 92_361_116,
64 sha256: "fbae9257e1e05ffc727e951ef9b9c98418e6d79f1c9b6b13bd59f5c9028a1478",
65 },
66];
67
68impl SpeakingVariant {
69 pub fn named(name: &str) -> Option<Self> {
72 if name == SPEAKING_MODEL {
73 return Some(Self::pinned());
74 }
75 VARIANTS.iter().copied().find(|variant| variant.variant == name)
76 }
77
78 pub fn pinned() -> Self {
80 VARIANTS
81 .iter()
82 .copied()
83 .find(|variant| variant.variant == PINNED_VARIANT)
84 .unwrap_or(VARIANTS[0])
85 }
86
87 pub fn spec(&self) -> ModelSpec {
89 ModelSpec {
90 file: self.file.to_owned(),
91 url: format!("https://huggingface.co/{HUB_REPO}/resolve/main/{}", self.path),
92 size: self.size,
93 sha256: Some(self.sha256.to_owned()),
94 facts: ModelFacts {
95 quant: Some(self.quant.to_owned()),
96 params: Some(82_000_000),
97 license: Some("Apache-2.0".to_owned()),
98 runtime: Some("onnx".to_owned()),
99 },
100 }
101 }
102}
103
104#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
106pub struct VoiceFile {
107 pub name: String,
109 pub size: u64,
110 pub sha256: String,
111}
112
113#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
115pub struct SpeakingInstall {
116 pub variant: String,
117 pub revision: String,
119 pub model: String,
121 pub model_size: u64,
122 pub model_sha256: String,
123 pub voices: Vec<VoiceFile>,
125 pub at: String,
127}
128
129impl SpeakingInstall {
130 pub fn spec(&self) -> ModelSpec {
132 let variant = SpeakingVariant::named(&self.variant).unwrap_or_else(SpeakingVariant::pinned);
133 ModelSpec {
134 file: self.model.clone(),
135 url: format!(
136 "https://huggingface.co/{HUB_REPO}/resolve/{}/{}",
137 self.revision, variant.path
138 ),
139 size: self.model_size,
140 sha256: Some(self.model_sha256.clone()),
141 facts: variant.spec().facts,
142 }
143 }
144
145 pub fn voice_names(&self) -> Vec<String> {
147 self.voices.iter().map(|voice| voice.name.clone()).collect()
148 }
149}
150
151#[derive(Debug, Clone, Deserialize, PartialEq)]
155#[serde(deny_unknown_fields)]
156pub struct TtsSettings {
157 #[serde(default)]
160 pub enabled: bool,
161 #[serde(default)]
163 pub models_dir: Option<String>,
164 #[serde(default)]
167 pub speaking_model: Option<String>,
168 #[serde(default)]
171 pub voice: Option<String>,
172 #[serde(default)]
176 pub language: Option<toml::Value>,
177 #[serde(default = "default_max_utterance_seconds")]
180 pub max_utterance_seconds: u64,
181}
182
183fn default_max_utterance_seconds() -> u64 {
184 120
185}
186
187impl Default for TtsSettings {
188 fn default() -> Self {
189 Self {
190 enabled: false,
191 models_dir: None,
192 speaking_model: None,
193 voice: None,
194 language: None,
195 max_utterance_seconds: default_max_utterance_seconds(),
196 }
197 }
198}
199
200impl TtsSettings {
201 pub fn configured_variant(&self) -> Result<Option<SpeakingVariant>, String> {
204 let Some(name) = self.speaking_model.as_deref() else {
205 return Ok(None);
206 };
207 SpeakingVariant::named(name)
208 .map(Some)
209 .ok_or_else(|| format!("[tts] speaking_model names {name}, which is not a known model"))
210 }
211
212 pub fn models_dir(&self) -> Option<PathBuf> {
214 match self.models_dir.as_deref() {
215 Some(dir) => Some(crate::config::expand_home(dir)),
216 None => dirs::cache_dir().map(|dir| dir.join("forge-tts")),
217 }
218 }
219
220 pub fn voices_dir(&self) -> Option<PathBuf> {
222 self.models_dir().map(|dir| dir.join("voices"))
223 }
224
225 pub fn pin(&self) -> ModelSpec {
228 self.configured_variant().ok().flatten().unwrap_or_else(SpeakingVariant::pinned).spec()
229 }
230}
231
232#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
234pub struct TtsSnapshot {
235 pub models: Vec<DictateModel>,
236}
237
238#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
240pub struct SpeakingVariantView {
241 pub variant: String,
242 pub file: String,
243 pub quant: String,
244 pub size: u64,
245 pub installed: bool,
246 pub active: bool,
247}
248
249#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
252#[serde(rename_all = "snake_case")]
253pub enum SpeakingFrom {
254 Config,
256 Installed,
258 Pin,
260}
261
262#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
265pub struct TtsView {
266 pub enabled: bool,
267 pub snapshot: TtsSnapshot,
268 pub models_dir: Option<PathBuf>,
269 pub speaking_model: String,
271 pub active: SpeakingVariantView,
273 pub from: SpeakingFrom,
275 pub variants: Vec<SpeakingVariantView>,
277 pub voices: Vec<String>,
279 pub voice: Option<String>,
281 pub preview_error: Option<String>,
284}
285
286impl crate::Workspace {
287 pub fn speaking_installs(&self) -> Vec<SpeakingInstall> {
289 let db = self.db.lock();
290 let Some(db) = db.as_ref() else { return Vec::new() };
291 match crate::store::speaking_models::installed(db) {
292 Ok(rows) => rows,
293 Err(error) => {
294 tracing::warn!(
295 event_name = "speaking_installed_read_failed",
296 %error,
297 "the speaking installs could not be read; the page shows none"
298 );
299 Vec::new()
300 }
301 }
302 }
303
304 pub fn speaking_install(&self, variant: &str) -> Option<SpeakingInstall> {
306 self.speaking_installs().into_iter().find(|install| install.variant == variant)
307 }
308
309 pub fn speaking_variant(&self) -> (SpeakingVariant, SpeakingFrom) {
312 if let Ok(Some(variant)) = self.config_arc().tts.configured_variant() {
313 return (variant, SpeakingFrom::Config);
314 }
315 if let Some(pick) = self.active_choice(DictateRole::Speaking)
316 && let Some(variant) = SpeakingVariant::named(&pick.variant)
317 {
318 return (variant, SpeakingFrom::Installed);
319 }
320 (SpeakingVariant::pinned(), SpeakingFrom::Pin)
321 }
322
323 pub fn speaking_files(&self) -> Option<forge_tts::ModelFiles> {
325 let (variant, _) = self.speaking_variant();
326 let install = self.speaking_install(variant.variant)?;
327 self.files_for_install(&install)
328 }
329
330 pub fn files_for_install(&self, install: &SpeakingInstall) -> Option<forge_tts::ModelFiles> {
332 let dir = self.config_arc().tts.models_dir()?;
333 let model = dir.join(&install.model);
334 let voices = dir.join("voices");
335 let digests =
336 std::iter::once((model.clone(), install.model_sha256.clone()))
337 .chain(install.voices.iter().map(|voice| {
338 (voices.join(format!("{}.bin", voice.name)), voice.sha256.clone())
339 }))
340 .collect();
341 Some(forge_tts::ModelFiles::new(model, voices).with_digests(digests))
342 }
343
344 pub fn voices(&self) -> TtsView {
346 let config = self.config_arc();
347 let settings = &config.tts;
348 let (active, from) = self.speaking_variant();
349 let installs = self.speaking_installs();
350 let installed = |variant: &str| installs.iter().any(|row| row.variant == variant);
351 let row = |variant: &SpeakingVariant| SpeakingVariantView {
352 variant: variant.variant.to_owned(),
353 file: variant.file.to_owned(),
354 quant: variant.quant.to_owned(),
355 size: variant.size,
356 installed: installed(variant.variant),
357 active: variant.variant == active.variant,
358 };
359 let active_installed = installed(active.variant);
360 let active_voices = installs
361 .iter()
362 .find(|row| row.variant == active.variant)
363 .map(SpeakingInstall::voice_names)
364 .unwrap_or_default();
365 let state =
366 if active_installed { DictateModelState::Ready } else { DictateModelState::Pending };
367 let enabled = self.voice_enabled();
370 TtsView {
371 enabled,
372 snapshot: TtsSnapshot {
373 models: if enabled {
374 vec![DictateModel {
375 role: DictateRole::Speaking,
376 file: active.file.to_owned(),
377 state,
378 }]
379 } else {
380 Vec::new()
381 },
382 },
383 models_dir: settings.models_dir(),
384 speaking_model: SPEAKING_MODEL.to_owned(),
385 active: row(&active),
386 from,
387 variants: VARIANTS.iter().map(row).collect(),
388 voices: active_voices,
389 voice: settings.voice.clone(),
390 preview_error: self.voice_preview_error.lock().ok().and_then(|slot| slot.clone()),
391 }
392 }
393}
394
395#[cfg(test)]
396mod tests {
397 use super::{PINNED_VARIANT, SPEAKING_MODEL, SpeakingVariant, TtsSettings};
398
399 #[test]
400 fn an_absent_tts_section_is_off() {
401 let settings: TtsSettings = toml::from_str("").expect("parses");
402 assert!(!settings.enabled);
403 assert_eq!(settings.max_utterance_seconds, 120);
404 assert!(settings.voice.is_none());
405 }
406
407 #[test]
408 fn an_unknown_tts_key_fails_the_load() {
409 let error = toml::from_str::<TtsSettings>("speek = \"x\"").expect_err("refused");
410 assert!(error.to_string().contains("speek"));
411 }
412
413 #[test]
414 fn the_models_own_name_selects_the_pinned_variant() {
415 let variant = SpeakingVariant::named(SPEAKING_MODEL).expect("the name is known");
416 assert_eq!(variant.variant, PINNED_VARIANT);
417 }
418
419 #[test]
420 fn a_variant_name_selects_itself() {
421 let variant = SpeakingVariant::named("kokoro-82m-fp16").expect("the name is known");
422 assert_eq!(variant.file, "model_fp16.onnx");
423 }
424
425 #[test]
429 fn the_speaking_source_crosses_as_a_bare_name() {
430 let value = serde_json::to_value(super::SpeakingFrom::Installed).expect("it serializes");
431 assert_eq!(value, serde_json::json!("installed"));
432 }
433
434 #[test]
435 fn a_configured_name_that_is_not_a_model_is_refused_by_name() {
436 let settings: TtsSettings =
437 toml::from_str("speaking_model = \"hexgrad/Kokoro-90M\"").expect("parses");
438 let error = settings.configured_variant().expect_err("refused");
439 assert!(error.contains("hexgrad/Kokoro-90M"), "got: {error}");
440 assert!(error.contains("speaking_model"), "got: {error}");
441 }
442
443 #[tokio::test]
447 async fn a_saved_tts_off_does_not_move_the_read_before_the_restart() {
448 let dir = tempfile::tempdir().expect("a temp dir");
449 let mut config = crate::config::LoadedConfig::empty_for_test();
450 config.tts.enabled = true;
451 let (workspace, _updates) =
452 crate::Workspace::testing_stub_with_config(dir.path().to_path_buf(), config)
453 .expect("the stub builds");
454 assert!(workspace.voice_enabled(), "the fixture starts on");
455 let _service = workspace.voice_service();
456
457 let mut saved = crate::config::LoadedConfig::empty_for_test();
459 saved.tts.enabled = false;
460 *workspace.config.write() = std::sync::Arc::new(saved);
461
462 assert!(workspace.voice_enabled(), "the running service still answers on");
463 assert!(workspace.voices().enabled, "the home read follows the service");
464 }
465}