diff --git a/crates/buzz-voice/src/pocket.rs b/crates/buzz-voice/src/pocket.rs index 7e15c44cb..e5bc60864 100644 --- a/crates/buzz-voice/src/pocket.rs +++ b/crates/buzz-voice/src/pocket.rs @@ -210,8 +210,9 @@ impl PocketTts { /// /// 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. + /// engine advances before PCM is available or between internal model-safe + /// text chunks. Returning `false` interrupts synthesis before the next + /// model or decoder step. pub fn synth_chunk_streaming( &self, text: &str, @@ -346,6 +347,19 @@ mod tests { ); } + #[test] + fn equal_length_pre_decoder_callback_can_cancel() { + let mut callback = |samples: &[f32], progress: f32| { + assert!(samples.is_empty()); + assert_eq!(progress, 0.25); + false + }; + let mut samples = Vec::new(); + + assert!(!append_and_callback(&mut samples, &[], &mut callback, 0.25) + .expect("pre-decoder cancellation callback")); + } + #[test] fn callback_panic_is_reported_without_unwinding() { let mut callback = |_: &[f32], _: f32| -> bool { @@ -432,7 +446,7 @@ mod tests { #[test] #[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"] - fn production_streaming_callback_interrupts_after_first_decoder_block() { + fn production_streaming_callback_interrupts_before_pcm_is_available() { 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"); @@ -450,7 +464,9 @@ mod tests { "en", &style, 1, - |_, _| { + |samples, progress| { + assert!(samples.is_empty()); + assert!(progress <= 0.5); callback_at = Some(started.elapsed()); false }, @@ -465,4 +481,41 @@ mod tests { assert!(matches!(outcome, SynthesisOutcome::Interrupted)); } + + #[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"); + 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, + |samples, progress| { + if samples.is_empty() { + return true; + } + assert!(progress > 0.5); + callback_at = Some(started.elapsed()); + false + }, + ) + .expect("interrupt production streaming after decoded PCM"); + let callback_at = callback_at.expect("decoder must produce a callback"); + let cancellation_latency = started.elapsed().saturating_sub(callback_at); + eprintln!( + "decoded_callback_to_cancel_return_ms={:.1}", + cancellation_latency.as_secs_f64() * 1000.0 + ); + + assert!(matches!(outcome, SynthesisOutcome::Interrupted)); + } } diff --git a/crates/buzz-voice/src/pocket_april.rs b/crates/buzz-voice/src/pocket_april.rs index 11344072d..5a5a7feb3 100644 --- a/crates/buzz-voice/src/pocket_april.rs +++ b/crates/buzz-voice/src/pocket_april.rs @@ -285,7 +285,7 @@ impl AprilPocketTts { &mut self, prepared: &AprilPreparedPrompt, style: &VoiceStyle, - callback: F, + mut callback: F, ) -> Result where F: FnMut(&[f32], f32) -> bool, @@ -316,9 +316,18 @@ impl AprilPocketTts { let text_embeddings = self.text_embeddings(token_ids)?; self.run_flow_main_prefix(&text_embeddings, &mut flow_state)?; let max_frames = estimate_max_frames(token_count, self.bundle.frame_rate); - let latents = - self.generate_latents(max_frames, prepared.frames_after_eos, &mut flow_state)?; - self.decode_latents(&latents, callback) + let (latents, interrupted) = self.generate_latents( + max_frames, + prepared.frames_after_eos, + &mut flow_state, + &mut callback, + )?; + if interrupted { + return Ok(AprilSynthesisOutcome::Interrupted); + } + self.decode_latents(&latents, |samples, progress| { + callback(samples, 0.5 + progress * 0.5) + }) } fn prepared_token_count(&self, text: &str) -> Result { @@ -478,18 +487,29 @@ impl AprilPocketTts { replace_state_from_outputs(state, &mut outputs) } - fn generate_latents( + fn generate_latents( &mut self, max_frames: usize, frames_after_eos: usize, state: &mut [StateValue], - ) -> Result, String> { + callback: &mut F, + ) -> Result<(Vec, bool), String> + where + F: FnMut(&[f32], f32) -> bool, + { let mut current = vec![f32::NAN; self.bundle.latent_dim]; let mut latents = Vec::with_capacity(max_frames * self.bundle.latent_dim); let mut eos_step = None; let mut rng = rand::rng(); for step in 0..max_frames { + // Preserve the mobile callback's pre-PCM cancellation point while + // reserving the second half of progress for decoder-block output. + // The empty block becomes an equal-length cumulative callback in + // `PocketTts::synth_chunk_streaming`. + if !callback(&[], step as f32 / max_frames as f32 * 0.5) { + return Ok((latents, true)); + } let sequence = Tensor::from_array(( vec![1_i64, 1, self.bundle.latent_dim as i64], current.clone().into_boxed_slice(), @@ -575,7 +595,7 @@ impl AprilPocketTts { current.clone_from(&noise); latents.extend_from_slice(&noise); } - Ok(latents) + Ok((latents, false)) } fn decode_latents(