Skip to main content

ravenrs_ai/
voice.rs

1//! A voice in use, and the screen for it. It speaks on its own thread, a sentence at a
2//! time, and lets go of its memory when nobody has asked anything for a while.
3
4use std::cell::RefCell;
5use std::collections::VecDeque;
6use std::fs;
7use std::path::PathBuf;
8use std::rc::Rc;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::mpsc::{self, RecvTimeoutError, Sender};
11use std::sync::{Arc, Mutex};
12use std::time::{Duration, Instant};
13
14use ravenrs_app_i18n::{Language, Texts};
15use ravenrs_layout::Align;
16use ravenrs_reactive::State;
17use ravenrs_style::{Semantic, TextStyle};
18use ravenrs_widgets::{
19    BuildCx, Button, ButtonSize, ButtonStyle, Disclosure, HStack, IconName, NodeId, SegmentedControl, Select,
20    Slider, Sound, SoundClip, Spacer, Text, Toggle, UiWaker, VStack, View,
21};
22use serde_json::{Value, json};
23
24use crate::{
25    KOKORO_SPEAKERS, LocalModel, Machine, ModelFamily, ModelFolder, ModelManager, ModelManagerLabels,
26    SPEECH_LANGUAGES, SPEECH_RATE, SpeechEngine, SpeechSettings, SpeechVoice,
27};
28
29const STAYS_LOADED: Duration = Duration::from_secs(120);
30const SETTINGS_FILE: &str = "voice-settings.json";
31
32struct Saying {
33    text: String,
34    settings: SpeechSettings,
35    waker: UiWaker,
36    turn: u64,
37}
38
39/// A model that speaks on its own thread. Dropping it ends the thread and frees the
40/// memory.
41#[derive(Clone)]
42pub struct LocalVoice {
43    jobs: Sender<Saying>,
44    turn: Arc<AtomicU64>,
45    ready: Arc<Mutex<VecDeque<SoundClip>>>,
46}
47
48impl std::fmt::Debug for LocalVoice {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        f.debug_struct("LocalVoice").finish_non_exhaustive()
51    }
52}
53
54impl LocalVoice {
55    /// A voice over `model`. Nothing is loaded until the first thing it says.
56    pub fn new(folder: ModelFolder, model: LocalModel) -> Self {
57        let (jobs, inbox) = mpsc::channel::<Saying>();
58        let (turn, ready) = (Arc::new(AtomicU64::new(0)), Arc::new(Mutex::new(VecDeque::new())));
59        let (now, out) = (Arc::clone(&turn), Arc::clone(&ready));
60        std::thread::spawn(move || {
61            let mut engine: Option<SpeechEngine> = None;
62            loop {
63                let job = if engine.is_some() {
64                    match inbox.recv_timeout(STAYS_LOADED) {
65                        Ok(job) => job,
66                        Err(RecvTimeoutError::Timeout) => {
67                            engine = None;
68                            continue;
69                        }
70                        Err(RecvTimeoutError::Disconnected) => return,
71                    }
72                } else {
73                    match inbox.recv() {
74                        Ok(job) => job,
75                        Err(_) => return,
76                    }
77                };
78                if engine.is_none() && job.turn == now.load(Ordering::Acquire) {
79                    engine = SpeechEngine::load(&folder, &model).ok();
80                }
81                let Some(speaking) = engine.as_mut() else { continue };
82                let mut free_at = Instant::now();
83                for sentence in speaking.pieces(&job.text) {
84                    let current = || job.turn == now.load(Ordering::Acquire);
85                    if !current() {
86                        break;
87                    }
88                    let Ok(sound) = speaking.speak(&sentence, &job.settings, &mut |_| current()) else {
89                        break;
90                    };
91                    if sound.is_empty() || !current() {
92                        continue;
93                    }
94                    // One after the other: the sound system plays clips over each other.
95                    std::thread::sleep(free_at.saturating_duration_since(Instant::now()));
96                    if !current() {
97                        break;
98                    }
99                    let clip = SoundClip::new(sound, SPEECH_RATE);
100                    free_at = Instant::now() + Duration::from_secs_f32(clip.seconds());
101                    if let Ok(mut ready) = out.lock() {
102                        ready.push_back(clip);
103                    }
104                    job.waker.wake();
105                }
106            }
107        });
108        Self { jobs, turn, ready }
109    }
110
111    /// Says `text`, in place of anything still being said. `waker` asks for a rebuild
112    /// when a sentence is ready, for [`LocalVoice::play_ready`] to start it.
113    pub fn say(&self, text: String, settings: SpeechSettings, waker: UiWaker) {
114        let turn = self.next_turn();
115        // The thread ends only with this handle, so this does not fail.
116        let _ = self.jobs.send(Saying { text, settings, waker, turn });
117    }
118
119    /// Stops saying anything still to come.
120    pub fn stop(&self) {
121        self.next_turn();
122    }
123
124    fn next_turn(&self) -> u64 {
125        if let Ok(mut ready) = self.ready.lock() {
126            ready.clear();
127        }
128        self.turn.fetch_add(1, Ordering::AcqRel) + 1
129    }
130
131    /// Plays what is ready. Call it when building a view that stays on the screen
132    /// while it speaks.
133    pub fn play_ready(&self, sound: &Sound) {
134        let Ok(mut ready) = self.ready.lock() else { return };
135        while let Some(clip) = ready.pop_front() {
136            sound.play(&clip);
137        }
138    }
139}
140
141/// The sentences of `text`, which is how the model says long texts best.
142pub fn sentences(text: &str) -> Vec<String> {
143    let mut all = Vec::new();
144    let mut sentence = String::new();
145    let mut letters = text.chars().peekable();
146    while let Some(letter) = letters.next() {
147        let ends = matches!(letter, '.' | '!' | '?' | '\n')
148            && letters.peek().is_none_or(|next| next.is_whitespace());
149        if letter != '\n' {
150            sentence.push(letter);
151        }
152        if ends {
153            let done = sentence.trim();
154            if done.chars().any(char::is_alphanumeric) {
155                all.push(done.to_owned());
156            }
157            sentence.clear();
158        }
159    }
160    let done = sentence.trim();
161    if done.chars().any(char::is_alphanumeric) {
162        all.push(done.to_owned());
163    }
164    all
165}
166
167impl SpeechSettings {
168    /// What the model's makers measured it with: steadier, and less varied.
169    pub fn steadier() -> Self {
170        Self { temperature: 0.4, top_k: 0, repetition_penalty: 1.2, repetition_window: 64, ..Self::default() }
171    }
172}
173
174fn to_json(settings: &SpeechSettings, read_aloud: bool) -> Value {
175    json!({
176        "language": settings.language,
177        "voice": match settings.voice { SpeechVoice::Female => "female", SpeechVoice::Male => "male" },
178        "speaker": settings.speaker,
179        "temperature": settings.temperature,
180        "top_k": settings.top_k,
181        "repetition_penalty": settings.repetition_penalty,
182        "repetition_window": settings.repetition_window,
183        "read_aloud": read_aloud,
184    })
185}
186
187/// Reads settings back. What is missing or strange stays as recommended.
188fn from_json(value: &Value) -> (SpeechSettings, bool) {
189    let mut settings = SpeechSettings::default();
190    if let Some(language) = value.get("language").and_then(Value::as_str)
191        && SPEECH_LANGUAGES.contains(&language)
192    {
193        language.clone_into(&mut settings.language);
194    }
195    if value.get("voice").and_then(Value::as_str) == Some("male") {
196        settings.voice = SpeechVoice::Male;
197    }
198    if let Some(speaker) = value.get("speaker").and_then(Value::as_str)
199        && KOKORO_SPEAKERS.contains(&speaker)
200    {
201        settings.speaker = Some(speaker.to_owned());
202    }
203    let number = |name: &str| value.get(name).and_then(Value::as_f64).filter(|number| *number >= 0.0);
204    if let Some(number) = number("temperature") {
205        settings.temperature = number as f32;
206    }
207    if let Some(number) = number("top_k") {
208        settings.top_k = number as usize;
209    }
210    if let Some(number) = number("repetition_penalty").filter(|number| *number >= 1.0) {
211        settings.repetition_penalty = number as f32;
212    }
213    if let Some(number) = number("repetition_window") {
214        settings.repetition_window = number as usize;
215    }
216    (settings, value.get("read_aloud").and_then(Value::as_bool).unwrap_or(true))
217}
218
219/// What the voice screen says.
220#[derive(Clone, Debug)]
221#[non_exhaustive]
222pub struct VoiceLabels {
223    /// For the list of models.
224    pub models: ModelManagerLabels,
225    /// The language's field.
226    pub language: String,
227    /// The name of each of [`SPEECH_LANGUAGES`], in its order.
228    pub languages: Vec<String>,
229    /// The voice's field.
230    pub voice: String,
231    /// The female voice.
232    pub female: String,
233    /// The male voice.
234    pub male: String,
235    /// The field for the model's own voices.
236    pub speaker: String,
237    /// The name of each of [`KOKORO_SPEAKERS`], in its order.
238    pub speakers: Vec<String>,
239    /// The switch that reads answers aloud.
240    pub read_aloud: String,
241    /// The button that plays a sample.
242    pub listen: String,
243    /// What the sample says.
244    pub sample: String,
245    /// What the fine settings are called.
246    pub advanced: String,
247    /// How free each code is to vary, and what that does.
248    pub temperature: (String, String),
249    /// How many codes it chooses among, and what that does.
250    pub top_k: (String, String),
251    /// How much it avoids repeating, and what that does.
252    pub repetition_penalty: (String, String),
253    /// How far back it looks, and what that does.
254    pub repetition_window: (String, String),
255    /// The button for the recommended settings.
256    pub recommended: String,
257    /// The button for the steadier settings.
258    pub steadier: String,
259    /// Whether numbers use a comma.
260    pub decimal_comma: bool,
261}
262
263impl Default for VoiceLabels {
264    fn default() -> Self {
265        Self::from_texts(&Texts::new(Language::English))
266    }
267}
268
269impl VoiceLabels {
270    /// The texts of a catalog, in its language. What it does not have comes from English.
271    pub fn from_texts(texts: &Texts) -> Self {
272        Self {
273            models: ModelManagerLabels::from_texts(texts),
274            language: texts.get("ai.voice.language"),
275            languages: texts.list("ai.voice.languages"),
276            voice: texts.get("ai.voice.voice"),
277            female: texts.get("ai.voice.female"),
278            male: texts.get("ai.voice.male"),
279            speaker: texts.get("ai.voice.speaker"),
280            speakers: texts.list("ai.voice.speakers"),
281            read_aloud: texts.get("ai.voice.read_aloud"),
282            listen: texts.get("ai.voice.listen"),
283            sample: texts.get("ai.voice.sample"),
284            advanced: texts.get("ai.voice.advanced"),
285            temperature: texts.pair("ai.voice.temperature"),
286            top_k: texts.pair("ai.voice.top_k"),
287            repetition_penalty: texts.pair("ai.voice.repetition_penalty"),
288            repetition_window: texts.pair("ai.voice.repetition_window"),
289            recommended: texts.get("ai.voice.recommended"),
290            steadier: texts.get("ai.voice.steadier"),
291            decimal_comma: texts.char("format.decimal", '.') == ',',
292        }
293    }
294
295    /// In Brazilian Portuguese.
296    pub fn pt_br() -> Self {
297        Self::from_texts(&Texts::new(Language::BrazilianPortuguese))
298    }
299}
300
301/// The screen for the voice: which model speaks, and how. The values start as its model
302/// card recommends, and are kept in the models' folder.
303#[derive(Clone)]
304pub struct VoiceManager {
305    models: ModelManager,
306    folder: ModelFolder,
307    file: PathBuf,
308    voice: Rc<RefCell<Option<(&'static str, LocalVoice)>>>,
309    saved: Rc<RefCell<Value>>,
310    language: State<usize>,
311    gender: State<usize>,
312    speaker: State<usize>,
313    temperature: State<f32>,
314    top_k: State<f32>,
315    repetition_penalty: State<f32>,
316    repetition_window: State<f32>,
317    read_aloud: State<bool>,
318    labels: VoiceLabels,
319}
320
321impl std::fmt::Debug for VoiceManager {
322    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
323        f.debug_struct("VoiceManager").finish_non_exhaustive()
324    }
325}
326
327impl VoiceManager {
328    /// The models of `catalog` in `root`, as many as fit this machine.
329    pub fn new(catalog: &[LocalModel], root: impl Into<PathBuf>) -> Self {
330        let root = root.into();
331        Self::for_machine(catalog, root.clone(), &Machine::detect(&root))
332    }
333
334    /// The same, for `machine`.
335    pub fn for_machine(catalog: &[LocalModel], root: impl Into<PathBuf>, machine: &Machine) -> Self {
336        let root = root.into();
337        let file = root.join(SETTINGS_FILE);
338        let found = fs::read(&file).ok().and_then(|bytes| serde_json::from_slice::<Value>(&bytes).ok());
339        let (settings, read_aloud) =
340            found.as_ref().map(from_json).unwrap_or((SpeechSettings::default(), true));
341        let folder = ModelFolder::new(root);
342        let language = SPEECH_LANGUAGES.iter().position(|name| *name == settings.language).unwrap_or(0);
343        Self {
344            models: ModelManager::for_machine(catalog, folder.clone(), machine),
345            folder,
346            file,
347            voice: Rc::default(),
348            saved: Rc::new(RefCell::new(to_json(&settings, read_aloud))),
349            language: State::new(language),
350            gender: State::new(usize::from(settings.voice == SpeechVoice::Male)),
351            speaker: State::new(
352                settings
353                    .speaker
354                    .as_deref()
355                    .and_then(|speaker| KOKORO_SPEAKERS.iter().position(|name| *name == speaker))
356                    .unwrap_or(0),
357            ),
358            temperature: State::new(settings.temperature),
359            top_k: State::new(settings.top_k as f32),
360            repetition_penalty: State::new(settings.repetition_penalty),
361            repetition_window: State::new(settings.repetition_window as f32),
362            read_aloud: State::new(read_aloud),
363            labels: VoiceLabels::default(),
364        }
365    }
366
367    /// What it says.
368    #[must_use]
369    pub fn labels(mut self, labels: VoiceLabels) -> Self {
370        self.models = self.models.labels(labels.models.clone());
371        self.labels = labels;
372        self
373    }
374
375    /// Whether the voice in use is Kokoro, which has voices of its own and speaks only
376    /// Portuguese.
377    fn kokoro(&self) -> bool {
378        self.models.model_in_use().is_some_and(|model| model.family == ModelFamily::Kokoro)
379    }
380
381    /// How to speak, as it is set now.
382    pub fn settings(&self) -> SpeechSettings {
383        let language = SPEECH_LANGUAGES.get(self.language.get()).copied().unwrap_or("Portuguese");
384        let kokoro = self.kokoro();
385        SpeechSettings {
386            language: if kokoro { "Portuguese" } else { language }.to_owned(),
387            voice: if self.gender.get() == 1 { SpeechVoice::Male } else { SpeechVoice::Female },
388            speaker: KOKORO_SPEAKERS
389                .get(self.speaker.get())
390                .filter(|_| kokoro)
391                .map(|name| (*name).to_owned()),
392            temperature: self.temperature.get(),
393            top_k: self.top_k.get().round() as usize,
394            repetition_penalty: self.repetition_penalty.get(),
395            repetition_window: self.repetition_window.get().round() as usize,
396            ..SpeechSettings::default()
397        }
398    }
399
400    /// The voice in use, if any.
401    pub fn voice(&self) -> Option<LocalVoice> {
402        let model = self.models.model_in_use()?;
403        let mut kept = self.voice.borrow_mut();
404        match kept.as_ref() {
405            Some((id, voice)) if *id == model.id => Some(voice.clone()),
406            _ => {
407                let voice = LocalVoice::new(self.folder.clone(), model);
408                *kept = Some((model.id, voice.clone()));
409                Some(voice)
410            }
411        }
412    }
413
414    /// Says `text` with the voice in use and the current settings, if there is one.
415    pub fn say(&self, text: String, waker: UiWaker) {
416        if let Some(voice) = self.voice() {
417            voice.say(text, self.settings(), waker);
418        }
419    }
420
421    /// Whether answers should be read when they arrive.
422    pub fn reads_aloud(&self) -> bool {
423        self.read_aloud.get() && self.models.model_in_use().is_some()
424    }
425
426    /// Plays what is ready. Call it when building a view that stays on the screen
427    /// while it speaks.
428    pub fn play_ready(&self, sound: &Sound) {
429        if let Some((_, voice)) = self.voice.borrow().as_ref() {
430            voice.play_ready(sound);
431        }
432    }
433
434    fn apply(&self, settings: &SpeechSettings) {
435        self.temperature.set(settings.temperature);
436        self.top_k.set(settings.top_k as f32);
437        self.repetition_penalty.set(settings.repetition_penalty);
438        self.repetition_window.set(settings.repetition_window as f32);
439    }
440
441    fn save_if_changed(&self) {
442        let now = to_json(&self.settings(), self.read_aloud.get());
443        if *self.saved.borrow() != now {
444            let bytes = serde_json::to_vec_pretty(&now).unwrap_or_default();
445            // Not saving only means going back to these values next time.
446            if let Some(parent) = self.file.parent() {
447                let _ = fs::create_dir_all(parent).and_then(|()| fs::write(&self.file, bytes));
448            }
449            *self.saved.borrow_mut() = now;
450        }
451    }
452
453    fn number(&self, value: f32) -> String {
454        let text = if value.fract().abs() < 1e-3 { format!("{value:.0}") } else { format!("{value:.2}") };
455        if self.labels.decimal_comma { text.replace('.', ",") } else { text }
456    }
457
458    fn slider(
459        &self,
460        names: &(String, String),
461        state: &State<f32>,
462        range: (f32, f32, f32),
463    ) -> impl View + 'static {
464        VStack::new((
465            HStack::new((
466                Text::new(names.0.clone()).style(TextStyle::Callout),
467                Spacer::new(),
468                Text::new(self.number(state.get())).style(TextStyle::Callout).foreground(Semantic::InkMuted),
469            ))
470            .align(Align::Center),
471            Slider::new(state.binding()).range(range.0, range.1).step(range.2).label(names.0.clone()),
472            Text::new(names.1.clone()).style(TextStyle::Caption).foreground(Semantic::InkMuted),
473        ))
474        .spacing(4.0)
475        .align(Align::Stretch)
476    }
477}
478
479impl View for VoiceManager {
480    fn build(&self, cx: &mut BuildCx<'_>) -> NodeId {
481        let theme = *cx.theme();
482        self.save_if_changed();
483        let sound = cx.sound();
484        self.play_ready(&sound);
485        let models = self.models.clone();
486        if self.models.model_in_use().is_none() {
487            return cx.child(&models);
488        }
489        let labels = &self.labels;
490        let waker = cx.ui_waker();
491        let (speaking, sample) = (self.clone(), labels.sample.clone());
492        let (recommended, steadier) = (self.clone(), self.clone());
493        let fine = VStack::new((
494            self.slider(&labels.temperature, &self.temperature, (0.0, 1.5, 0.05)),
495            self.slider(&labels.top_k, &self.top_k, (0.0, 512.0, 1.0)),
496            self.slider(&labels.repetition_penalty, &self.repetition_penalty, (1.0, 1.5, 0.01)),
497            self.slider(&labels.repetition_window, &self.repetition_window, (0.0, 4096.0, 64.0)),
498            HStack::new((
499                Button::new(labels.recommended.clone())
500                    .size(ButtonSize::Small)
501                    .style(ButtonStyle::Secondary)
502                    .on_click(move || recommended.apply(&SpeechSettings::default())),
503                Button::new(labels.steadier.clone())
504                    .size(ButtonSize::Small)
505                    .style(ButtonStyle::Secondary)
506                    .on_click(move || steadier.apply(&SpeechSettings::steadier())),
507            ))
508            .spacing(theme.spacing.s2)
509            .wrap(),
510        ))
511        .spacing(theme.spacing.s3)
512        .align(Align::Stretch);
513        let listen = HStack::new((Button::new(labels.listen.clone())
514            .icon(IconName::Volume2)
515            .size(ButtonSize::Small)
516            .on_click(move || speaking.say(sample.clone(), waker.clone())),))
517        .align(Align::Center);
518        let read_aloud = Toggle::new(self.read_aloud.binding()).label(labels.read_aloud.clone());
519        if self.kokoro() {
520            // Kokoro has its own voices and nothing fine to set.
521            let controls = VStack::new((
522                SegmentedControl::new(self.speaker.binding(), labels.speakers.clone())
523                    .label(labels.speaker.clone()),
524                read_aloud,
525                listen,
526            ))
527            .spacing(theme.spacing.s3)
528            .align(Align::Stretch);
529            return cx
530                .child(&VStack::new((models, controls)).spacing(theme.spacing.s4).align(Align::Stretch));
531        }
532        let controls = VStack::new((
533            HStack::new((
534                Select::new(self.language.binding(), labels.languages.clone()).label(labels.language.clone()),
535                SegmentedControl::new(self.gender.binding(), [labels.female.clone(), labels.male.clone()])
536                    .label(labels.voice.clone()),
537            ))
538            .spacing(theme.spacing.s3)
539            .align(Align::Center)
540            .wrap(),
541            read_aloud,
542            listen,
543            Disclosure::new(labels.advanced.clone(), fine),
544        ))
545        .spacing(theme.spacing.s3)
546        .align(Align::Stretch);
547        cx.child(&VStack::new((models, controls)).spacing(theme.spacing.s4).align(Align::Stretch))
548    }
549}
550
551#[cfg(test)]
552mod tests {
553    use super::*;
554
555    #[test]
556    fn text_is_said_a_sentence_at_a_time() {
557        assert_eq!(
558            sentences("Você tem R$ 1.234,56. Quer ver mais?\nOk"),
559            ["Você tem R$ 1.234,56.", "Quer ver mais?", "Ok"]
560        );
561        assert!(sentences(" ... ").is_empty());
562    }
563
564    #[test]
565    fn the_settings_come_back_as_they_were_saved() {
566        let settings = SpeechSettings {
567            language: "Spanish".to_owned(),
568            voice: SpeechVoice::Male,
569            temperature: 0.5,
570            top_k: 64,
571            repetition_penalty: 1.2,
572            repetition_window: 128,
573            speaker: Some("pm_santa".to_owned()),
574            ..SpeechSettings::default()
575        };
576        assert_eq!(from_json(&to_json(&settings, false)), (settings, false));
577    }
578
579    #[test]
580    fn what_is_strange_stays_as_recommended() {
581        let strange = json!({"language": "Klingon", "temperature": -1, "repetition_penalty": 0.1, "speaker": "af_heart"});
582        assert_eq!(from_json(&strange), (SpeechSettings::default(), true));
583    }
584}