forge_tts/
engine.rs

1//! The in-process Kokoro-82M engine.
2
3use std::sync::{Mutex, mpsc};
4
5use tokio::sync::mpsc as tokio_mpsc;
6
7use crate::model::ModelFiles;
8use crate::{Chunk, Engine, Error, SpeakRequest};
9
10/// The wire's sample rate: 24 kHz mono 16-bit. Kokoro renders at this rate
11/// natively; anything else is resampled once, here.
12pub const SAMPLE_RATE: u32 = 24_000;
13
14/// The voice an utterance renders in when the request names none.
15pub const DEFAULT_VOICE: &str = "af_heart";
16
17/// A loaded Kokoro-82M pipeline. One render thread per engine, so a second
18/// utterance queues behind the first; the thread owns the ONNX session and
19/// the runtime the engine's async API needs.
20///
21/// Kokoro renders a whole clip, so one `speak` answers one chunk.
22pub struct KokoroTts {
23    jobs: Option<tokio_mpsc::UnboundedSender<Job>>,
24    worker: Mutex<Option<std::thread::JoinHandle<()>>>,
25}
26
27struct Job {
28    request: SpeakRequest,
29    reply: tokio_mpsc::Sender<Result<Chunk, Error>>,
30}
31
32/// The installed voices a request is checked against.
33///
34/// The engine's own lookup falls back to a lone loaded voice when the name
35/// is unknown, which is the silent default the wire must not have.
36pub fn select_voice(requested: Option<&str>, installed: &[String]) -> Result<String, Error> {
37    let id = requested.unwrap_or(DEFAULT_VOICE);
38    if installed.iter().any(|voice| voice == id) {
39        Ok(id.to_owned())
40    } else {
41        Err(Error::VoiceMissing(id.to_owned()))
42    }
43}
44
45impl KokoroTts {
46    /// Load the model and every voice in `files.voices`, on a thread of their
47    /// own. A caller may warm it ahead of the first `say`, and a failed warm
48    /// is retried by the next one. Every declared digest is checked before
49    /// the load.
50    pub fn load(files: &ModelFiles) -> Result<Self, Error> {
51        files.verify()?;
52        let installed = files.installed_voices()?;
53        let (jobs, mut queue) = tokio_mpsc::unbounded_channel::<Job>();
54        let (ready, loaded) = mpsc::channel();
55        let model = files.model.clone();
56        let voices = files.voices.clone();
57        let worker_installed = installed.clone();
58        let worker = std::thread::Builder::new()
59            .name("forge-tts".to_owned())
60            .spawn(move || {
61                let runtime =
62                    match tokio::runtime::Builder::new_current_thread().enable_all().build() {
63                        Ok(runtime) => runtime,
64                        Err(error) => {
65                            let _ = ready.send(Err(Error::Load(error.to_string())));
66                            return;
67                        }
68                    };
69                let engine = match runtime.block_on(kokoro_en::KokoroTts::new(&model, &voices)) {
70                    Ok(engine) => engine,
71                    Err(error) => {
72                        let _ = ready.send(Err(Error::Load(error.to_string())));
73                        return;
74                    }
75                };
76                let mut renderer = Kokoro { runtime, engine };
77                if ready.send(Ok(())).is_err() {
78                    return;
79                }
80                while let Some(job) = queue.blocking_recv() {
81                    // A send failure is the host stopping playback, so this
82                    // render's result has nowhere to go.
83                    let _ = job.reply.blocking_send(chunk_for(
84                        &mut renderer,
85                        &worker_installed,
86                        &job.request,
87                    ));
88                }
89            })
90            .map_err(|error| Error::Load(error.to_string()))?;
91
92        match loaded.recv() {
93            Ok(Ok(())) => Ok(Self { jobs: Some(jobs), worker: Mutex::new(Some(worker)) }),
94            Ok(Err(error)) => Err(error),
95            Err(_) => Err(Error::Load("the render thread stopped while loading".to_owned())),
96        }
97    }
98}
99
100impl Drop for KokoroTts {
101    fn drop(&mut self) {
102        // Close the queue first, so the thread ends its loop, drops the
103        // session and only then lets the join return.
104        drop(self.jobs.take());
105        let worker = self.worker.get_mut().ok().and_then(Option::take);
106        if let Some(worker) = worker {
107            let _ = worker.join();
108        }
109    }
110}
111
112impl Engine for KokoroTts {
113    fn speak(&self, request: SpeakRequest) -> tokio_mpsc::Receiver<Result<Chunk, Error>> {
114        let Some(jobs) = self.jobs.as_ref() else {
115            return refused();
116        };
117        let (reply, chunks) = tokio_mpsc::channel(2);
118        if jobs.send(Job { request, reply }).is_err() {
119            return refused();
120        }
121        chunks
122    }
123}
124
125/// A channel carrying one failure, for a render thread that is gone.
126fn refused() -> tokio_mpsc::Receiver<Result<Chunk, Error>> {
127    let (reply, chunks) = tokio_mpsc::channel(1);
128    let _ = reply.try_send(Err(Error::Render("the render thread is no longer running".to_owned())));
129    chunks
130}
131
132/// What the engine's inference looks like to the wire conversion, so a test
133/// can drive the conversion with no model load.
134trait Renderer: Send {
135    fn native_rate(&self) -> u32;
136    fn render(&mut self, text: &str, voice: &str) -> Result<Vec<f32>, Error>;
137}
138
139/// One render: the voice the request names, the model's samples, and the
140/// wire's own shape.
141fn chunk_for(
142    renderer: &mut dyn Renderer,
143    installed: &[String],
144    request: &SpeakRequest,
145) -> Result<Chunk, Error> {
146    let voice = select_voice(request.voice.as_deref(), installed)?;
147    let samples = renderer.render(&request.text, &voice)?;
148    let samples = resample(&samples, renderer.native_rate(), SAMPLE_RATE);
149    Ok(Chunk { seq: 0, sample_rate: SAMPLE_RATE, pcm: pcm16(&samples) })
150}
151
152/// Linear resampling to the wire rate; identity when the rates match.
153fn resample(samples: &[f32], from: u32, to: u32) -> Vec<f32> {
154    if from == to || from == 0 || samples.is_empty() {
155        return samples.to_vec();
156    }
157    let ratio = f64::from(from) / f64::from(to);
158    let length = ((samples.len() as f64) * f64::from(to) / f64::from(from)).round() as usize;
159    let last = samples.len() - 1;
160    let mut out = Vec::with_capacity(length);
161    for index in 0..length {
162        let position = index as f64 * ratio;
163        let left = position.floor() as usize;
164        let fraction = (position - left as f64) as f32;
165        let a = samples[left.min(last)];
166        let b = samples[(left + 1).min(last)];
167        out.push(a + (b - a) * fraction);
168    }
169    out
170}
171
172/// The wire's sample: clamp to [-1, 1], then the signed 16-bit range.
173fn pcm16(samples: &[f32]) -> Vec<i16> {
174    samples
175        .iter()
176        .map(|sample| {
177            let clamped = sample.clamp(-1.0, 1.0);
178            (clamped * f32::from(i16::MAX)) as i16
179        })
180        .collect()
181}
182
183/// The loaded engine, as the wire conversion sees it.
184struct Kokoro {
185    runtime: tokio::runtime::Runtime,
186    engine: kokoro_en::KokoroTts,
187}
188
189impl Renderer for Kokoro {
190    fn native_rate(&self) -> u32 {
191        SAMPLE_RATE
192    }
193
194    fn render(&mut self, text: &str, voice: &str) -> Result<Vec<f32>, Error> {
195        self.runtime
196            .block_on(self.engine.synth(text, voice))
197            .map(|(audio, _)| audio)
198            .map_err(|error| Error::Render(error.to_string()))
199    }
200}
201
202#[cfg(test)]
203mod tests {
204    use super::{DEFAULT_VOICE, Renderer, SAMPLE_RATE, chunk_for, select_voice};
205    use crate::{Engine as _, Error, SpeakRequest};
206
207    /// A renderer that records the voice it was asked for and answers with
208    /// `samples` at `rate`.
209    struct Stub {
210        rate: u32,
211        samples: Vec<f32>,
212        spoken: Vec<String>,
213    }
214
215    impl Stub {
216        fn new(rate: u32, samples: Vec<f32>) -> Self {
217            Self { rate, samples, spoken: Vec::new() }
218        }
219    }
220
221    impl Renderer for Stub {
222        fn native_rate(&self) -> u32 {
223            self.rate
224        }
225
226        fn render(&mut self, _text: &str, voice: &str) -> Result<Vec<f32>, Error> {
227            self.spoken.push(voice.to_owned());
228            Ok(self.samples.clone())
229        }
230    }
231
232    fn installed() -> Vec<String> {
233        vec!["af_heart".to_owned(), "bf_emma".to_owned()]
234    }
235
236    #[test]
237    fn a_chosen_voice_is_the_one_the_render_uses() {
238        let mut stub = Stub::new(SAMPLE_RATE, vec![0.0; 8]);
239        let request = SpeakRequest { text: "hello".to_owned(), voice: Some("bf_emma".to_owned()) };
240        chunk_for(&mut stub, &installed(), &request).expect("the render lands");
241        assert_eq!(stub.spoken, vec!["bf_emma".to_owned()]);
242    }
243
244    #[test]
245    fn no_chosen_voice_renders_in_the_default() {
246        let mut stub = Stub::new(SAMPLE_RATE, vec![0.0; 8]);
247        let request = SpeakRequest { text: "hello".to_owned(), voice: None };
248        chunk_for(&mut stub, &installed(), &request).expect("the render lands");
249        assert_eq!(stub.spoken, vec![DEFAULT_VOICE.to_owned()]);
250    }
251
252    #[test]
253    fn an_unknown_voice_is_refused_by_name() {
254        let error = select_voice(Some("zf_nobody"), &installed()).expect_err("refused");
255        assert_eq!(error, Error::VoiceMissing("zf_nobody".to_owned()));
256        assert!(error.to_string().contains("zf_nobody"), "got: {error}");
257    }
258
259    #[test]
260    fn a_default_voice_that_is_not_installed_is_refused_rather_than_replaced() {
261        let error = select_voice(None, &["bf_emma".to_owned()]).expect_err("refused");
262        assert_eq!(error, Error::VoiceMissing(DEFAULT_VOICE.to_owned()));
263    }
264
265    /// The model's own rate never reaches the wire: a 16 kHz render arrives
266    /// at 24 kHz, mono, as signed 16-bit samples.
267    #[test]
268    fn the_wire_chunk_is_24k_mono_sixteen_bit_whatever_the_model_rate() {
269        let mut stub = Stub::new(16_000, vec![0.5; 16]);
270        let request = SpeakRequest { text: "hello".to_owned(), voice: None };
271        let chunk = chunk_for(&mut stub, &installed(), &request).expect("the render lands");
272        assert_eq!(chunk.sample_rate, SAMPLE_RATE);
273        assert_eq!(chunk.pcm.len(), 24, "1 ms at 16 kHz is 1 ms at 24 kHz");
274        assert!(chunk.pcm.iter().all(|sample| *sample == 16_383), "0.5 scales to half the range");
275    }
276
277    #[test]
278    fn a_sample_past_full_scale_is_clamped_not_wrapped() {
279        let mut stub = Stub::new(SAMPLE_RATE, vec![2.0, -2.0]);
280        let request = SpeakRequest { text: "loud".to_owned(), voice: None };
281        let chunk = chunk_for(&mut stub, &installed(), &request).expect("the render lands");
282        assert_eq!(chunk.pcm, vec![i16::MAX, -i16::MAX]);
283    }
284
285    /// Loads the real weights, so it runs only where they are on disk:
286    /// `FORGE_TTS_MODEL_DIR` must hold `model_quantized.onnx` and a
287    /// `voices/` directory with at least the default voice.
288    #[tokio::test]
289    #[ignore = "loads the real model"]
290    async fn the_engine_speaks_a_short_sentence() {
291        let dir = std::path::PathBuf::from(
292            std::env::var("FORGE_TTS_MODEL_DIR").expect("FORGE_TTS_MODEL_DIR"),
293        );
294        let files = crate::ModelFiles::new(dir.join("model_quantized.onnx"), dir.join("voices"));
295        let engine = crate::KokoroTts::load(&files).expect("loads");
296        let mut chunks =
297            engine.speak(SpeakRequest { text: "Ready when you are.".to_owned(), voice: None });
298        let first = chunks.recv().await.expect("a chunk").expect("a render");
299        assert_eq!(first.sample_rate, SAMPLE_RATE);
300        assert!(!first.pcm.is_empty());
301    }
302}
forge_tts/engine.rs

forge_tts/
engine.rs
