Skip to main content

forge_workspace/
tts.rs

1//! The `[tts]` section: what the speaking role runs, the Hugging Face
2//! install that supplies it, and the read the models page and the settings
3//! card draw.
4
5use std::path::PathBuf;
6
7use forge_dictate::{ModelFacts, ModelSpec};
8use serde::Deserialize;
9
10use crate::dictate::{DictateModel, DictateModelState, DictateRole};
11
12/// The speaking model's own name. `[tts] speaking_model` accepts this, or one
13/// of [`VARIANTS`]' names to pin a quantisation.
14pub const SPEAKING_MODEL: &str = "hexgrad/Kokoro-82M";
15
16/// The Hub repository publishing the ONNX export and the style files the
17/// engine loads.
18pub const HUB_REPO: &str = "onnx-community/Kokoro-82M-v1.0-ONNX";
19
20/// The variant the role runs when nothing chooses another: the int8 export.
21pub const PINNED_VARIANT: &str = "kokoro-82m-int8";
22
23/// One ONNX variant of the model the Hub publishes.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct SpeakingVariant {
26    /// The name `[tts] speaking_model` accepts.
27    pub variant: &'static str,
28    /// The file's path in the repository.
29    pub path: &'static str,
30    /// The file's own name under the models directory.
31    pub file: &'static str,
32    pub quant: &'static str,
33    /// The release's byte length, for the row drawn before an install.
34    pub size: u64,
35    /// The release's LFS digest, for the row drawn before an install.
36    pub sha256: &'static str,
37}
38
39/// The variants the speaking role can run, the int8 model the benchmark
40/// measured among them.
41pub 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    /// The variant a name selects: the model's own name picks the pin, and a
70    /// variant name picks itself.
71    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    /// The compiled default.
79    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    /// The row a spec-driven read draws before an install lands.
88    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/// One style file the install fetched.
105#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
106pub struct VoiceFile {
107    /// The voice id: the file's stem.
108    pub name: String,
109    pub size: u64,
110    pub sha256: String,
111}
112
113/// What one installed variant is made of, as the install recorded it.
114#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
115pub struct SpeakingInstall {
116    pub variant: String,
117    /// The Hub revision every file was fetched at, so the set is pinned.
118    pub revision: String,
119    /// The model's file name under the models directory.
120    pub model: String,
121    pub model_size: u64,
122    pub model_sha256: String,
123    /// The style files, by voice id.
124    pub voices: Vec<VoiceFile>,
125    /// RFC 3339.
126    pub at: String,
127}
128
129impl SpeakingInstall {
130    /// The model's spec, off the record rather than off the release.
131    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    /// The voice ids this install carried.
146    pub fn voice_names(&self) -> Vec<String> {
147        self.voices.iter().map(|voice| voice.name.clone()).collect()
148    }
149}
150
151/// Deliberately narrower than a model spec, for the same reason
152/// [`crate::dictate::DictateSettings`] is: a hand-edited URL, byte length or
153/// digest is how a file arrives that nothing can verify.
154#[derive(Debug, Clone, Deserialize, PartialEq)]
155#[serde(deny_unknown_fields)]
156pub struct TtsSettings {
157    /// Off unless asked for. Turning it on makes the speaking role available;
158    /// the model installs from the models page or the settings card.
159    #[serde(default)]
160    pub enabled: bool,
161    /// Where the model files live. Absent uses the platform cache directory.
162    #[serde(default)]
163    pub models_dir: Option<String>,
164    /// The model the role runs. Absent uses the compiled pin; an unknown
165    /// name is refused by name.
166    #[serde(default)]
167    pub speaking_model: Option<String>,
168    /// The voice every utterance renders in. Absent uses the engine's own
169    /// default; an id that is not installed is refused by name.
170    #[serde(default)]
171    pub voice: Option<String>,
172    /// Ghost of the retired `[tts] language` hint: read only so a stale
173    /// config keeps loading. The voice id is the language choice (the 54
174    /// voices cover eight languages), so there is nothing for the key to set.
175    #[serde(default)]
176    pub language: Option<toml::Value>,
177    /// Upper bound on one utterance's text, as a character estimate at a
178    /// rough speaking pace, not one sentence.
179    #[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    /// The variant a configured name selects, refusing an unknown name the
202    /// way `[dictate] transcribe_model` does.
203    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    /// Where the model files land.
213    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    /// The directory the style files live in.
221    pub fn voices_dir(&self) -> Option<PathBuf> {
222        self.models_dir().map(|dir| dir.join("voices"))
223    }
224
225    /// The model a role runs before anything is installed: the configured
226    /// variant, else the pin.
227    pub fn pin(&self) -> ModelSpec {
228        self.configured_variant().ok().flatten().unwrap_or_else(SpeakingVariant::pinned).spec()
229    }
230}
231
232/// The speaking role's model rows.
233#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
234pub struct TtsSnapshot {
235    pub models: Vec<DictateModel>,
236}
237
238/// One variant the page draws: whether it is installed, and whether it runs.
239#[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/// Where the active speaking model came from. A bare name on the wire: the
250/// speaking read has one `from`, and a tag object would nest it.
251#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
252#[serde(rename_all = "snake_case")]
253pub enum SpeakingFrom {
254    /// `[tts] speaking_model` names it.
255    Config,
256    /// The last pick made on this page.
257    Installed,
258    /// The compiled default.
259    Pin,
260}
261
262/// The voice read: whether `[tts]` is on, where the files land, the speaking
263/// model's state, and the voices that model installed.
264#[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    /// The model's own name, which `[tts] speaking_model` accepts.
270    pub speaking_model: String,
271    /// The variant the role runs.
272    pub active: SpeakingVariantView,
273    /// Where the active variant's choice came from.
274    pub from: SpeakingFrom,
275    /// Every variant the role can run.
276    pub variants: Vec<SpeakingVariantView>,
277    /// The installed voice ids, sorted.
278    pub voices: Vec<String>,
279    /// The configured voice; absent means the engine's default.
280    pub voice: Option<String>,
281    /// The last preview's render failure, when one failed after the command
282    /// was accepted. Cleared by the next preview.
283    pub preview_error: Option<String>,
284}
285
286impl crate::Workspace {
287    /// Every speaking install this machine has recorded.
288    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    /// One installed variant, when this machine has it.
305    pub fn speaking_install(&self, variant: &str) -> Option<SpeakingInstall> {
306        self.speaking_installs().into_iter().find(|install| install.variant == variant)
307    }
308
309    /// The variant the role runs: the configured name, else the runtime
310    /// pick, else the compiled pin.
311    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    /// The files the active variant loads, when this machine has them.
324    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    /// The files one installed variant is made of.
331    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    /// The models page's and the settings card's read.
345    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        // The running system's flag, not a saved value the restart has not
368        // taken: the tool gate and the registered set read the same one.
369        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    /// The speaking read's `from` is a bare name: the client compares it
426    /// against a list of strings, and a tag object would read as unknown on
427    /// both pages.
428    #[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    /// A save marked restart cannot move the read away from the running
444    /// service: the home read and the registered set answer the same flag
445    /// until the restart takes it.
446    #[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        // The save lands, and the key is marked restart.
458        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}