perf(desktop): stream Pocket TTS decoder chunks

crates/buzz-voice/src/lib.rs

@@ -5,7 +5,7 @@ pub mod pocket;
 
 pub use pocket::{
     april_model_info, load_text_to_speech, load_voice_style, PocketModelInfo, PocketTts,
-    VoiceStyle, DEFAULT_VOICE, SAMPLE_RATE, VOICE_FILE_EXT,
+    SynthesisOutcome, VoiceStyle, DEFAULT_VOICE, SAMPLE_RATE, VOICE_FILE_EXT,
 };
 
 /// One immutable artifact required by the April Pocket bundle.

crates/buzz-voice/src/pocket.rs

@@ -23,7 +23,7 @@ mod pocket_april;
 #[path = "pocket_models.rs"]
 mod pocket_models;
 
-use pocket_april::{prepare_april_prompt, AprilPocketTts};
+use pocket_april::{prepare_april_prompt, AprilPocketTts, AprilSynthesisOutcome};
 pub use pocket_models::{
     april_model_info, PocketModelArtifact, PocketModelInfo, APRIL_BUNDLE_ID, APRIL_MODEL_ID,
     APRIL_MODEL_REVISION,
@@ -69,6 +69,16 @@ pub struct PocketTts {
     inner: Mutex<AprilPocketTts>,
 }
 
+/// Result of a callback-driven Pocket synthesis request.
+#[derive(Debug, Clone, PartialEq)]
+pub enum SynthesisOutcome {
+    /// Synthesis finished and contains the same PCM exposed cumulatively to
+    /// the callback.
+    Complete(Vec<f32>),
+    /// The callback requested cancellation before synthesis completed.
+    Interrupted,
+}
+
 /// Load Buzz Desktop's pinned April INT8 model.
 pub fn load_text_to_speech(model_dir: &str) -> Result<PocketTts, String> {
     let dir = PathBuf::from(model_dir);
@@ -90,42 +100,104 @@ pub fn load_text_to_speech(model_dir: &str) -> Result<PocketTts, String> {
 impl PocketTts {
     /// Split text into synthesis units that satisfy the bundle's exact
     /// 50-token input limit.
+    ///
+    /// The first sentence remains its own unit when it fits. Oversized
+    /// sentences fall back to clause, word, and UTF-8 scalar boundaries, while
+    /// later sentences pack into the largest natural unit that fits.
+    ///
+    /// Chunks are contiguous substrings of the prepared model prompt and may
+    /// retain boundary whitespace. Concatenating them with `chunks.concat()`
+    /// reconstructs that prompt exactly, and each chunk's prepared token count
+    /// is at most 50.
     pub fn split_text_into_chunks(&self, text: &str) -> Result<Vec<String>, String> {
         let Some(prepared) = prepare_april_prompt(text) else {
             return Ok(Vec::new());
         };
         self.inner
             .lock()
             .map_err(|_| "Pocket TTS engine lock poisoned".to_string())?
-            .split_prompt(&prepared)
+            .split_playback_prompt(&prepared)
+    }
+
+    /// Group prepared text into playback units.
+    ///
+    /// A leading sentence remains separate so playback can begin promptly.
+    /// The remaining prompt stays in one playback unit; synthesis applies the
+    /// model's 50-token split internally without introducing playback pauses at
+    /// those model-only boundaries.
+    pub fn split_text_into_playback_chunks(&self, text: &str) -> Result<Vec<String>, String> {
+        let Some(prepared) = prepare_april_prompt(text) else {
+            return Ok(Vec::new());
+        };
+        Ok(self
+            .inner
+            .lock()
+            .map_err(|_| "Pocket TTS engine lock poisoned".to_string())?
+            .group_playback_prompt(&prepared))
     }
 
     /// Synthesize text with the supplied reference voice.
     ///
     /// Pocket detects language from text and this model uses one synthesis
     /// step, so `_lang` and `_steps` intentionally do not affect output.
     pub fn synth_chunk(
+        &self,
+        text: &str,
+        lang: &str,
+        style: &VoiceStyle,
+        steps: usize,
+    ) -> Result<Vec<f32>, String> {
+        match self.synth_chunk_streaming(text, lang, style, steps, |_, _| true)? {
+            SynthesisOutcome::Complete(samples) => Ok(samples),
+            SynthesisOutcome::Interrupted => Ok(Vec::new()),
+        }
+    }
+
+    /// Synthesize text while reporting cumulative PCM as decoder blocks finish.
+    ///
+    /// Callback sample buffers contain all PCM produced for this call so far.
+    /// Their lengths never decrease, but equal lengths are allowed while the
+    /// engine advances between internal model-safe text chunks. Returning
+    /// `false` interrupts synthesis before the next decoder block.
+    pub fn synth_chunk_streaming<F>(
         &self,
         text: &str,
         _lang: &str,
         style: &VoiceStyle,
         _steps: usize,
-    ) -> Result<Vec<f32>, String> {
+        mut callback: F,
+    ) -> Result<SynthesisOutcome, String>
+    where
+        F: FnMut(&[f32], f32) -> bool,
+    {
         let Some(prepared) = prepare_april_prompt(text) else {
-            return Ok(Vec::new());
+            return Ok(SynthesisOutcome::Complete(Vec::new()));
         };
         let mut engine = self
             .inner
             .lock()
             .map_err(|_| "Pocket TTS engine lock poisoned".to_string())?;
         let chunks = engine.split_prompt(&prepared)?;
+        let chunk_count = chunks.len();
         let mut samples = Vec::new();
-        for chunk in chunks {
+        for (chunk_index, chunk) in chunks.into_iter().enumerate() {
+            if chunk_index > 0 && !callback(&samples, chunk_index as f32 / chunk_count as f32) {
+                return Ok(SynthesisOutcome::Interrupted);
+            }
             let prepared = prepare_april_prompt(&chunk)
                 .ok_or_else(|| "Pocket TTS prompt chunk became empty".to_string())?;
-            samples.extend(engine.synth_chunk(&prepared, style)?);
+            let outcome = engine.synth_chunk_streaming(&prepared, style, |block, progress| {
+                samples.extend_from_slice(block);
+                callback(
+                    &samples,
+                    (chunk_index as f32 + progress) / chunk_count as f32,
+                )
+            })?;
+            if matches!(outcome, AprilSynthesisOutcome::Interrupted) {
+                return Ok(SynthesisOutcome::Interrupted);
+            }
         }
-        Ok(samples)
+        Ok(SynthesisOutcome::Complete(samples))
     }
 }
 
@@ -164,4 +236,89 @@ mod tests {
         assert!(samples.iter().all(|sample| sample.is_finite()));
         assert!(samples.iter().any(|sample| sample.abs() > 1.0e-6));
     }
+
+    #[test]
+    #[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"]
+    fn production_streaming_callbacks_are_cumulative_across_model_chunks() {
+        let dir = std::env::var("BUZZ_POCKET_TEST_MODEL_DIR")
+            .expect("set BUZZ_POCKET_TEST_MODEL_DIR to an April INT8 model directory");
+        let engine = load_text_to_speech(&dir).expect("load April INT8 engine");
+        let style = load_voice_style(&Path::new(&dir).join("reference_sample.wav"))
+            .expect("load reference voice");
+        engine
+            .synth_chunk("Warm up.", "en", &style, 1)
+            .expect("warm production engine");
+        let text = "And sometimes, when I am certain the reader is rested, I will engage him with a sentence of considerable length, a sentence that burns with energy and builds with all the impetus of a crescendo, the roll of the drums, the crash of the cymbals–sounds that say listen to this, it is important.";
+        let mut reconstructed = Vec::new();
+        let mut previous_len = 0;
+        let mut saw_equal_repeat = false;
+        let mut callback_count = 0;
+        let mut first_callback = None;
+        let started = std::time::Instant::now();
+
+        let outcome = engine
+            .synth_chunk_streaming(text, "en", &style, 1, |cumulative, _| {
+                callback_count += 1;
+                first_callback.get_or_insert_with(|| started.elapsed());
+                assert!(cumulative.len() >= previous_len);
+                saw_equal_repeat |= cumulative.len() == previous_len;
+                reconstructed.extend_from_slice(&cumulative[previous_len..]);
+                previous_len = cumulative.len();
+                true
+            })
+            .expect("stream through the production API");
+        let SynthesisOutcome::Complete(samples) = outcome else {
+            panic!("uninterrupted synthesis must complete");
+        };
+        let total = started.elapsed();
+        let first_callback = first_callback.expect("decoder must produce a callback");
+        let audio_duration =
+            std::time::Duration::from_secs_f64(samples.len() as f64 / SAMPLE_RATE as f64);
+        eprintln!(
+            "first_callback_ms={:.1} total_ms={:.1} audio_seconds={:.3} rtf={:.3} callbacks={callback_count}",
+            first_callback.as_secs_f64() * 1000.0,
+            total.as_secs_f64() * 1000.0,
+            audio_duration.as_secs_f64(),
+            total.as_secs_f64() / audio_duration.as_secs_f64(),
+        );
+
+        assert!(saw_equal_repeat);
+        assert_eq!(reconstructed, samples);
+    }
+
+    #[test]
+    #[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"]
+    fn production_streaming_callback_interrupts_after_first_decoder_block() {
+        let dir = std::env::var("BUZZ_POCKET_TEST_MODEL_DIR")
+            .expect("set BUZZ_POCKET_TEST_MODEL_DIR to an April INT8 model directory");
+        let engine = load_text_to_speech(&dir).expect("load April INT8 engine");
+        let style = load_voice_style(&Path::new(&dir).join("reference_sample.wav"))
+            .expect("load reference voice");
+        engine
+            .synth_chunk("Warm up.", "en", &style, 1)
+            .expect("warm production engine");
+        let mut callback_at = None;
+        let started = std::time::Instant::now();
+
+        let outcome = engine
+            .synth_chunk_streaming(
+                "This sentence is long enough to require more than one decoder block.",
+                "en",
+                &style,
+                1,
+                |_, _| {
+                    callback_at = Some(started.elapsed());
+                    false
+                },
+            )
+            .expect("interrupt production streaming");
+        let callback_at = callback_at.expect("decoder must produce a callback");
+        let cancellation_latency = started.elapsed().saturating_sub(callback_at);
+        eprintln!(
+            "callback_to_cancel_return_ms={:.1}",
+            cancellation_latency.as_secs_f64() * 1000.0
+        );
+
+        assert!(matches!(outcome, SynthesisOutcome::Interrupted));
+    }
 }