From 450d36b682f1c7fd4f8ec2ca78012a907a87a9bf Mon Sep 17 00:00:00 2001 From: John Tennant Date: Wed, 19 Aug 2026 16:55:45 -0400 Subject: [PATCH 01/10] feat(voice): stream Pocket audio during synthesis --- src-tauri/crates/berd-voice/src/lib.rs | 4 +- src-tauri/crates/berd-voice/src/pocket.rs | 384 ++++---- .../crates/berd-voice/src/pocket_april.rs | 846 ++++++++++++++++-- src-tauri/src/commands/pocket_voice.rs | 117 +-- 4 files changed, 989 insertions(+), 362 deletions(-) diff --git a/src-tauri/crates/berd-voice/src/lib.rs b/src-tauri/crates/berd-voice/src/lib.rs index 651f0fd59..db34aec2c 100644 --- a/src-tauri/crates/berd-voice/src/lib.rs +++ b/src-tauri/crates/berd-voice/src/lib.rs @@ -3,5 +3,7 @@ mod pocket; pub use pocket::{ - load_text_to_speech, load_voice_style, PocketTts, SynthesisOutcome, VoiceStyle, SAMPLE_RATE, + april_model_info, load_text_to_speech, load_voice_style, PocketModelArtifact, PocketModelInfo, + PocketTts, VoiceStyle, APRIL_BUNDLE_ID, APRIL_MODEL_ID, APRIL_MODEL_REVISION, DEFAULT_VOICE, + SAMPLE_RATE, VOICE_FILE_EXT, }; diff --git a/src-tauri/crates/berd-voice/src/pocket.rs b/src-tauri/crates/berd-voice/src/pocket.rs index c7f8b8207..bbf456823 100644 --- a/src-tauri/crates/berd-voice/src/pocket.rs +++ b/src-tauri/crates/berd-voice/src/pocket.rs @@ -14,51 +14,35 @@ //! Berd's Pocket model installer writes the complete attribution beside the //! cached model files. -use std::cell::RefCell; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::Mutex; use sherpa_onnx::Wave; #[path = "pocket_april.rs"] mod pocket_april; -use pocket_april::{prepare_april_prompt, AprilPocketTts, AprilSynthesisOutcome}; +#[path = "pocket_models.rs"] +mod pocket_models; + +use pocket_april::{prepare_april_prompt, AprilPocketTts}; +pub use pocket_models::{ + april_model_info, PocketModelArtifact, PocketModelInfo, APRIL_BUNDLE_ID, APRIL_MODEL_ID, + APRIL_MODEL_REVISION, +}; /// Pocket TTS emits 24 kHz mono PCM. pub const SAMPLE_RATE: u32 = 24_000; const TTS_NUM_THREADS: usize = 1; -thread_local! { - static ACTIVE_SYNTHESIS_ENGINES: RefCell> = const { RefCell::new(Vec::new()) }; -} - -struct SynthesisCallGuard { - engine_id: usize, -} - -impl SynthesisCallGuard { - fn enter(engine_id: usize) -> Result { - ACTIVE_SYNTHESIS_ENGINES.with(|active| { - let mut active = active.borrow_mut(); - if active.contains(&engine_id) { - return Err("Pocket TTS callback re-entered the active engine".to_string()); - } - active.push(engine_id); - Ok(Self { engine_id }) - }) - } -} - -impl Drop for SynthesisCallGuard { - fn drop(&mut self) { - ACTIVE_SYNTHESIS_ENGINES.with(|active| { - let mut active = active.borrow_mut(); - if let Some(index) = active.iter().rposition(|engine| *engine == self.engine_id) { - active.remove(index); - } - }); - } +/// EXPERIMENTAL (latency): override ONNX intra-op threads for the Pocket +/// sessions via `BERD_TTS_THREADS`. Default preserves production's 1. +fn tts_num_threads() -> usize { + std::env::var("BERD_TTS_THREADS") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&n| n >= 1) + .unwrap_or(TTS_NUM_THREADS) } /// Loaded reference voice samples and their original sample rate. @@ -90,113 +74,101 @@ pub struct PocketTts { inner: Mutex, } -/// 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), - /// The callback requested cancellation before synthesis completed. - Interrupted, -} - /// Load Berd's pinned April INT8 model. pub fn load_text_to_speech(model_dir: &str) -> Result { let dir = Path::new(model_dir); Ok(PocketTts { - inner: Mutex::new(AprilPocketTts::load(dir, TTS_NUM_THREADS)?), + inner: Mutex::new(AprilPocketTts::load(&dir, tts_num_threads())?), }) } impl PocketTts { - /// Synthesize text while reporting cumulative PCM as decoder blocks finish. + /// Split text into model-safe synthesis units that satisfy the bundle's + /// exact 50-token input limit, packing sentences whenever they fit. + pub fn split_text_into_chunks(&self, text: &str) -> Result, 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 text into ordered playback units, keeping the first sentence + /// separate so it reaches synthesis before the remainder is packed. /// - /// 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 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( + /// Units are contiguous substrings of the prepared model prompt and may + /// retain boundary whitespace. Concatenating them with `chunks.concat()` + /// reconstructs that prompt exactly, and each unit's prepared token count + /// is at most 50. + pub fn split_text_for_playback(&self, text: &str) -> Result, 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_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, - mut callback: F, - ) -> Result - where - F: FnMut(&[f32], f32) -> bool, - { - let _call_guard = SynthesisCallGuard::enter(self as *const Self as usize)?; + _steps: usize, + ) -> Result, String> { let Some(prepared) = prepare_april_prompt(text) else { - return Ok(SynthesisOutcome::Complete(Vec::new())); + return Ok(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_index, chunk) in chunks.into_iter().enumerate() { - if chunk_index > 0 - && !callback_allows_progress( - &mut callback, - &samples, - chunk_index as f32 / chunk_count as f32, - )? - { - return Ok(SynthesisOutcome::Interrupted); - } + for chunk in chunks { let prepared = prepare_april_prompt(&chunk) .ok_or_else(|| "Pocket TTS prompt chunk became empty".to_string())?; - let mut callback_error = None; - let outcome = engine.synth_chunk_streaming(&prepared, style, |block, progress| { - match append_and_callback( - &mut samples, - block, - &mut callback, - (chunk_index as f32 + progress) / chunk_count as f32, - ) { - Ok(allowed) => allowed, - Err(error) => { - callback_error = Some(error); - false - } - } - })?; - if let Some(error) = callback_error { - return Err(error); - } - if matches!(outcome, AprilSynthesisOutcome::Interrupted) { - return Ok(SynthesisOutcome::Interrupted); - } + samples.extend(engine.synth_chunk(&prepared, style)?); } - Ok(SynthesisOutcome::Complete(samples)) + Ok(samples) } -} - -fn append_and_callback( - samples: &mut Vec, - block: &[f32], - callback: &mut F, - progress: f32, -) -> Result -where - F: FnMut(&[f32], f32) -> bool, -{ - samples.extend_from_slice(block); - callback_allows_progress(callback, samples, progress) -} -fn callback_allows_progress( - callback: &mut F, - samples: &[f32], - progress: f32, -) -> Result -where - F: FnMut(&[f32], f32) -> bool, -{ - std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| callback(samples, progress))) - .map_err(|_| "Pocket TTS synthesis callback panicked".to_string()) + /// EXPERIMENTAL (latency): streaming synthesis. Invokes `on_audio` with + /// PCM deltas as soon as roughly `emit_frames` Flow LM frames (80 ms of + /// audio each) have been generated and decoded. Concatenated deltas equal + /// one `synth_chunk` result. The callback runs on the caller thread and + /// returns `false` to cancel; the function then returns Ok(false). + pub fn synth_chunk_streaming( + &self, + text: &str, + style: &VoiceStyle, + emit_frames: usize, + on_audio: &mut dyn FnMut(Vec) -> bool, + ) -> Result { + let Some(prepared) = prepare_april_prompt(text) else { + return Ok(true); + }; + let mut engine = self + .inner + .lock() + .map_err(|_| "Pocket TTS engine lock poisoned".to_string())?; + let chunks = engine.split_prompt(&prepared)?; + for chunk in chunks { + let prepared = prepare_april_prompt(&chunk) + .ok_or_else(|| "Pocket TTS prompt chunk became empty".to_string())?; + if !engine.synth_chunk_streaming(&prepared, style, emit_frames, on_audio)? { + return Ok(false); + } + } + Ok(true) + } } #[cfg(test)] @@ -204,109 +176,119 @@ mod tests { use super::*; #[test] - fn cumulative_callback_allows_growth_equal_repeats_and_cancellation() { - let mut observed = Vec::new(); - let mut callback = |samples: &[f32], progress: f32| { - observed.push((samples.to_vec(), progress)); - progress < 0.75 - }; - let mut samples = Vec::new(); - - assert!( - append_and_callback(&mut samples, &[1.0, 2.0], &mut callback, 0.25) - .expect("first callback") - ); - assert!(append_and_callback(&mut samples, &[], &mut callback, 0.5) - .expect("equal-length callback")); - assert!( - !append_and_callback(&mut samples, &[3.0], &mut callback, 0.75) - .expect("cancelling callback") - ); - - assert_eq!( - observed, - vec![ - (vec![1.0, 2.0], 0.25), - (vec![1.0, 2.0], 0.5), - (vec![1.0, 2.0, 3.0], 0.75), - ] - ); + fn desktop_model_is_april_int8_only() { + let info = april_model_info(); + assert_eq!(info.max_token_per_chunk, 50); + assert_eq!(info.sample_rate, SAMPLE_RATE); + assert!(info + .artifacts + .iter() + .any(|artifact| artifact.filename == "flow_lm_main_int8.onnx")); + assert!(!info + .artifacts + .iter() + .any(|artifact| artifact.filename == "flow_lm_main.onnx")); } - #[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")); + /// Which splitter each production function delegates to, across the whole + /// file rather than one hand-picked window. + /// + /// A wrong delegation can reinstate either shipped defect in one token: + /// removing first-sentence priority from playback, or re-isolating sentence + /// one inside units that already fit. Asserting the whole map means a new + /// delegation must be declared here to compile green. + fn splitter_delegations(source: &str) -> Vec<(String, Vec)> { + let production = source + .split_once("\n#[cfg(test)]") + .map_or(source, |(production, _)| production); + // Scan code only. Prose cannot call a splitter, but it can contain + // ` fn `, which would end a body early and hide a call after it, and it + // can name a splitter, which would report a call the code never makes. + let production: String = production + .lines() + .map(|line| line.split_once("//").map_or(line, |(code, _)| code)) + .collect::>() + .join("\n"); + let mut out = Vec::new(); + let mut rest = production.as_str(); + while let Some((_, after)) = rest.split_once(" fn ") { + let (name, body) = after + .split_once('(') + .expect("a function signature has an argument list"); + // End at this function's own closing brace, not at the next ` fn `: + // a body provably stops where its braces balance, so no later + // function's calls are attributed here and none of this one's are + // dropped. + let inner = body.split_once('{').map_or("", |(_, inner)| inner); + let mut depth = 1usize; + let body = inner + .char_indices() + .find(|&(_, ch)| { + depth = match ch { + '{' => depth + 1, + '}' => depth - 1, + _ => depth, + }; + depth == 0 + }) + .map_or(inner, |(end, _)| &inner[..end]); + let mut calls = Vec::new(); + // Check the isolating spelling first: ".split_prompt(" is a + // substring of neither, but a naive contains() on the shorter name + // would also match the longer one. + for _ in 0..body.matches(".split_playback_prompt(").count() { + calls.push("split_playback_prompt".to_string()); + } + let plain = body.matches(".split_prompt(").count(); + for _ in 0..plain { + calls.push("split_prompt".to_string()); + } + if !calls.is_empty() { + out.push((name.trim().to_string(), calls)); + } + rest = after; + } + out } #[test] - fn callback_panic_is_reported_without_unwinding() { - let mut callback = |_: &[f32], _: f32| -> bool { - panic!("callback failure"); - }; + fn every_production_splitter_delegation_is_declared() { + let source = include_str!("pocket.rs"); + let actual = splitter_delegations(source); + let expected: Vec<(String, Vec)> = vec![ + // Model units: pack sentences, never isolate. + ("split_text_into_chunks".into(), vec!["split_prompt".into()]), + // Playback units: isolate sentence one for time-to-first-audio. + ( + "split_text_for_playback".into(), + vec!["split_playback_prompt".into()], + ), + // Synthesis receives an already-packed unit: re-isolating here + // re-adds the per-sentence seam this PR removes. + ("synth_chunk".into(), vec!["split_prompt".into()]), + ("synth_chunk_streaming".into(), vec!["split_prompt".into()]), + ]; assert_eq!( - callback_allows_progress(&mut callback, &[], 0.0).unwrap_err(), - "Pocket TTS synthesis callback panicked" + actual, expected, + "a production function changed which splitter it calls (or a new \ + one appeared); isolating outside split_text_for_playback delays \ + first audio, packing inside it removes the guarantee" ); } - #[test] - fn active_engine_reentry_is_rejected() { - let _guard = SynthesisCallGuard::enter(42).expect("first call"); - assert!(SynthesisCallGuard::enter(42).is_err()); - } - #[test] #[ignore = "requires BERD_POCKET_TEST_MODEL_DIR"] - fn production_streaming_callbacks_are_cumulative_across_model_chunks() { + fn production_api_emits_non_silent_april_int8_pcm() { let dir = std::env::var("BERD_POCKET_TEST_MODEL_DIR") .expect("set BERD_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 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, &style, |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(), - ); + let samples = engine + .synth_chunk("Bright birds begin beside the bay.", "en", &style, 1) + .expect("synthesize through the production API"); - assert!(saw_equal_repeat); - assert_eq!(reconstructed, samples); + assert!(!samples.is_empty()); assert!(samples.iter().all(|sample| sample.is_finite())); assert!(samples.iter().any(|sample| sample.abs() > 1.0e-6)); } diff --git a/src-tauri/crates/berd-voice/src/pocket_april.rs b/src-tauri/crates/berd-voice/src/pocket_april.rs index d2b259a16..1de647a07 100644 --- a/src-tauri/crates/berd-voice/src/pocket_april.rs +++ b/src-tauri/crates/berd-voice/src/pocket_april.rs @@ -96,13 +96,129 @@ struct StateValue { value: DynValue, } -struct CachedVoice { - samples_ptr: usize, +/// Stable identity for a reference voice: a content hash of the sample +/// buffer plus its length and rate. Buffer addresses are NOT part of the +/// key — voice switching clones and drops sample buffers, so the allocator +/// can hand a different voice the same address, and an address-based key +/// would then restore the previous voice's cached state. +#[derive(PartialEq, Eq, Clone, Copy, Debug)] +struct VoiceKey { + content_hash: u64, samples_len: usize, sample_rate: i32, +} + +fn voice_key(style: &VoiceStyle) -> VoiceKey { + use std::hash::Hasher; + let mut hasher = std::hash::DefaultHasher::new(); + for sample in &style.samples { + hasher.write_u32(sample.to_bits()); + } + VoiceKey { + content_hash: hasher.finish(), + samples_len: style.samples.len(), + sample_rate: style.sample_rate, + } +} + +struct CachedVoice { + key: VoiceKey, embeddings: Vec, } +/// EXPERIMENTAL (latency): a dtype-tagged copy of one recurrent state tensor, +/// used to snapshot the Flow LM state right after voice conditioning so +/// subsequent chunks skip the ~160 ms `condition_voice` pass entirely. +enum SnapshotTensor { + F32(Vec, Vec), + I64(Vec, Vec), + Bool(Vec, Vec), +} + +struct CachedConditioning { + key: VoiceKey, + state: Vec<(StateSpec, SnapshotTensor)>, +} + +fn snapshot_state(state: &[StateValue]) -> Result, String> { + state + .iter() + .map(|value| { + let tensor = match value.spec.dtype { + StateDtype::Float32 => { + let (shape, data) = value + .value + .try_extract_tensor::() + .map_err(ort_error("snapshot f32 state"))?; + SnapshotTensor::F32(shape.to_vec(), data.to_vec()) + } + StateDtype::Int64 => { + let (shape, data) = value + .value + .try_extract_tensor::() + .map_err(ort_error("snapshot i64 state"))?; + SnapshotTensor::I64(shape.to_vec(), data.to_vec()) + } + StateDtype::Bool => { + let (shape, data) = value + .value + .try_extract_tensor::() + .map_err(ort_error("snapshot bool state"))?; + SnapshotTensor::Bool(shape.to_vec(), data.to_vec()) + } + }; + Ok((value.spec.clone(), tensor)) + }) + .collect() +} + +fn restore_state(snapshot: &[(StateSpec, SnapshotTensor)]) -> Result, String> { + snapshot + .iter() + .map(|(spec, tensor)| { + let value = match tensor { + SnapshotTensor::F32(shape, data) => { + if data.is_empty() { + Tensor::::new(&ort::memory::Allocator::default(), shape.clone()) + .map_err(ort_error("restore empty f32 state"))? + .into_dyn() + } else { + Tensor::from_array((shape.clone(), data.clone().into_boxed_slice())) + .map_err(ort_error("restore f32 state"))? + .into_dyn() + } + } + SnapshotTensor::I64(shape, data) => { + if data.is_empty() { + Tensor::::new(&ort::memory::Allocator::default(), shape.clone()) + .map_err(ort_error("restore empty i64 state"))? + .into_dyn() + } else { + Tensor::from_array((shape.clone(), data.clone().into_boxed_slice())) + .map_err(ort_error("restore i64 state"))? + .into_dyn() + } + } + SnapshotTensor::Bool(shape, data) => { + if data.is_empty() { + Tensor::::new(&ort::memory::Allocator::default(), shape.clone()) + .map_err(ort_error("restore empty bool state"))? + .into_dyn() + } else { + Tensor::from_array((shape.clone(), data.clone().into_boxed_slice())) + .map_err(ort_error("restore bool state"))? + .into_dyn() + } + } + }; + Ok(StateValue { + spec: spec.clone(), + value, + }) + }) + .collect() +} + pub(crate) struct AprilPocketTts { bundle: Bundle, tokenizer: Tokenizer, @@ -113,11 +229,10 @@ pub(crate) struct AprilPocketTts { flow: Session, mimi_decoder: Session, cached_voice: Option, -} - -pub(crate) enum AprilSynthesisOutcome { - Complete, - Interrupted, + /// EXPERIMENTAL (latency): post-`condition_voice` Flow LM state, cached + /// per reference voice. Restoring it replaces the ~160 ms conditioning + /// pass on every chunk after the first for a given voice. + cached_conditioning: Option, } #[derive(Debug, Clone, PartialEq)] @@ -251,6 +366,7 @@ impl AprilPocketTts { tokenizer, bos_embedding, cached_voice: None, + cached_conditioning: None, }) } @@ -261,25 +377,110 @@ impl AprilPocketTts { if self.prepared_token_count(&prepared.text)? <= self.bundle.max_token_per_chunk { return Ok(vec![prepared.text.clone()]); } - split_at_natural_boundaries( + split_model_at_natural_boundaries(&prepared.text, self.bundle.max_token_per_chunk, |text| { + self.prepared_token_count(text) + }) + } + + pub(crate) fn split_playback_prompt( + &self, + prepared: &AprilPreparedPrompt, + ) -> Result, String> { + split_playback_at_natural_boundaries( &prepared.text, self.bundle.max_token_per_chunk, - false, |text| self.prepared_token_count(text), ) } - pub(crate) fn synth_chunk_streaming( + pub(crate) fn synth_chunk( &mut self, prepared: &AprilPreparedPrompt, style: &VoiceStyle, - mut callback: F, - ) -> Result - where - F: FnMut(&[f32], f32) -> bool, - { + ) -> Result, String> { + // EXPERIMENTAL (latency bench): phase timing, enabled by BERD_TTS_PHASE_LOG=1. + let phase_log = std::env::var("BERD_TTS_PHASE_LOG").is_ok_and(|v| v == "1"); + let t0 = std::time::Instant::now(); + let mut flow_state = self.conditioned_flow_state(style)?; + let t_condition = t0.elapsed(); + let token_ids = self + .tokenizer + .encode(prepared.text.as_str(), false) + .map_err(|err| format!("tokenize Pocket TTS prompt: {err}"))? + .get_ids() + .iter() + .copied() + .map(i64::from) + .collect::>(); + if token_ids.is_empty() { + return Ok(Vec::new()); + } + if token_ids.len() > self.bundle.max_token_per_chunk { + return Err(format!( + "Pocket TTS prompt has {} tokens; split_text_into_chunks maximum is {}", + token_ids.len(), + self.bundle.max_token_per_chunk + )); + } + + let token_count = token_ids.len(); + let text_embeddings = self.text_embeddings(token_ids)?; + self.run_flow_main_prefix(&text_embeddings, &mut flow_state)?; + let t_prefix = t0.elapsed(); + 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)?; + let t_generate = t0.elapsed(); + let audio = self.decode_latents(&latents)?; + if phase_log { + eprintln!( + "tts-phase: condition={:.0}ms prefix={:.0}ms generate={:.0}ms decode={:.0}ms frames={} audio_s={:.2}", + t_condition.as_secs_f64() * 1e3, + (t_prefix - t_condition).as_secs_f64() * 1e3, + (t_generate - t_prefix).as_secs_f64() * 1e3, + (t0.elapsed() - t_generate).as_secs_f64() * 1e3, + latents.len() / self.bundle.latent_dim, + audio.len() as f64 / self.bundle.sample_rate as f64, + ); + } + Ok(audio) + } + + /// EXPERIMENTAL (latency): return a fresh Flow LM state conditioned on + /// the reference voice, restoring a cached snapshot when the same voice + /// samples were conditioned before. Keyed by voice content, like + /// `cached_voice` — never by buffer address. + fn conditioned_flow_state(&mut self, style: &VoiceStyle) -> Result, String> { + let key = voice_key(style); + if let Some(cached) = &self.cached_conditioning { + if cached.key == key { + return restore_state(&cached.state); + } + } let voice_embeddings = self.voice_embeddings(style)?; - let mut flow_state = self.condition_voice(&voice_embeddings)?; + let state = self.condition_voice(&voice_embeddings)?; + self.cached_conditioning = Some(CachedConditioning { + key, + state: snapshot_state(&state)?, + }); + Ok(state) + } + + /// EXPERIMENTAL (latency): streaming synthesis — interleaves the Flow LM + /// frame loop with incremental stateful Mimi decoding, invoking + /// `on_audio` with each decoded delta as soon as ~`emit_frames` latent + /// frames exist (80 ms of audio per frame). The Mimi decoder carries its + /// recurrent state across deltas, so the concatenated deltas are the same + /// audio `synth_chunk` would return. Returns Ok(false) when the callback + /// requested cancellation. + pub(crate) fn synth_chunk_streaming( + &mut self, + prepared: &AprilPreparedPrompt, + style: &VoiceStyle, + emit_frames: usize, + on_audio: &mut dyn FnMut(Vec) -> bool, + ) -> Result { + let mut flow_state = self.conditioned_flow_state(style)?; let token_ids = self .tokenizer .encode(prepared.text.as_str(), false) @@ -290,7 +491,7 @@ impl AprilPocketTts { .map(i64::from) .collect::>(); if token_ids.is_empty() { - return Ok(AprilSynthesisOutcome::Complete); + return Ok(true); } if token_ids.len() > self.bundle.max_token_per_chunk { return Err(format!( @@ -304,18 +505,165 @@ 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, interrupted) = self.generate_latents( - max_frames, - prepared.frames_after_eos, - &mut flow_state, - &mut callback, - )?; - if interrupted { - return Ok(AprilSynthesisOutcome::Interrupted); + let emit_frames = emit_frames.max(1); + + let mut mimi_state = initialize_state(&self.bundle.mimi_state_manifest)?; + let mut pending: Vec = Vec::with_capacity(emit_frames * self.bundle.latent_dim); + let mut current = vec![f32::NAN; self.bundle.latent_dim]; + let mut eos_step = None; + let mut rng = rand::rng(); + + for step in 0..max_frames { + let sequence = Tensor::from_array(( + vec![1_i64, 1, self.bundle.latent_dim as i64], + current.clone().into_boxed_slice(), + )) + .map_err(ort_error("create latent input"))?; + let text_embeddings = Tensor::::new( + &ort::memory::Allocator::default(), + [1_i64, 0, self.bundle.conditioning_dim as i64], + ) + .map_err(ort_error("create empty text input"))?; + let mut inputs = vec![ + (Cow::Borrowed("sequence"), SessionInputValue::from(sequence)), + ( + Cow::Borrowed("text_embeddings"), + SessionInputValue::from(text_embeddings), + ), + ]; + append_state_inputs(&mut inputs, &flow_state); + // Scoped: `outputs` borrows `self.flow_main`; it must drop before + // `decode_frames` takes `&mut self` below. + let (conditioning, eos_logit) = { + let mut outputs = self + .flow_main + .run(inputs) + .map_err(ort_error("run Pocket TTS Flow LM"))?; + let conditioning = outputs[0] + .try_extract_tensor::() + .map_err(ort_error("extract Flow LM conditioning"))? + .1 + .to_vec(); + let eos_logit = outputs[1] + .try_extract_tensor::() + .map_err(ort_error("extract Flow LM EOS logit"))? + .1 + .first() + .copied() + .ok_or_else(|| "Flow LM returned empty EOS logit".to_string())?; + replace_state_from_outputs(&mut flow_state, &mut outputs)?; + (conditioning, eos_logit) + }; + + if eos_logit > EOS_LOGIT_THRESHOLD && eos_step.is_none() { + eos_step = Some(step); + } + if eos_step.is_some_and(|eos| step >= eos + prepared.frames_after_eos) { + break; + } + + let mut noise = + normal_noise(&mut rng, self.bundle.latent_dim, DEFAULT_TEMPERATURE.sqrt()); + let conditioning = Tensor::from_array(( + vec![1_i64, self.bundle.conditioning_dim as i64], + conditioning.into_boxed_slice(), + )) + .map_err(ort_error("create flow conditioning"))?; + let s = Tensor::from_array((vec![1_i64, 1], vec![0.0_f32].into_boxed_slice())) + .map_err(ort_error("create flow start tensor"))?; + let t = Tensor::from_array((vec![1_i64, 1], vec![1.0_f32].into_boxed_slice())) + .map_err(ort_error("create flow end tensor"))?; + let x = Tensor::from_array(( + vec![1_i64, self.bundle.latent_dim as i64], + noise.clone().into_boxed_slice(), + )) + .map_err(ort_error("create flow noise tensor"))?; + let outputs = self + .flow + .run(ort::inputs![ + "c" => conditioning, + "s" => s, + "t" => t, + "x" => x, + ]) + .map_err(ort_error("run Pocket TTS flow"))?; + let flow = outputs[0] + .try_extract_tensor::() + .map_err(ort_error("extract Pocket TTS flow"))? + .1; + if flow.len() != noise.len() { + return Err(format!( + "flow returned {} values; expected {}", + flow.len(), + noise.len() + )); + } + for (sample, delta) in noise.iter_mut().zip(flow) { + *sample += *delta; + } + drop(outputs); + current.clone_from(&noise); + pending.extend_from_slice(&noise); + + if pending.len() >= emit_frames * self.bundle.latent_dim { + let audio = self.decode_frames(&pending, &mut mimi_state)?; + pending.clear(); + if !audio.is_empty() && !on_audio(audio) { + return Ok(false); + } + } } - self.decode_latents(&latents, |samples, progress| { - callback(samples, 0.5 + progress * 0.5) - }) + if !pending.is_empty() { + let audio = self.decode_frames(&pending, &mut mimi_state)?; + if !audio.is_empty() && !on_audio(audio) { + return Ok(false); + } + } + Ok(true) + } + + /// EXPERIMENTAL (latency): decode a batch of latent frames with a + /// caller-held Mimi state, so successive calls continue one stream. + fn decode_frames( + &mut self, + latents: &[f32], + state: &mut [StateValue], + ) -> Result, String> { + if latents.is_empty() { + return Ok(Vec::new()); + } + if !latents.len().is_multiple_of(self.bundle.latent_dim) { + return Err(format!( + "latent buffer has {} values, not divisible by {}", + latents.len(), + self.bundle.latent_dim + )); + } + let frame_count = latents.len() / self.bundle.latent_dim; + let mut audio = Vec::new(); + for start in (0..frame_count).step_by(DECODER_CHUNK_FRAMES) { + let end = (start + DECODER_CHUNK_FRAMES).min(frame_count); + let values = + latents[start * self.bundle.latent_dim..end * self.bundle.latent_dim].to_vec(); + let latent = Tensor::from_array(( + vec![1_i64, (end - start) as i64, self.bundle.latent_dim as i64], + values.into_boxed_slice(), + )) + .map_err(ort_error("create Mimi latent tensor"))?; + let mut inputs = vec![(Cow::Borrowed("latent"), SessionInputValue::from(latent))]; + append_state_inputs(&mut inputs, state); + let mut outputs = self + .mimi_decoder + .run(inputs) + .map_err(ort_error("run Mimi decoder"))?; + let samples = outputs[0] + .try_extract_tensor::() + .map_err(ort_error("extract Mimi audio"))? + .1; + audio.extend_from_slice(samples); + replace_state_from_outputs(state, &mut outputs)?; + } + Ok(audio) } fn prepared_token_count(&self, text: &str) -> Result { @@ -334,13 +682,9 @@ impl AprilPocketTts { } fn voice_embeddings(&mut self, style: &VoiceStyle) -> Result, String> { - let key = ( - style.samples.as_ptr() as usize, - style.samples.len(), - style.sample_rate, - ); + let key = voice_key(style); if let Some(cached) = &self.cached_voice { - if (cached.samples_ptr, cached.samples_len, cached.sample_rate) == key { + if cached.key == key { return Ok(cached.embeddings.clone()); } } @@ -381,9 +725,7 @@ impl AprilPocketTts { embeddings.extend_from_slice(&self.bos_embedding); embeddings.extend_from_slice(encoded); self.cached_voice = Some(CachedVoice { - samples_ptr: key.0, - samples_len: key.1, - sample_rate: key.2, + key, embeddings: embeddings.clone(), }); Ok(embeddings) @@ -475,29 +817,18 @@ 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], - callback: &mut F, - ) -> Result<(Vec, bool), String> - where - F: FnMut(&[f32], f32) -> bool, - { + ) -> Result, String> { 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(), @@ -583,19 +914,12 @@ impl AprilPocketTts { current.clone_from(&noise); latents.extend_from_slice(&noise); } - Ok((latents, false)) + Ok(latents) } - fn decode_latents( - &mut self, - latents: &[f32], - mut callback: F, - ) -> Result - where - F: FnMut(&[f32], f32) -> bool, - { + fn decode_latents(&mut self, latents: &[f32]) -> Result, String> { if latents.is_empty() { - return Ok(AprilSynthesisOutcome::Complete); + return Ok(Vec::new()); } if !latents.len().is_multiple_of(self.bundle.latent_dim) { return Err(format!( @@ -606,6 +930,7 @@ impl AprilPocketTts { } let frame_count = latents.len() / self.bundle.latent_dim; let mut state = initialize_state(&self.bundle.mimi_state_manifest)?; + let mut audio = Vec::new(); for start in (0..frame_count).step_by(DECODER_CHUNK_FRAMES) { let end = (start + DECODER_CHUNK_FRAMES).min(frame_count); @@ -625,17 +950,36 @@ impl AprilPocketTts { let samples = outputs[0] .try_extract_tensor::() .map_err(ort_error("extract Mimi audio"))? - .1 - .to_vec(); + .1; + audio.extend_from_slice(samples); replace_state_from_outputs(&mut state, &mut outputs)?; - if !callback(&samples, end as f32 / frame_count as f32) { - return Ok(AprilSynthesisOutcome::Interrupted); - } } - Ok(AprilSynthesisOutcome::Complete) + Ok(audio) } } +fn split_model_at_natural_boundaries( + text: &str, + max_tokens: usize, + token_count: F, +) -> Result, String> +where + F: FnMut(&str) -> Result, +{ + split_at_natural_boundaries(text, max_tokens, false, token_count) +} + +fn split_playback_at_natural_boundaries( + text: &str, + max_tokens: usize, + token_count: F, +) -> Result, String> +where + F: FnMut(&str) -> Result, +{ + split_at_natural_boundaries(text, max_tokens, true, token_count) +} + fn split_at_natural_boundaries( text: &str, max_tokens: usize, @@ -680,9 +1024,17 @@ where .chars() .next() .is_some_and(is_closing_punctuation); - if (!at_word_end && !at_clause_end) || token_count(&text[start..end])? > max_tokens { + if !at_word_end && !at_clause_end { continue; } + // Prepared token counts are monotonic in prefix length, so once a + // candidate overflows the limit no longer candidate can fit. Stop + // scanning instead of tokenizing every remaining boundary: that + // kept this loop superlinear in prompt length, and the cost landed + // before the first chunk reached synthesis. + if token_count(&text[start..end])? > max_tokens { + break; + } word_end = Some(end); match natural_boundary(&text[start..end], end == text.len()) { @@ -1003,14 +1355,91 @@ mod tests { assert_eq!(shape_len(&[2, 1, 8, 1000, 64]).expect("shape"), 1_024_000); } + /// The two engine splitters must keep OPPOSITE isolation polarity. + /// + /// The guards in `pocket.rs` pin which engine method each public API calls, + /// but they cannot see what the method itself does: pointing + /// `split_playback_prompt` at the model wrapper leaves every call site's + /// source text untouched while first-sentence isolation silently stops + /// happening, so the first playback unit becomes the whole utterance and + /// first audio waits on generating all of it. + #[test] + fn engine_splitters_keep_opposite_isolation_polarity() { + let source = include_str!("pocket_april.rs"); + let production = source + .split_once("\n#[cfg(test)]") + .map_or(source, |(production, _)| production); + + // A method's own code, and nothing else. Ending at the method's own + // closing brace keeps the NEXT method's doc comment out, and stripping + // `//` to end of line keeps prose out: neither can call a splitter, so + // scanning either reports drift in a method that has not changed. + let method_code = |name: &str| -> String { + let (_, body) = production + .split_once(name) + .unwrap_or_else(|| panic!("{name} exists")); + let (body, _) = body + .split_once("\n }\n") + .unwrap_or_else(|| panic!("{name} has a closing brace")); + body.lines() + .map(|line| line.split_once("//").map_or(line, |(code, _)| code)) + .collect::>() + .join("\n") + }; + let model = method_code("fn split_prompt"); + let model = model.as_str(); + let playback = method_code("fn split_playback_prompt"); + let playback = playback.as_str(); + + assert_eq!( + ( + model.matches("split_model_at_natural_boundaries(").count(), + model + .matches("split_playback_at_natural_boundaries(") + .count(), + ), + (1, 0), + "split_prompt must pack sentences: isolating here peels sentence \ + one off every already-packed unit" + ); + assert_eq!( + ( + playback + .matches("split_playback_at_natural_boundaries(") + .count(), + playback + .matches("split_model_at_natural_boundaries(") + .count(), + ), + (1, 0), + "split_playback_prompt must isolate sentence one: packing here \ + makes the first playback unit the whole utterance and delays \ + first audio by the full generation" + ); + + // Calling the isolating splitter is necessary but not sufficient: a + // short circuit before the call can return the whole utterance as one + // unit while leaving the delegated splitter unchanged. Playback must + // delegate unconditionally so sentence one remains the first unit. + for control_flow in ["if ", "match ", "else", "return"] { + assert!( + !playback.contains(control_flow), + "split_playback_prompt must delegate unconditionally, found \ + `{control_flow}`: a branch before the split can return the \ + whole utterance as the first playback unit, delaying first \ + audio by the full generation" + ); + } + } + fn whitespace_token_count(text: &str) -> Result { Ok(text.split_whitespace().count()) } #[test] - fn natural_split_keeps_first_sentence_separate_then_packs_the_remainder() { + fn playback_split_keeps_first_sentence_separate_then_packs_the_remainder() { let text = "One two. Three four. Five six."; - let chunks = split_at_natural_boundaries(text, 4, true, whitespace_token_count).unwrap(); + let chunks = split_playback_at_natural_boundaries(text, 4, whitespace_token_count).unwrap(); assert_eq!(chunks, ["One two. ", "Three four. Five six."]); assert_eq!(chunks.concat(), text); } @@ -1018,11 +1447,27 @@ mod tests { #[test] fn model_split_packs_multiple_sentences_within_limit() { let text = "One two. Three four. Five six."; - let chunks = split_at_natural_boundaries(text, 4, false, whitespace_token_count).unwrap(); + let chunks = split_model_at_natural_boundaries(text, 4, whitespace_token_count).unwrap(); assert_eq!(chunks, ["One two. Three four. ", "Five six."]); assert_eq!(chunks.concat(), text); } + #[test] + fn playback_then_model_split_does_not_isolate_later_sentences_again() { + let text = "Alpha one. Beta two. Gamma three."; + let playback = + split_playback_at_natural_boundaries(text, 50, whitespace_token_count).unwrap(); + assert_eq!(playback, ["Alpha one. ", "Beta two. Gamma three."]); + + let model: Vec<_> = playback + .iter() + .flat_map(|chunk| { + split_model_at_natural_boundaries(chunk.trim(), 50, whitespace_token_count).unwrap() + }) + .collect(); + assert_eq!(model, ["Alpha one.", "Beta two. Gamma three."]); + } + #[test] fn natural_split_prefers_preceding_sentence_boundary() { let text = "One two. Three four five six."; @@ -1086,6 +1531,39 @@ mod tests { assert_eq!(chunks.concat(), text); } + #[test] + fn natural_split_stops_counting_tokens_past_the_limit() { + // Each boundary scan must stop at the first overflowing candidate + // rather than tokenizing every remaining boundary. Scanning to + // end-of-text makes tokenizer input grow superlinearly in prompt + // length, and that cost is paid before the first chunk reaches + // synthesis, taxing time-to-first-audio on long prompts. + let sentence = "The relay finished its migration and the channel list refreshed. "; + let tokenized_bytes = |repeats: usize| -> usize { + let text = sentence.repeat(repeats).trim_end().to_string(); + let total = std::cell::Cell::new(0_usize); + let chunks = split_at_natural_boundaries(&text, 50, true, |chunk| { + total.set(total.get() + chunk.len()); + whitespace_token_count(chunk) + }) + .expect("split repeated sentences"); + assert_eq!(chunks.concat(), text); + assert!(chunks.len() > 1); + total.get() + }; + + // Doubling the prompt must not multiply tokenizer work superlinearly. + // Bounded scans grow ~2x here; scanning to end-of-text grows ~5.5x. + let single = tokenized_bytes(12); + let double = tokenized_bytes(24); + assert!( + double < single * 3, + "doubling the prompt grew tokenizer input from {single} to {double} bytes \ + ({:.1}x); bounded scans stay near 2x", + double as f64 / single as f64, + ); + } + #[test] fn normal_noise_has_requested_length() { let mut rng = rand::rng(); @@ -1099,6 +1577,234 @@ mod tests { assert_eq!(estimate_max_frames(300, 12.5), 1_275); } + /// Regression (review finding): the voice caches must key on CONTENT. + /// Voice switching clones and drops sample buffers, so a new voice with + /// the same length and rate can land at a recycled address — an + /// address-based key would then restore the previous voice's state and + /// speak with the wrong voice. + #[test] + fn voice_key_is_content_based_not_address_based() { + let style_a = VoiceStyle { + samples: vec![0.1, -0.2, 0.3, -0.4], + sample_rate: 24_000, + }; + // Same length, same rate, different content — MUST key differently, + // regardless of what address the allocator hands out. + let style_b = VoiceStyle { + samples: vec![0.4, -0.3, 0.2, -0.1], + sample_rate: 24_000, + }; + assert_ne!(voice_key(&style_a), voice_key(&style_b)); + + // Same content in a fresh allocation — MUST key identically, so the + // cache still hits across clones of the same voice. + let style_a_clone = VoiceStyle { + samples: style_a.samples.clone(), + sample_rate: style_a.sample_rate, + }; + assert_ne!( + style_a.samples.as_ptr(), + style_a_clone.samples.as_ptr(), + "clone must be a distinct allocation for this test to mean anything" + ); + assert_eq!(voice_key(&style_a), voice_key(&style_a_clone)); + + // Same content at a different rate is a different voice identity. + let style_a_resampled = VoiceStyle { + samples: style_a.samples.clone(), + sample_rate: 16_000, + }; + assert_ne!(voice_key(&style_a), voice_key(&style_a_resampled)); + } + + #[test] + #[ignore = "requires BERD_POCKET_TEST_MODEL_DIR"] + fn switching_between_equal_length_voices_reconditions_the_flow_state() { + let dir = std::env::var("BERD_POCKET_TEST_MODEL_DIR") + .expect("set BERD_POCKET_TEST_MODEL_DIR to the verified April bundle"); + let style_a = + crate::pocket::load_voice_style(&Path::new(&dir).join("reference_sample.wav")) + .expect("load reference voice"); + // Voice B: same length, same rate, different content (reversed + // samples) — the exact shape an address-recycling collision takes. + let style_b = VoiceStyle { + samples: style_a.samples.iter().rev().copied().collect(), + sample_rate: style_a.sample_rate, + }; + assert_eq!(style_a.samples.len(), style_b.samples.len()); + + // Engine 1: condition A (primes both caches), then switch to B. + let mut engine = AprilPocketTts::load(Path::new(&dir), 1).expect("load April bundle"); + let state_a = snapshot_state( + &engine + .conditioned_flow_state(&style_a) + .expect("condition A"), + ) + .expect("snapshot A"); + let state_b_after_switch = snapshot_state( + &engine + .conditioned_flow_state(&style_b) + .expect("condition B"), + ) + .expect("snapshot B after switch"); + // Warm hit on the SAME voice: the cached restore must reproduce the + // original conditioning bit-for-bit (cache warm == cache cold). + let state_b_warm_hit = snapshot_state( + &engine + .conditioned_flow_state(&style_b) + .expect("condition B warm"), + ) + .expect("snapshot B warm hit"); + + // Engine 2: fresh process conditions B with no cache in play. + let mut fresh = AprilPocketTts::load(Path::new(&dir), 1).expect("load April bundle"); + let state_b_fresh = snapshot_state( + &fresh + .conditioned_flow_state(&style_b) + .expect("condition B fresh"), + ) + .expect("snapshot B fresh"); + + // The switched state must equal a from-scratch conditioning of B and + // must NOT be A's cached state. + assert!( + snapshots_equal(&state_b_after_switch, &state_b_fresh), + "switching voices must recondition, not replay the cache" + ); + assert!( + !snapshots_equal(&state_b_after_switch, &state_a), + "equal-length distinct voices must produce distinct conditioning" + ); + // And the warm cache hit must be indistinguishable from recomputing. + assert!( + snapshots_equal(&state_b_warm_hit, &state_b_fresh), + "a warm conditioning-cache hit must equal a cold recompute" + ); + } + + fn snapshots_equal( + a: &[(StateSpec, SnapshotTensor)], + b: &[(StateSpec, SnapshotTensor)], + ) -> bool { + // f32 compares bitwise: state tensors legitimately contain NaN fill, + // and NaN != NaN under float equality would make identical states + // compare unequal. + a.len() == b.len() + && a.iter().zip(b).all(|((_, ta), (_, tb))| match (ta, tb) { + (SnapshotTensor::F32(sa, da), SnapshotTensor::F32(sb, db)) => { + sa == sb + && da.len() == db.len() + && da.iter().zip(db).all(|(x, y)| x.to_bits() == y.to_bits()) + } + (SnapshotTensor::I64(sa, da), SnapshotTensor::I64(sb, db)) => sa == sb && da == db, + (SnapshotTensor::Bool(sa, da), SnapshotTensor::Bool(sb, db)) => { + sa == sb && da == db + } + _ => false, + }) + } + + #[test] + #[ignore = "requires BERD_POCKET_TEST_MODEL_DIR"] + fn incremental_stateful_decode_matches_batch_decode() { + let dir = std::env::var("BERD_POCKET_TEST_MODEL_DIR") + .expect("set BERD_POCKET_TEST_MODEL_DIR to the verified April bundle"); + let mut engine = AprilPocketTts::load(Path::new(&dir), 1).expect("load April bundle"); + let style = crate::pocket::load_voice_style(&Path::new(&dir).join("reference_sample.wav")) + .expect("load reference voice"); + + // Generate one real latent sequence (the RNG makes repeat synths + // differ, so both decode paths must consume the SAME latents). + let prepared = + prepare_april_prompt("The relay deploy finished and every check passed cleanly.") + .expect("prepare prompt"); + let mut flow_state = engine + .conditioned_flow_state(&style) + .expect("condition voice"); + let token_ids = engine + .tokenizer + .encode(prepared.text.as_str(), false) + .expect("tokenize") + .get_ids() + .iter() + .copied() + .map(i64::from) + .collect::>(); + let token_count = token_ids.len(); + let text_embeddings = engine.text_embeddings(token_ids).expect("text embeddings"); + engine + .run_flow_main_prefix(&text_embeddings, &mut flow_state) + .expect("prefix"); + let max_frames = estimate_max_frames(token_count, engine.bundle.frame_rate); + let latents = engine + .generate_latents(max_frames, prepared.frames_after_eos, &mut flow_state) + .expect("generate latents"); + let frame_count = latents.len() / engine.bundle.latent_dim; + assert!( + frame_count > DECODER_CHUNK_FRAMES, + "need a multi-chunk case" + ); + + // Batch: the production decode (fresh state, 12-frame steps). + let batch = engine.decode_latents(&latents).expect("batch decode"); + + // Incremental chunkings: 12-frame deltas through one carried Mimi + // state must be bit-exact (the production batch path itself steps by + // DECODER_CHUNK_FRAMES=12 through one state). Sub-12 chunkings are + // measured for the record but are NOT exact — the decoder has + // intra-chunk lookahead — so streaming must emit at >= 12 frames. + for delta_frames in [6usize, 4, 2, 1] { + let mut state = + initialize_state(&engine.bundle.mimi_state_manifest).expect("mimi state"); + let mut streamed = Vec::new(); + for chunk in latents.chunks(delta_frames * engine.bundle.latent_dim) { + streamed.extend( + engine + .decode_frames(chunk, &mut state) + .expect("delta decode"), + ); + } + assert_eq!(batch.len(), streamed.len(), "sample count must match"); + let max_diff = batch + .iter() + .zip(&streamed) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + let rms_batch = (batch.iter().map(|s| s * s).sum::() / batch.len() as f32).sqrt(); + let rms_err = (batch + .iter() + .zip(&streamed) + .map(|(a, b)| (a - b) * (a - b)) + .sum::() + / batch.len() as f32) + .sqrt(); + eprintln!( + "delta_frames={delta_frames}: max|diff|={max_diff:.6} rms_err={rms_err:.6} snr_db={:.1}", + 20.0 * (rms_batch / rms_err.max(1e-12)).log10() + ); + } + let mut state = initialize_state(&engine.bundle.mimi_state_manifest).expect("mimi state"); + let mut streamed = Vec::new(); + for chunk in latents.chunks(DECODER_CHUNK_FRAMES * engine.bundle.latent_dim) { + streamed.extend( + engine + .decode_frames(chunk, &mut state) + .expect("delta decode"), + ); + } + + assert_eq!(batch.len(), streamed.len(), "sample count must match"); + let max_diff = batch + .iter() + .zip(&streamed) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + assert!( + max_diff <= 1.0e-4, + "incremental decode diverged from batch decode: max |diff| = {max_diff}" + ); + } + #[test] #[ignore = "requires BERD_POCKET_TEST_MODEL_DIR"] fn tokenizer_matches_sentencepiece_reference_including_unknown_words() { diff --git a/src-tauri/src/commands/pocket_voice.rs b/src-tauri/src/commands/pocket_voice.rs index 4345f4f09..85fdea2a2 100644 --- a/src-tauri/src/commands/pocket_voice.rs +++ b/src-tauri/src/commands/pocket_voice.rs @@ -13,7 +13,7 @@ use std::time::{Duration, Instant, SystemTime}; #[cfg(target_os = "macos")] use berd_voice::SAMPLE_RATE; #[cfg(target_os = "macos")] -use berd_voice::{load_text_to_speech, load_voice_style, SynthesisOutcome}; +use berd_voice::{load_text_to_speech, load_voice_style}; use futures_util::StreamExt; #[cfg(target_os = "macos")] use rodio::buffer::SamplesBuffer; @@ -36,6 +36,8 @@ const DOWNLOAD_PROGRESS_EMIT_INTERVAL: Duration = Duration::from_millis(100); const DOWNLOAD_CONNECT_TIMEOUT: Duration = Duration::from_secs(15); const DOWNLOAD_READ_TIMEOUT: Duration = Duration::from_secs(30); const DOWNLOAD_TOTAL_TIMEOUT: Duration = Duration::from_secs(30 * 60); +#[cfg(target_os = "macos")] +const STREAMING_EMIT_FRAMES: usize = 12; const PARAKEET_ARCHIVE: Artifact = Artifact { filename: "parakeet.tar.bz2", size: 104_337_827, @@ -1597,31 +1599,6 @@ async fn download_artifact( Ok(()) } -#[cfg(any(target_os = "macos", test))] -fn cumulative_delta<'a>( - previous_len: &mut usize, - samples: &'a [f32], -) -> Result, String> { - // Pocket invokes the callback with an empty slice while generating the - // next internal model unit. It is progress-only, not a cumulative reset. - if samples.is_empty() { - return Ok(None); - } - if samples.len() < *previous_len { - return Err(format!( - "Pocket cumulative callback length decreased from {} to {}", - *previous_len, - samples.len() - )); - } - if samples.len() == *previous_len { - return Ok(None); - } - let delta = &samples[*previous_len..]; - *previous_len = samples.len(); - Ok(Some(delta)) -} - #[cfg(target_os = "macos")] fn synthesize_and_stream( base: &Path, @@ -1676,7 +1653,6 @@ fn synthesize_and_stream( let rate = NonZero::new(SAMPLE_RATE) .ok_or_else(|| "Pocket sample rate invariant failed".to_string())?; let player = Arc::new(Player::connect_new(sink.mixer())); - let previous_len = Arc::new(Mutex::new(0_usize)); let speed_processor = Rc::new(RefCell::new(StreamingSpeedProcessor::new( speed, SAMPLE_RATE, @@ -1686,53 +1662,40 @@ fn synthesize_and_stream( let callback_player = player.clone(); let callback_active = active.clone(); - let callback_previous_len = previous_len.clone(); let callback_speed_processor = speed_processor.clone(); let callback_error_slot = callback_error.clone(); let callback_started = playback_started.clone(); - let outcome = - engine.synth_chunk_streaming(text, &style, move |samples: &[f32], _progress: f32| { - if !callback_active.load(Ordering::SeqCst) { + let mut on_audio = move |samples: Vec| { + if !callback_active.load(Ordering::SeqCst) { + return false; + } + let delta = match callback_speed_processor.borrow_mut().process(&samples) { + Ok(processed) => processed, + Err(error) => { + if let Ok(mut callback_error) = callback_error_slot.lock() { + *callback_error = Some(error); + } return false; } - let Ok(mut previous_len) = callback_previous_len.lock() else { - return false; - }; - let delta = match cumulative_delta(&mut previous_len, samples) { - Ok(Some(delta)) => match callback_speed_processor.borrow_mut().process(delta) { - Ok(processed) => processed, - Err(error) => { - if let Ok(mut callback_error) = callback_error_slot.lock() { - *callback_error = Some(error); - } - return false; - } - }, - Ok(None) => return true, - Err(error) => { + }; + if !delta.is_empty() { + callback_player.append(SamplesBuffer::new(channels, rate, delta)); + if !callback_started.swap(true, Ordering::SeqCst) { + println!("VOICE_CONVERSATION_PLAYBACK_STARTED"); + if let Err(error) = std::io::stdout().flush() { if let Ok(mut callback_error) = callback_error_slot.lock() { - *callback_error = Some(error); + *callback_error = Some(format!("signal Pocket playback start: {error}")); } return false; } - }; - if !delta.is_empty() { - callback_player.append(SamplesBuffer::new(channels, rate, delta)); - if !callback_started.swap(true, Ordering::SeqCst) { - println!("VOICE_CONVERSATION_PLAYBACK_STARTED"); - if let Err(error) = std::io::stdout().flush() { - if let Ok(mut callback_error) = callback_error_slot.lock() { - *callback_error = - Some(format!("signal Pocket playback start: {error}")); - } - return false; - } - } } - true - })?; + } + true + }; + let completed = + engine.synth_chunk_streaming(text, &style, STREAMING_EMIT_FRAMES, &mut on_audio)?; - if matches!(outcome, SynthesisOutcome::Complete(_)) { + if completed { let tail = speed_processor.borrow_mut().finish()?; if !tail.is_empty() { player.append(SamplesBuffer::new(channels, rate, tail)); @@ -1753,7 +1716,7 @@ fn synthesize_and_stream( player.stop(); return Err(error); } - if matches!(outcome, SynthesisOutcome::Interrupted) { + if !completed { player.stop(); return Ok(()); } @@ -1956,32 +1919,6 @@ mod tests { assert_ne!(original, replacement); } - #[test] - fn cumulative_callback_emits_only_growth_and_rejects_regression() { - let mut previous_len = 0; - assert_eq!(cumulative_delta(&mut previous_len, &[]), Ok(None)); - assert_eq!( - cumulative_delta(&mut previous_len, &[1.0, 2.0, 3.0]), - Ok(Some(&[1.0, 2.0, 3.0][..])) - ); - assert_eq!(previous_len, 3); - assert_eq!(cumulative_delta(&mut previous_len, &[]), Ok(None)); - assert_eq!(previous_len, 3); - assert_eq!( - cumulative_delta(&mut previous_len, &[1.0, 2.0, 3.0]), - Ok(None) - ); - assert_eq!( - cumulative_delta(&mut previous_len, &[1.0, 2.0, 3.0, 4.0]), - Ok(Some(&[4.0][..])) - ); - assert_eq!(previous_len, 4); - assert_eq!( - cumulative_delta(&mut previous_len, &[1.0, 2.0]), - Err("Pocket cumulative callback length decreased from 4 to 2".to_string()) - ); - } - #[test] fn failed_atomic_publication_restores_previous_cache() { let directory = tempfile::tempdir().expect("temporary directory"); From 5204fe485605cafed71da22746dc1c786909dd68 Mon Sep 17 00:00:00 2001 From: John Tennant Date: Wed, 19 Aug 2026 17:42:11 -0400 Subject: [PATCH 02/10] fix(voice): preserve streaming cancellation contracts --- src-tauri/crates/berd-voice/src/pocket.rs | 196 ++++++++--------- .../crates/berd-voice/src/pocket_april.rs | 197 ++++++------------ src-tauri/src/commands/pocket_voice.rs | 3 + 3 files changed, 152 insertions(+), 244 deletions(-) diff --git a/src-tauri/crates/berd-voice/src/pocket.rs b/src-tauri/crates/berd-voice/src/pocket.rs index bbf456823..1886549b8 100644 --- a/src-tauri/crates/berd-voice/src/pocket.rs +++ b/src-tauri/crates/berd-voice/src/pocket.rs @@ -14,6 +14,7 @@ //! Berd's Pocket model installer writes the complete attribution beside the //! cached model files. +use std::cell::RefCell; use std::path::{Path, PathBuf}; use std::sync::Mutex; @@ -35,8 +36,44 @@ pub const SAMPLE_RATE: u32 = 24_000; const TTS_NUM_THREADS: usize = 1; -/// EXPERIMENTAL (latency): override ONNX intra-op threads for the Pocket -/// sessions via `BERD_TTS_THREADS`. Default preserves production's 1. +thread_local! { + static ACTIVE_SYNTHESIS_ENGINES: RefCell> = const { RefCell::new(Vec::new()) }; +} + +struct SynthesisCallGuard { + engine_id: usize, +} + +impl SynthesisCallGuard { + fn enter(engine_id: usize) -> Result { + ACTIVE_SYNTHESIS_ENGINES.with(|active| { + let mut active = active.borrow_mut(); + if active.contains(&engine_id) { + return Err("Pocket TTS callback re-entered the active engine".to_string()); + } + active.push(engine_id); + Ok(Self { engine_id }) + }) + } + + fn is_active(engine_id: usize) -> bool { + ACTIVE_SYNTHESIS_ENGINES.with(|active| active.borrow().contains(&engine_id)) + } +} + +impl Drop for SynthesisCallGuard { + fn drop(&mut self) { + ACTIVE_SYNTHESIS_ENGINES.with(|active| { + let mut active = active.borrow_mut(); + if let Some(index) = active.iter().rposition(|engine| *engine == self.engine_id) { + active.remove(index); + } + }); + } +} + +/// Return the configured ONNX intra-op thread count for Pocket sessions. +/// `BERD_TTS_THREADS` overrides the single-thread default when set. fn tts_num_threads() -> usize { std::env::var("BERD_TTS_THREADS") .ok() @@ -86,6 +123,7 @@ impl PocketTts { /// Split text into model-safe synthesis units that satisfy the bundle's /// exact 50-token input limit, packing sentences whenever they fit. pub fn split_text_into_chunks(&self, text: &str) -> Result, String> { + self.reject_reentry()?; let Some(prepared) = prepare_april_prompt(text) else { return Ok(Vec::new()); }; @@ -95,23 +133,6 @@ impl PocketTts { .split_prompt(&prepared) } - /// Split text into ordered playback units, keeping the first sentence - /// separate so it reaches synthesis before the remainder is packed. - /// - /// Units are contiguous substrings of the prepared model prompt and may - /// retain boundary whitespace. Concatenating them with `chunks.concat()` - /// reconstructs that prompt exactly, and each unit's prepared token count - /// is at most 50. - pub fn split_text_for_playback(&self, text: &str) -> Result, 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_playback_prompt(&prepared) - } - /// Synthesize text with the supplied reference voice. /// /// Pocket detects language from text and this model uses one synthesis @@ -123,6 +144,7 @@ impl PocketTts { style: &VoiceStyle, _steps: usize, ) -> Result, String> { + self.reject_reentry()?; let Some(prepared) = prepare_april_prompt(text) else { return Ok(Vec::new()); }; @@ -140,11 +162,12 @@ impl PocketTts { Ok(samples) } - /// EXPERIMENTAL (latency): streaming synthesis. Invokes `on_audio` with - /// PCM deltas as soon as roughly `emit_frames` Flow LM frames (80 ms of - /// audio each) have been generated and decoded. Concatenated deltas equal - /// one `synth_chunk` result. The callback runs on the caller thread and - /// returns `false` to cancel; the function then returns Ok(false). + /// Stream synthesis as PCM deltas become decoder-safe. `emit_frames` is + /// rounded down to a positive multiple of the Mimi decoder's 12-frame + /// chunk size. Concatenated non-empty deltas equal one `synth_chunk` + /// result. The callback runs on the caller thread and may receive empty + /// deltas so cancellation is observed before PCM is available. Returning + /// `false` cancels synthesis and makes the function return `Ok(false)`. pub fn synth_chunk_streaming( &self, text: &str, @@ -152,6 +175,7 @@ impl PocketTts { emit_frames: usize, on_audio: &mut dyn FnMut(Vec) -> bool, ) -> Result { + let _call_guard = SynthesisCallGuard::enter(self as *const Self as usize)?; let Some(prepared) = prepare_april_prompt(text) else { return Ok(true); }; @@ -159,16 +183,48 @@ impl PocketTts { .inner .lock() .map_err(|_| "Pocket TTS engine lock poisoned".to_string())?; - let chunks = engine.split_prompt(&prepared)?; + let chunks = engine.split_playback_prompt(&prepared)?; for chunk in chunks { + if !callback_allows_audio(on_audio, Vec::new())? { + return Ok(false); + } let prepared = prepare_april_prompt(&chunk) .ok_or_else(|| "Pocket TTS prompt chunk became empty".to_string())?; - if !engine.synth_chunk_streaming(&prepared, style, emit_frames, on_audio)? { + let mut callback_error = None; + let completed = + engine.synth_chunk_streaming(&prepared, style, emit_frames, &mut |audio| { + match callback_allows_audio(on_audio, audio) { + Ok(allowed) => allowed, + Err(error) => { + callback_error = Some(error); + false + } + } + })?; + if let Some(error) = callback_error { + return Err(error); + } + if !completed { return Ok(false); } } Ok(true) } + + fn reject_reentry(&self) -> Result<(), String> { + if SynthesisCallGuard::is_active(self as *const Self as usize) { + return Err("Pocket TTS callback re-entered the active engine".to_string()); + } + Ok(()) + } +} + +fn callback_allows_audio( + callback: &mut dyn FnMut(Vec) -> bool, + audio: Vec, +) -> Result { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| callback(audio))) + .map_err(|_| "Pocket TTS synthesis callback panicked".to_string()) } #[cfg(test)] @@ -190,89 +246,19 @@ mod tests { .any(|artifact| artifact.filename == "flow_lm_main.onnx")); } - /// Which splitter each production function delegates to, across the whole - /// file rather than one hand-picked window. - /// - /// A wrong delegation can reinstate either shipped defect in one token: - /// removing first-sentence priority from playback, or re-isolating sentence - /// one inside units that already fit. Asserting the whole map means a new - /// delegation must be declared here to compile green. - fn splitter_delegations(source: &str) -> Vec<(String, Vec)> { - let production = source - .split_once("\n#[cfg(test)]") - .map_or(source, |(production, _)| production); - // Scan code only. Prose cannot call a splitter, but it can contain - // ` fn `, which would end a body early and hide a call after it, and it - // can name a splitter, which would report a call the code never makes. - let production: String = production - .lines() - .map(|line| line.split_once("//").map_or(line, |(code, _)| code)) - .collect::>() - .join("\n"); - let mut out = Vec::new(); - let mut rest = production.as_str(); - while let Some((_, after)) = rest.split_once(" fn ") { - let (name, body) = after - .split_once('(') - .expect("a function signature has an argument list"); - // End at this function's own closing brace, not at the next ` fn `: - // a body provably stops where its braces balance, so no later - // function's calls are attributed here and none of this one's are - // dropped. - let inner = body.split_once('{').map_or("", |(_, inner)| inner); - let mut depth = 1usize; - let body = inner - .char_indices() - .find(|&(_, ch)| { - depth = match ch { - '{' => depth + 1, - '}' => depth - 1, - _ => depth, - }; - depth == 0 - }) - .map_or(inner, |(end, _)| &inner[..end]); - let mut calls = Vec::new(); - // Check the isolating spelling first: ".split_prompt(" is a - // substring of neither, but a naive contains() on the shorter name - // would also match the longer one. - for _ in 0..body.matches(".split_playback_prompt(").count() { - calls.push("split_playback_prompt".to_string()); - } - let plain = body.matches(".split_prompt(").count(); - for _ in 0..plain { - calls.push("split_prompt".to_string()); - } - if !calls.is_empty() { - out.push((name.trim().to_string(), calls)); - } - rest = after; - } - out + #[test] + fn active_engine_reentry_is_rejected() { + let _guard = SynthesisCallGuard::enter(42).expect("first call"); + assert!(SynthesisCallGuard::enter(42).is_err()); + assert!(SynthesisCallGuard::is_active(42)); } #[test] - fn every_production_splitter_delegation_is_declared() { - let source = include_str!("pocket.rs"); - let actual = splitter_delegations(source); - let expected: Vec<(String, Vec)> = vec![ - // Model units: pack sentences, never isolate. - ("split_text_into_chunks".into(), vec!["split_prompt".into()]), - // Playback units: isolate sentence one for time-to-first-audio. - ( - "split_text_for_playback".into(), - vec!["split_playback_prompt".into()], - ), - // Synthesis receives an already-packed unit: re-isolating here - // re-adds the per-sentence seam this PR removes. - ("synth_chunk".into(), vec!["split_prompt".into()]), - ("synth_chunk_streaming".into(), vec!["split_prompt".into()]), - ]; + fn callback_panic_is_reported_without_unwinding() { + let mut callback = |_: Vec| -> bool { panic!("callback failure") }; assert_eq!( - actual, expected, - "a production function changed which splitter it calls (or a new \ - one appeared); isolating outside split_text_for_playback delays \ - first audio, packing inside it removes the guarantee" + callback_allows_audio(&mut callback, Vec::new()).unwrap_err(), + "Pocket TTS synthesis callback panicked" ); } diff --git a/src-tauri/crates/berd-voice/src/pocket_april.rs b/src-tauri/crates/berd-voice/src/pocket_april.rs index 1de647a07..324ee0103 100644 --- a/src-tauri/crates/berd-voice/src/pocket_april.rs +++ b/src-tauri/crates/berd-voice/src/pocket_april.rs @@ -36,6 +36,11 @@ const DECODER_CHUNK_FRAMES: usize = 12; const TOKENS_PER_SECOND_ESTIMATE: f32 = 3.0; const GENERATION_SECONDS_PADDING: f32 = 2.0; +fn decoder_safe_emit_frames(requested: usize) -> usize { + let requested = requested.max(DECODER_CHUNK_FRAMES); + requested - requested % DECODER_CHUNK_FRAMES +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum TextBoundary { Sentence, @@ -126,9 +131,8 @@ struct CachedVoice { embeddings: Vec, } -/// EXPERIMENTAL (latency): a dtype-tagged copy of one recurrent state tensor, -/// used to snapshot the Flow LM state right after voice conditioning so -/// subsequent chunks skip the ~160 ms `condition_voice` pass entirely. +/// A dtype-tagged copy of one recurrent state tensor used to snapshot the Flow +/// LM state after voice conditioning so subsequent chunks can restore it. enum SnapshotTensor { F32(Vec, Vec), I64(Vec, Vec), @@ -229,9 +233,8 @@ pub(crate) struct AprilPocketTts { flow: Session, mimi_decoder: Session, cached_voice: Option, - /// EXPERIMENTAL (latency): post-`condition_voice` Flow LM state, cached - /// per reference voice. Restoring it replaces the ~160 ms conditioning - /// pass on every chunk after the first for a given voice. + /// Post-`condition_voice` Flow LM state cached per reference voice. + /// Restoring it avoids repeating conditioning after the first chunk. cached_conditioning: Option, } @@ -377,18 +380,22 @@ impl AprilPocketTts { if self.prepared_token_count(&prepared.text)? <= self.bundle.max_token_per_chunk { return Ok(vec![prepared.text.clone()]); } - split_model_at_natural_boundaries(&prepared.text, self.bundle.max_token_per_chunk, |text| { - self.prepared_token_count(text) - }) + split_at_natural_boundaries( + &prepared.text, + self.bundle.max_token_per_chunk, + false, + |text| self.prepared_token_count(text), + ) } pub(crate) fn split_playback_prompt( &self, prepared: &AprilPreparedPrompt, ) -> Result, String> { - split_playback_at_natural_boundaries( + split_at_natural_boundaries( &prepared.text, self.bundle.max_token_per_chunk, + true, |text| self.prepared_token_count(text), ) } @@ -398,7 +405,7 @@ impl AprilPocketTts { prepared: &AprilPreparedPrompt, style: &VoiceStyle, ) -> Result, String> { - // EXPERIMENTAL (latency bench): phase timing, enabled by BERD_TTS_PHASE_LOG=1. + // Optional synthesis phase timing, enabled by BERD_TTS_PHASE_LOG=1. let phase_log = std::env::var("BERD_TTS_PHASE_LOG").is_ok_and(|v| v == "1"); let t0 = std::time::Instant::now(); let mut flow_state = self.conditioned_flow_state(style)?; @@ -446,10 +453,10 @@ impl AprilPocketTts { Ok(audio) } - /// EXPERIMENTAL (latency): return a fresh Flow LM state conditioned on - /// the reference voice, restoring a cached snapshot when the same voice - /// samples were conditioned before. Keyed by voice content, like - /// `cached_voice` — never by buffer address. + /// Return a fresh Flow LM state conditioned on the reference voice, + /// restoring a cached snapshot when the same voice samples were + /// conditioned before. Keyed by voice content, like `cached_voice` — + /// never by buffer address. fn conditioned_flow_state(&mut self, style: &VoiceStyle) -> Result, String> { let key = voice_key(style); if let Some(cached) = &self.cached_conditioning { @@ -466,13 +473,13 @@ impl AprilPocketTts { Ok(state) } - /// EXPERIMENTAL (latency): streaming synthesis — interleaves the Flow LM - /// frame loop with incremental stateful Mimi decoding, invoking - /// `on_audio` with each decoded delta as soon as ~`emit_frames` latent - /// frames exist (80 ms of audio per frame). The Mimi decoder carries its - /// recurrent state across deltas, so the concatenated deltas are the same - /// audio `synth_chunk` would return. Returns Ok(false) when the callback - /// requested cancellation. + /// Interleave the Flow LM frame loop with incremental stateful Mimi + /// decoding, invoking `on_audio` with each decoded delta as soon as the + /// configured latent-frame interval is available. Empty callbacks expose + /// cancellation points before decoded PCM exists. The Mimi decoder carries + /// its recurrent state across deltas, so concatenated non-empty deltas are + /// the same audio `synth_chunk` would return. Returns `Ok(false)` when the + /// callback requests cancellation. pub(crate) fn synth_chunk_streaming( &mut self, prepared: &AprilPreparedPrompt, @@ -480,6 +487,9 @@ impl AprilPocketTts { emit_frames: usize, on_audio: &mut dyn FnMut(Vec) -> bool, ) -> Result { + if !on_audio(Vec::new()) { + return Ok(false); + } let mut flow_state = self.conditioned_flow_state(style)?; let token_ids = self .tokenizer @@ -505,7 +515,7 @@ 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 emit_frames = emit_frames.max(1); + let emit_frames = decoder_safe_emit_frames(emit_frames); let mut mimi_state = initialize_state(&self.bundle.mimi_state_manifest)?; let mut pending: Vec = Vec::with_capacity(emit_frames * self.bundle.latent_dim); @@ -514,6 +524,9 @@ impl AprilPocketTts { let mut rng = rand::rng(); for step in 0..max_frames { + if !on_audio(Vec::new()) { + return Ok(false); + } let sequence = Tensor::from_array(( vec![1_i64, 1, self.bundle.latent_dim as i64], current.clone().into_boxed_slice(), @@ -622,8 +635,8 @@ impl AprilPocketTts { Ok(true) } - /// EXPERIMENTAL (latency): decode a batch of latent frames with a - /// caller-held Mimi state, so successive calls continue one stream. + /// Decode a batch of latent frames with a caller-held Mimi state, so + /// successive calls continue one stream. fn decode_frames( &mut self, latents: &[f32], @@ -958,28 +971,6 @@ impl AprilPocketTts { } } -fn split_model_at_natural_boundaries( - text: &str, - max_tokens: usize, - token_count: F, -) -> Result, String> -where - F: FnMut(&str) -> Result, -{ - split_at_natural_boundaries(text, max_tokens, false, token_count) -} - -fn split_playback_at_natural_boundaries( - text: &str, - max_tokens: usize, - token_count: F, -) -> Result, String> -where - F: FnMut(&str) -> Result, -{ - split_at_natural_boundaries(text, max_tokens, true, token_count) -} - fn split_at_natural_boundaries( text: &str, max_tokens: usize, @@ -1027,11 +1018,9 @@ where if !at_word_end && !at_clause_end { continue; } - // Prepared token counts are monotonic in prefix length, so once a - // candidate overflows the limit no longer candidate can fit. Stop - // scanning instead of tokenizing every remaining boundary: that - // kept this loop superlinear in prompt length, and the cost landed - // before the first chunk reached synthesis. + // Prepared token counts are monotonic in prefix length, so no + // longer candidate can fit after the first overflow. Stopping here + // keeps boundary scanning bounded before synthesis starts. if token_count(&text[start..end])? > max_tokens { break; } @@ -1355,83 +1344,6 @@ mod tests { assert_eq!(shape_len(&[2, 1, 8, 1000, 64]).expect("shape"), 1_024_000); } - /// The two engine splitters must keep OPPOSITE isolation polarity. - /// - /// The guards in `pocket.rs` pin which engine method each public API calls, - /// but they cannot see what the method itself does: pointing - /// `split_playback_prompt` at the model wrapper leaves every call site's - /// source text untouched while first-sentence isolation silently stops - /// happening, so the first playback unit becomes the whole utterance and - /// first audio waits on generating all of it. - #[test] - fn engine_splitters_keep_opposite_isolation_polarity() { - let source = include_str!("pocket_april.rs"); - let production = source - .split_once("\n#[cfg(test)]") - .map_or(source, |(production, _)| production); - - // A method's own code, and nothing else. Ending at the method's own - // closing brace keeps the NEXT method's doc comment out, and stripping - // `//` to end of line keeps prose out: neither can call a splitter, so - // scanning either reports drift in a method that has not changed. - let method_code = |name: &str| -> String { - let (_, body) = production - .split_once(name) - .unwrap_or_else(|| panic!("{name} exists")); - let (body, _) = body - .split_once("\n }\n") - .unwrap_or_else(|| panic!("{name} has a closing brace")); - body.lines() - .map(|line| line.split_once("//").map_or(line, |(code, _)| code)) - .collect::>() - .join("\n") - }; - let model = method_code("fn split_prompt"); - let model = model.as_str(); - let playback = method_code("fn split_playback_prompt"); - let playback = playback.as_str(); - - assert_eq!( - ( - model.matches("split_model_at_natural_boundaries(").count(), - model - .matches("split_playback_at_natural_boundaries(") - .count(), - ), - (1, 0), - "split_prompt must pack sentences: isolating here peels sentence \ - one off every already-packed unit" - ); - assert_eq!( - ( - playback - .matches("split_playback_at_natural_boundaries(") - .count(), - playback - .matches("split_model_at_natural_boundaries(") - .count(), - ), - (1, 0), - "split_playback_prompt must isolate sentence one: packing here \ - makes the first playback unit the whole utterance and delays \ - first audio by the full generation" - ); - - // Calling the isolating splitter is necessary but not sufficient: a - // short circuit before the call can return the whole utterance as one - // unit while leaving the delegated splitter unchanged. Playback must - // delegate unconditionally so sentence one remains the first unit. - for control_flow in ["if ", "match ", "else", "return"] { - assert!( - !playback.contains(control_flow), - "split_playback_prompt must delegate unconditionally, found \ - `{control_flow}`: a branch before the split can return the \ - whole utterance as the first playback unit, delaying first \ - audio by the full generation" - ); - } - } - fn whitespace_token_count(text: &str) -> Result { Ok(text.split_whitespace().count()) } @@ -1439,7 +1351,7 @@ mod tests { #[test] fn playback_split_keeps_first_sentence_separate_then_packs_the_remainder() { let text = "One two. Three four. Five six."; - let chunks = split_playback_at_natural_boundaries(text, 4, whitespace_token_count).unwrap(); + let chunks = split_at_natural_boundaries(text, 4, true, whitespace_token_count).unwrap(); assert_eq!(chunks, ["One two. ", "Three four. Five six."]); assert_eq!(chunks.concat(), text); } @@ -1447,7 +1359,7 @@ mod tests { #[test] fn model_split_packs_multiple_sentences_within_limit() { let text = "One two. Three four. Five six."; - let chunks = split_model_at_natural_boundaries(text, 4, whitespace_token_count).unwrap(); + let chunks = split_at_natural_boundaries(text, 4, false, whitespace_token_count).unwrap(); assert_eq!(chunks, ["One two. Three four. ", "Five six."]); assert_eq!(chunks.concat(), text); } @@ -1455,14 +1367,14 @@ mod tests { #[test] fn playback_then_model_split_does_not_isolate_later_sentences_again() { let text = "Alpha one. Beta two. Gamma three."; - let playback = - split_playback_at_natural_boundaries(text, 50, whitespace_token_count).unwrap(); + let playback = split_at_natural_boundaries(text, 50, true, whitespace_token_count).unwrap(); assert_eq!(playback, ["Alpha one. ", "Beta two. Gamma three."]); let model: Vec<_> = playback .iter() .flat_map(|chunk| { - split_model_at_natural_boundaries(chunk.trim(), 50, whitespace_token_count).unwrap() + split_at_natural_boundaries(chunk.trim(), 50, false, whitespace_token_count) + .unwrap() }) .collect(); assert_eq!(model, ["Alpha one.", "Beta two. Gamma three."]); @@ -1577,11 +1489,18 @@ mod tests { assert_eq!(estimate_max_frames(300, 12.5), 1_275); } - /// Regression (review finding): the voice caches must key on CONTENT. - /// Voice switching clones and drops sample buffers, so a new voice with - /// the same length and rate can land at a recycled address — an - /// address-based key would then restore the previous voice's state and - /// speak with the wrong voice. + #[test] + fn streaming_emit_interval_uses_decoder_chunk_boundaries() { + assert_eq!(decoder_safe_emit_frames(0), DECODER_CHUNK_FRAMES); + assert_eq!(decoder_safe_emit_frames(1), DECODER_CHUNK_FRAMES); + assert_eq!(decoder_safe_emit_frames(12), DECODER_CHUNK_FRAMES); + assert_eq!(decoder_safe_emit_frames(13), DECODER_CHUNK_FRAMES); + assert_eq!(decoder_safe_emit_frames(24), DECODER_CHUNK_FRAMES * 2); + } + + /// Voice caches key on content because switching voices clones and drops + /// sample buffers. An address-based key could restore another voice's + /// state when the allocator reuses a buffer address. #[test] fn voice_key_is_content_based_not_address_based() { let style_a = VoiceStyle { diff --git a/src-tauri/src/commands/pocket_voice.rs b/src-tauri/src/commands/pocket_voice.rs index 85fdea2a2..4a60ee59c 100644 --- a/src-tauri/src/commands/pocket_voice.rs +++ b/src-tauri/src/commands/pocket_voice.rs @@ -1669,6 +1669,9 @@ fn synthesize_and_stream( if !callback_active.load(Ordering::SeqCst) { return false; } + if samples.is_empty() { + return true; + } let delta = match callback_speed_processor.borrow_mut().process(&samples) { Ok(processed) => processed, Err(error) => { From d329e965355738cc1335790781918c0199d38c30 Mon Sep 17 00:00:00 2001 From: John Tennant Date: Wed, 19 Aug 2026 17:51:56 -0400 Subject: [PATCH 03/10] refactor(voice): share the streaming synthesis path --- .../crates/berd-voice/src/pocket_april.rs | 291 +----------------- 1 file changed, 5 insertions(+), 286 deletions(-) diff --git a/src-tauri/crates/berd-voice/src/pocket_april.rs b/src-tauri/crates/berd-voice/src/pocket_april.rs index 324ee0103..1feb5eb82 100644 --- a/src-tauri/crates/berd-voice/src/pocket_april.rs +++ b/src-tauri/crates/berd-voice/src/pocket_april.rs @@ -405,51 +405,11 @@ impl AprilPocketTts { prepared: &AprilPreparedPrompt, style: &VoiceStyle, ) -> Result, String> { - // Optional synthesis phase timing, enabled by BERD_TTS_PHASE_LOG=1. - let phase_log = std::env::var("BERD_TTS_PHASE_LOG").is_ok_and(|v| v == "1"); - let t0 = std::time::Instant::now(); - let mut flow_state = self.conditioned_flow_state(style)?; - let t_condition = t0.elapsed(); - let token_ids = self - .tokenizer - .encode(prepared.text.as_str(), false) - .map_err(|err| format!("tokenize Pocket TTS prompt: {err}"))? - .get_ids() - .iter() - .copied() - .map(i64::from) - .collect::>(); - if token_ids.is_empty() { - return Ok(Vec::new()); - } - if token_ids.len() > self.bundle.max_token_per_chunk { - return Err(format!( - "Pocket TTS prompt has {} tokens; split_text_into_chunks maximum is {}", - token_ids.len(), - self.bundle.max_token_per_chunk - )); - } - - let token_count = token_ids.len(); - let text_embeddings = self.text_embeddings(token_ids)?; - self.run_flow_main_prefix(&text_embeddings, &mut flow_state)?; - let t_prefix = t0.elapsed(); - 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)?; - let t_generate = t0.elapsed(); - let audio = self.decode_latents(&latents)?; - if phase_log { - eprintln!( - "tts-phase: condition={:.0}ms prefix={:.0}ms generate={:.0}ms decode={:.0}ms frames={} audio_s={:.2}", - t_condition.as_secs_f64() * 1e3, - (t_prefix - t_condition).as_secs_f64() * 1e3, - (t_generate - t_prefix).as_secs_f64() * 1e3, - (t0.elapsed() - t_generate).as_secs_f64() * 1e3, - latents.len() / self.bundle.latent_dim, - audio.len() as f64 / self.bundle.sample_rate as f64, - ); - } + let mut audio = Vec::new(); + self.synth_chunk_streaming(prepared, style, DECODER_CHUNK_FRAMES, &mut |delta| { + audio.extend(delta); + true + })?; Ok(audio) } @@ -829,146 +789,6 @@ impl AprilPocketTts { .map_err(ort_error("prime Pocket TTS text state"))?; replace_state_from_outputs(state, &mut outputs) } - - fn generate_latents( - &mut self, - max_frames: usize, - frames_after_eos: usize, - state: &mut [StateValue], - ) -> Result, String> { - 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 { - let sequence = Tensor::from_array(( - vec![1_i64, 1, self.bundle.latent_dim as i64], - current.clone().into_boxed_slice(), - )) - .map_err(ort_error("create latent input"))?; - let text_embeddings = Tensor::::new( - &ort::memory::Allocator::default(), - [1_i64, 0, self.bundle.conditioning_dim as i64], - ) - .map_err(ort_error("create empty text input"))?; - let mut inputs = vec![ - (Cow::Borrowed("sequence"), SessionInputValue::from(sequence)), - ( - Cow::Borrowed("text_embeddings"), - SessionInputValue::from(text_embeddings), - ), - ]; - append_state_inputs(&mut inputs, state); - let mut outputs = self - .flow_main - .run(inputs) - .map_err(ort_error("run Pocket TTS Flow LM"))?; - let conditioning = outputs[0] - .try_extract_tensor::() - .map_err(ort_error("extract Flow LM conditioning"))? - .1 - .to_vec(); - let eos_logit = outputs[1] - .try_extract_tensor::() - .map_err(ort_error("extract Flow LM EOS logit"))? - .1 - .first() - .copied() - .ok_or_else(|| "Flow LM returned empty EOS logit".to_string())?; - replace_state_from_outputs(state, &mut outputs)?; - - if eos_logit > EOS_LOGIT_THRESHOLD && eos_step.is_none() { - eos_step = Some(step); - } - if eos_step.is_some_and(|eos| step >= eos + frames_after_eos) { - break; - } - - let mut noise = - normal_noise(&mut rng, self.bundle.latent_dim, DEFAULT_TEMPERATURE.sqrt()); - let conditioning = Tensor::from_array(( - vec![1_i64, self.bundle.conditioning_dim as i64], - conditioning.into_boxed_slice(), - )) - .map_err(ort_error("create flow conditioning"))?; - let s = Tensor::from_array((vec![1_i64, 1], vec![0.0_f32].into_boxed_slice())) - .map_err(ort_error("create flow start tensor"))?; - let t = Tensor::from_array((vec![1_i64, 1], vec![1.0_f32].into_boxed_slice())) - .map_err(ort_error("create flow end tensor"))?; - let x = Tensor::from_array(( - vec![1_i64, self.bundle.latent_dim as i64], - noise.clone().into_boxed_slice(), - )) - .map_err(ort_error("create flow noise tensor"))?; - let outputs = self - .flow - .run(ort::inputs![ - "c" => conditioning, - "s" => s, - "t" => t, - "x" => x, - ]) - .map_err(ort_error("run Pocket TTS flow"))?; - let flow = outputs[0] - .try_extract_tensor::() - .map_err(ort_error("extract Pocket TTS flow"))? - .1; - if flow.len() != noise.len() { - return Err(format!( - "flow returned {} values; expected {}", - flow.len(), - noise.len() - )); - } - for (sample, delta) in noise.iter_mut().zip(flow) { - *sample += *delta; - } - current.clone_from(&noise); - latents.extend_from_slice(&noise); - } - Ok(latents) - } - - fn decode_latents(&mut self, latents: &[f32]) -> Result, String> { - if latents.is_empty() { - return Ok(Vec::new()); - } - if !latents.len().is_multiple_of(self.bundle.latent_dim) { - return Err(format!( - "latent buffer has {} values, not divisible by {}", - latents.len(), - self.bundle.latent_dim - )); - } - let frame_count = latents.len() / self.bundle.latent_dim; - let mut state = initialize_state(&self.bundle.mimi_state_manifest)?; - let mut audio = Vec::new(); - - for start in (0..frame_count).step_by(DECODER_CHUNK_FRAMES) { - let end = (start + DECODER_CHUNK_FRAMES).min(frame_count); - let values = - latents[start * self.bundle.latent_dim..end * self.bundle.latent_dim].to_vec(); - let latent = Tensor::from_array(( - vec![1_i64, (end - start) as i64, self.bundle.latent_dim as i64], - values.into_boxed_slice(), - )) - .map_err(ort_error("create Mimi latent tensor"))?; - let mut inputs = vec![(Cow::Borrowed("latent"), SessionInputValue::from(latent))]; - append_state_inputs(&mut inputs, &state); - let mut outputs = self - .mimi_decoder - .run(inputs) - .map_err(ort_error("run Mimi decoder"))?; - let samples = outputs[0] - .try_extract_tensor::() - .map_err(ort_error("extract Mimi audio"))? - .1; - audio.extend_from_slice(samples); - replace_state_from_outputs(&mut state, &mut outputs)?; - } - Ok(audio) - } } fn split_at_natural_boundaries( @@ -1623,107 +1443,6 @@ mod tests { }) } - #[test] - #[ignore = "requires BERD_POCKET_TEST_MODEL_DIR"] - fn incremental_stateful_decode_matches_batch_decode() { - let dir = std::env::var("BERD_POCKET_TEST_MODEL_DIR") - .expect("set BERD_POCKET_TEST_MODEL_DIR to the verified April bundle"); - let mut engine = AprilPocketTts::load(Path::new(&dir), 1).expect("load April bundle"); - let style = crate::pocket::load_voice_style(&Path::new(&dir).join("reference_sample.wav")) - .expect("load reference voice"); - - // Generate one real latent sequence (the RNG makes repeat synths - // differ, so both decode paths must consume the SAME latents). - let prepared = - prepare_april_prompt("The relay deploy finished and every check passed cleanly.") - .expect("prepare prompt"); - let mut flow_state = engine - .conditioned_flow_state(&style) - .expect("condition voice"); - let token_ids = engine - .tokenizer - .encode(prepared.text.as_str(), false) - .expect("tokenize") - .get_ids() - .iter() - .copied() - .map(i64::from) - .collect::>(); - let token_count = token_ids.len(); - let text_embeddings = engine.text_embeddings(token_ids).expect("text embeddings"); - engine - .run_flow_main_prefix(&text_embeddings, &mut flow_state) - .expect("prefix"); - let max_frames = estimate_max_frames(token_count, engine.bundle.frame_rate); - let latents = engine - .generate_latents(max_frames, prepared.frames_after_eos, &mut flow_state) - .expect("generate latents"); - let frame_count = latents.len() / engine.bundle.latent_dim; - assert!( - frame_count > DECODER_CHUNK_FRAMES, - "need a multi-chunk case" - ); - - // Batch: the production decode (fresh state, 12-frame steps). - let batch = engine.decode_latents(&latents).expect("batch decode"); - - // Incremental chunkings: 12-frame deltas through one carried Mimi - // state must be bit-exact (the production batch path itself steps by - // DECODER_CHUNK_FRAMES=12 through one state). Sub-12 chunkings are - // measured for the record but are NOT exact — the decoder has - // intra-chunk lookahead — so streaming must emit at >= 12 frames. - for delta_frames in [6usize, 4, 2, 1] { - let mut state = - initialize_state(&engine.bundle.mimi_state_manifest).expect("mimi state"); - let mut streamed = Vec::new(); - for chunk in latents.chunks(delta_frames * engine.bundle.latent_dim) { - streamed.extend( - engine - .decode_frames(chunk, &mut state) - .expect("delta decode"), - ); - } - assert_eq!(batch.len(), streamed.len(), "sample count must match"); - let max_diff = batch - .iter() - .zip(&streamed) - .map(|(a, b)| (a - b).abs()) - .fold(0.0f32, f32::max); - let rms_batch = (batch.iter().map(|s| s * s).sum::() / batch.len() as f32).sqrt(); - let rms_err = (batch - .iter() - .zip(&streamed) - .map(|(a, b)| (a - b) * (a - b)) - .sum::() - / batch.len() as f32) - .sqrt(); - eprintln!( - "delta_frames={delta_frames}: max|diff|={max_diff:.6} rms_err={rms_err:.6} snr_db={:.1}", - 20.0 * (rms_batch / rms_err.max(1e-12)).log10() - ); - } - let mut state = initialize_state(&engine.bundle.mimi_state_manifest).expect("mimi state"); - let mut streamed = Vec::new(); - for chunk in latents.chunks(DECODER_CHUNK_FRAMES * engine.bundle.latent_dim) { - streamed.extend( - engine - .decode_frames(chunk, &mut state) - .expect("delta decode"), - ); - } - - assert_eq!(batch.len(), streamed.len(), "sample count must match"); - let max_diff = batch - .iter() - .zip(&streamed) - .map(|(a, b)| (a - b).abs()) - .fold(0.0f32, f32::max); - assert!( - max_diff <= 1.0e-4, - "incremental decode diverged from batch decode: max |diff| = {max_diff}" - ); - } - #[test] #[ignore = "requires BERD_POCKET_TEST_MODEL_DIR"] fn tokenizer_matches_sentencepiece_reference_including_unknown_words() { From 0f7dbd03c106d029d23333cf295d448789503866 Mon Sep 17 00:00:00 2001 From: John Tennant Date: Wed, 19 Aug 2026 22:16:03 -0400 Subject: [PATCH 04/10] feat(voice): stream utterances through Pocket playback --- src-tauri/crates/berd-voice/src/lib.rs | 4 +- src-tauri/crates/berd-voice/src/pocket.rs | 36 +- .../crates/berd-voice/src/pocket_april.rs | 208 ++++++++- src-tauri/src/commands/pocket_voice.rs | 412 +++++++++++++++++- src-tauri/src/lib.rs | 3 + 5 files changed, 647 insertions(+), 16 deletions(-) diff --git a/src-tauri/crates/berd-voice/src/lib.rs b/src-tauri/crates/berd-voice/src/lib.rs index db34aec2c..88ddda099 100644 --- a/src-tauri/crates/berd-voice/src/lib.rs +++ b/src-tauri/crates/berd-voice/src/lib.rs @@ -4,6 +4,6 @@ mod pocket; pub use pocket::{ april_model_info, load_text_to_speech, load_voice_style, PocketModelArtifact, PocketModelInfo, - PocketTts, VoiceStyle, APRIL_BUNDLE_ID, APRIL_MODEL_ID, APRIL_MODEL_REVISION, DEFAULT_VOICE, - SAMPLE_RATE, VOICE_FILE_EXT, + PocketTts, StreamingTextChunks, VoiceStyle, APRIL_BUNDLE_ID, APRIL_MODEL_ID, + APRIL_MODEL_REVISION, DEFAULT_VOICE, SAMPLE_RATE, VOICE_FILE_EXT, }; diff --git a/src-tauri/crates/berd-voice/src/pocket.rs b/src-tauri/crates/berd-voice/src/pocket.rs index 1886549b8..c68b2ce0a 100644 --- a/src-tauri/crates/berd-voice/src/pocket.rs +++ b/src-tauri/crates/berd-voice/src/pocket.rs @@ -111,6 +111,14 @@ pub struct PocketTts { inner: Mutex, } +/// Stable synthesis units drained from a growing assistant response. +#[derive(Debug, PartialEq, Eq)] +pub struct StreamingTextChunks { + pub ready: Vec, + pub pending: String, + pub first_chunk_pending: bool, +} + /// Load Berd's pinned April INT8 model. pub fn load_text_to_speech(model_dir: &str) -> Result { let dir = Path::new(model_dir); @@ -133,6 +141,32 @@ impl PocketTts { .split_prompt(&prepared) } + /// Drain model-safe units from text that may still be growing. + /// + /// The first complete sentence is made ready immediately. Later text stays + /// pending until it overflows the model's exact token limit, at which point + /// every stable natural chunk except the growing tail is returned. `flush` + /// makes the tail ready at a response or tool boundary. + pub fn take_streaming_text_chunks( + &self, + text: &str, + first_chunk_pending: bool, + flush: bool, + ) -> Result { + self.reject_reentry()?; + let mut engine = self + .inner + .lock() + .map_err(|_| "Pocket TTS engine lock poisoned".to_string())?; + let (ready, pending, first_chunk_pending) = + engine.take_streaming_text_chunks(text, first_chunk_pending, flush)?; + Ok(StreamingTextChunks { + ready, + pending, + first_chunk_pending, + }) + } + /// Synthesize text with the supplied reference voice. /// /// Pocket detects language from text and this model uses one synthesis @@ -183,7 +217,7 @@ impl PocketTts { .inner .lock() .map_err(|_| "Pocket TTS engine lock poisoned".to_string())?; - let chunks = engine.split_playback_prompt(&prepared)?; + let chunks = engine.split_prompt(&prepared)?; for chunk in chunks { if !callback_allows_audio(on_audio, Vec::new())? { return Ok(false); diff --git a/src-tauri/crates/berd-voice/src/pocket_april.rs b/src-tauri/crates/berd-voice/src/pocket_april.rs index 1feb5eb82..8beaf6aa8 100644 --- a/src-tauri/crates/berd-voice/src/pocket_april.rs +++ b/src-tauri/crates/berd-voice/src/pocket_april.rs @@ -388,15 +388,23 @@ impl AprilPocketTts { ) } - pub(crate) fn split_playback_prompt( - &self, - prepared: &AprilPreparedPrompt, - ) -> Result, String> { - split_at_natural_boundaries( - &prepared.text, + pub(crate) fn take_streaming_text_chunks( + &mut self, + text: &str, + first_chunk_pending: bool, + flush: bool, + ) -> Result<(Vec, String, bool), String> { + take_streaming_chunks_at_natural_boundaries( + text, self.bundle.max_token_per_chunk, - true, - |text| self.prepared_token_count(text), + first_chunk_pending, + flush, + |candidate| { + let Some(prepared) = prepare_april_prompt(candidate) else { + return Ok(0); + }; + self.prepared_token_count(&prepared.text) + }, ) } @@ -903,6 +911,97 @@ where Ok(chunks) } +fn take_streaming_chunks_at_natural_boundaries( + text: &str, + max_tokens: usize, + mut first_chunk_pending: bool, + flush: bool, + mut token_count: F, +) -> Result<(Vec, String, bool), String> +where + F: FnMut(&str) -> Result, +{ + let mut pending = text.trim_start().to_string(); + let mut ready = Vec::new(); + + if pending.is_empty() { + return Ok((ready, pending, first_chunk_pending)); + } + + if first_chunk_pending { + let first_sentence_end = pending.char_indices().find_map(|(offset, ch)| { + let end = offset + ch.len_utf8(); + let at_boundary = end == pending.len() + || pending[end..] + .chars() + .next() + .is_some_and(char::is_whitespace); + (at_boundary && natural_boundary(&pending[..end], false) == TextBoundary::Sentence) + .then_some(end) + }); + + if let Some(mut end) = first_sentence_end { + while pending[end..] + .chars() + .next() + .is_some_and(char::is_whitespace) + { + end += pending[end..] + .chars() + .next() + .expect("checked above") + .len_utf8(); + } + let sentence = pending[..end].to_string(); + if token_count(&sentence)? <= max_tokens { + ready.push(sentence); + } else { + ready.extend(split_at_natural_boundaries( + &sentence, + max_tokens, + false, + &mut token_count, + )?); + } + pending = pending[end..].to_string(); + first_chunk_pending = false; + } + } + + if pending.is_empty() { + return Ok((ready, pending, first_chunk_pending)); + } + + if flush { + ready.extend(split_at_natural_boundaries( + &pending, + max_tokens, + first_chunk_pending, + &mut token_count, + )?); + pending.clear(); + first_chunk_pending = false; + return Ok((ready, pending, first_chunk_pending)); + } + + if token_count(&pending)? > max_tokens { + let chunks = split_at_natural_boundaries( + &pending, + max_tokens, + first_chunk_pending, + &mut token_count, + )?; + if chunks.len() > 1 { + let stable_count = chunks.len() - 1; + ready.extend(chunks[..stable_count].iter().cloned()); + pending = chunks[stable_count].clone(); + first_chunk_pending = false; + } + } + + Ok((ready, pending, first_chunk_pending)) +} + fn natural_boundary(candidate: &str, is_end_of_text: bool) -> TextBoundary { if is_end_of_text { return TextBoundary::Sentence; @@ -934,11 +1033,16 @@ fn looks_like_abbreviation(candidate: &str) -> bool { let last_word = candidate .rsplit_once(char::is_whitespace) .map_or(candidate, |(_, word)| word); + let dotted = last_word.strip_suffix('.'); ABBREVIATIONS.contains(&last_word) - || (last_word.ends_with('.') - && last_word[..last_word.len() - 1] - .chars() - .all(|ch| ch.is_ascii_digit())) + || dotted.is_some_and(|stem| { + stem.chars().all(|ch| ch.is_ascii_digit()) + || (stem.chars().count() == 1 && stem.chars().all(|ch| ch.is_alphabetic())) + || (stem.contains('.') + && stem.split('.').all(|part| { + part.chars().count() == 1 && part.chars().all(char::is_alphabetic) + })) + }) } fn load_session(path: PathBuf, num_threads: usize) -> Result { @@ -1200,6 +1304,86 @@ mod tests { assert_eq!(model, ["Alpha one.", "Beta two. Gamma three."]); } + #[test] + fn streaming_text_releases_the_first_sentence_immediately() { + let (ready, pending, first_pending) = take_streaming_chunks_at_natural_boundaries( + "One two. Three four", + 50, + true, + false, + whitespace_token_count, + ) + .unwrap(); + + assert_eq!(ready, ["One two. "]); + assert_eq!(pending, "Three four"); + assert!(!first_pending); + } + + #[test] + fn streaming_text_keeps_later_sentences_pending_under_the_token_limit() { + let (ready, pending, first_pending) = take_streaming_chunks_at_natural_boundaries( + "Three four. Five six.", + 5, + false, + false, + whitespace_token_count, + ) + .unwrap(); + + assert!(ready.is_empty()); + assert_eq!(pending, "Three four. Five six."); + assert!(!first_pending); + } + + #[test] + fn streaming_text_drains_stable_natural_chunks_after_overflow() { + let (ready, pending, first_pending) = take_streaming_chunks_at_natural_boundaries( + "Three four. Five six. Seven eight.", + 4, + false, + false, + whitespace_token_count, + ) + .unwrap(); + + assert_eq!(ready, ["Three four. Five six. "]); + assert_eq!(pending, "Seven eight."); + assert!(!first_pending); + } + + #[test] + fn streaming_text_flushes_the_growing_tail() { + let (ready, pending, first_pending) = take_streaming_chunks_at_natural_boundaries( + "Three four. Five six.", + 5, + false, + true, + whitespace_token_count, + ) + .unwrap(); + + assert_eq!(ready, ["Three four. Five six."]); + assert!(pending.is_empty()); + assert!(!first_pending); + } + + #[test] + fn streaming_text_waits_for_a_real_first_sentence_boundary() { + let (ready, pending, first_pending) = take_streaming_chunks_at_natural_boundaries( + "Dr. J. Smith is still speaking", + 50, + true, + false, + whitespace_token_count, + ) + .unwrap(); + + assert!(ready.is_empty()); + assert_eq!(pending, "Dr. J. Smith is still speaking"); + assert!(first_pending); + } + #[test] fn natural_split_prefers_preceding_sentence_boundary() { let text = "One two. Three four five six."; diff --git a/src-tauri/src/commands/pocket_voice.rs b/src-tauri/src/commands/pocket_voice.rs index 4a60ee59c..18329fbad 100644 --- a/src-tauri/src/commands/pocket_voice.rs +++ b/src-tauri/src/commands/pocket_voice.rs @@ -7,13 +7,15 @@ use std::io::Read; use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; +#[cfg(target_os = "macos")] +use std::sync::mpsc; use std::sync::{Arc, Mutex, OnceLock}; use std::time::{Duration, Instant, SystemTime}; #[cfg(target_os = "macos")] use berd_voice::SAMPLE_RATE; #[cfg(target_os = "macos")] -use berd_voice::{load_text_to_speech, load_voice_style}; +use berd_voice::{load_text_to_speech, load_voice_style, PocketTts, VoiceStyle}; use futures_util::StreamExt; #[cfg(target_os = "macos")] use rodio::buffer::SamplesBuffer; @@ -31,6 +33,8 @@ use tokio::io::AsyncWriteExt; const CACHE_VERSION: &str = "native-voice-v2"; const VERIFIED_MARKER: &str = ".verified"; const POCKET_EVENT: &str = "pocket-voice:event"; +#[cfg(target_os = "macos")] +const POCKET_STREAM_EVENT: &str = "pocket-voice:stream-event"; const DEFAULT_VOICE: &str = "mary"; const DOWNLOAD_PROGRESS_EMIT_INTERVAL: Duration = Duration::from_millis(100); const DOWNLOAD_CONNECT_TIMEOUT: Duration = Duration::from_secs(15); @@ -89,6 +93,7 @@ const MODEL_ARTIFACTS: &[Artifact] = &[ Artifact { filename: "LICENSE", size: 18_655, sha256: "fe7b4ce83b8381cc5b216bbb4af73c570688d1b819c73bbaed8ca401f4677cd6", url: "https://huggingface.co/KevinAHM/pocket-tts-onnx/resolve/58a6d00cf13d239b6748cb0769f35c580a8f606c/onnx/LICENSE" }, ]; +#[cfg(target_os = "macos")] #[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq)] #[serde(rename_all = "camelCase")] pub struct PocketVoice { @@ -129,6 +134,41 @@ pub struct PocketVoiceState { #[derive(Debug, Default)] struct PlaybackRuntime { active: Option>, + #[cfg(target_os = "macos")] + stream: Option, +} + +#[cfg(target_os = "macos")] +#[derive(Debug)] +struct ActivePocketStream { + id: String, + sender: mpsc::Sender, +} + +#[cfg(target_os = "macos")] +#[derive(Debug)] +enum PocketStreamCommand { + Append(String), + Finish, + Stop, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +enum PocketStreamEventState { + Started, + Completed, + Interrupted, + Failed, +} + +#[cfg(target_os = "macos")] +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct PocketStreamEvent { + stream_id: String, + state: PocketStreamEventState, + error: Option, } #[derive(Clone, Debug, Default)] @@ -651,6 +691,147 @@ pub async fn speak_pocket_voice( .map_err(|error| format!("Pocket playback task failed: {error}"))? } +#[tauri::command] +pub fn start_pocket_voice_stream( + app: AppHandle, + state: State<'_, PocketVoiceState>, + native_voice: State<'_, NativeVoiceState>, + stream_id: String, +) -> Result<(), String> { + #[cfg(not(target_os = "macos"))] + { + let _ = (app, state, native_voice, stream_id); + return Err("Pocket voice playback is currently supported on macOS only".to_string()); + } + + #[cfg(target_os = "macos")] + { + if stream_id.trim().is_empty() { + return Err("Pocket voice stream id cannot be empty".to_string()); + } + let base = cache_base(&app)?; + if !pocket_installation_valid(&base) { + return Err("Pocket TTS installation is incomplete or corrupt".to_string()); + } + let voice_id = selected_voice(&base); + let voice = VOICES + .iter() + .find(|voice| voice.id == voice_id) + .copied() + .ok_or_else(|| format!("Unknown selected Pocket voice: {voice_id}"))?; + let active = begin_playback(&state, "Pocket voice playback is already active")?; + let speed = playback_speed(&base); + let output_device = selected_output_device(); + let effective_output_device = effective_output_device_name(output_device.as_deref()); + let capture_suppression = output_device_uses_speakers(effective_output_device.as_deref()) + .then(|| { + log::info!("[voice-echo-guard] speaker output detected"); + native_voice.suppress_capture() + }); + let (sender, receiver) = mpsc::channel(); + { + let mut playback = state + .playback + .lock() + .map_err(|_| "Pocket TTS playback state lock was poisoned".to_string())?; + playback.stream = Some(ActivePocketStream { + id: stream_id.clone(), + sender, + }); + } + + let playback = state.playback.clone(); + let playback_active = active.clone(); + tauri::async_runtime::spawn_blocking(move || { + let _capture_suppression = capture_suppression; + let result = run_pocket_voice_stream( + &app, + &stream_id, + &base, + voice, + output_device.as_deref(), + active.clone(), + speed, + receiver, + ); + let (event_state, error) = match result { + Ok(PocketStreamEventState::Completed) => (PocketStreamEventState::Completed, None), + Ok(PocketStreamEventState::Interrupted) => { + (PocketStreamEventState::Interrupted, None) + } + Ok(other) => (other, None), + Err(error) if !active.load(Ordering::SeqCst) => { + log::debug!("Pocket voice stream stopped after error: {error}"); + (PocketStreamEventState::Interrupted, None) + } + Err(error) => (PocketStreamEventState::Failed, Some(error)), + }; + emit_pocket_stream_event(&app, &stream_id, event_state, error); + finish_playback(&playback, &playback_active); + }); + Ok(()) + } +} + +#[tauri::command] +pub fn append_pocket_voice_stream( + state: State<'_, PocketVoiceState>, + stream_id: String, + text: String, +) -> Result<(), String> { + #[cfg(not(target_os = "macos"))] + { + let _ = (state, stream_id, text); + return Err("Pocket voice playback is currently supported on macOS only".to_string()); + } + + #[cfg(target_os = "macos")] + { + if text.is_empty() { + return Ok(()); + } + send_pocket_stream_command(&state, &stream_id, PocketStreamCommand::Append(text)) + } +} + +#[tauri::command] +pub fn finish_pocket_voice_stream( + state: State<'_, PocketVoiceState>, + stream_id: String, +) -> Result<(), String> { + #[cfg(not(target_os = "macos"))] + { + let _ = (state, stream_id); + return Err("Pocket voice playback is currently supported on macOS only".to_string()); + } + + #[cfg(target_os = "macos")] + { + send_pocket_stream_command(&state, &stream_id, PocketStreamCommand::Finish) + } +} + +#[cfg(target_os = "macos")] +fn send_pocket_stream_command( + state: &PocketVoiceState, + stream_id: &str, + command: PocketStreamCommand, +) -> Result<(), String> { + let playback = state + .playback + .lock() + .map_err(|_| "Pocket TTS playback state lock was poisoned".to_string())?; + let stream = playback + .stream + .as_ref() + .filter(|stream| stream.id == stream_id) + .ok_or_else(|| format!("Pocket voice stream is not active: {stream_id}"))?; + stream + .sender + .send(command) + .map_err(|_| format!("Pocket voice stream worker stopped: {stream_id}")) +} + #[tauri::command] pub fn stop_pocket_voice(state: State<'_, PocketVoiceState>) -> Result { stop_pocket_playback(&state) @@ -665,6 +846,10 @@ fn stop_pocket_playback(state: &PocketVoiceState) -> Result { return Ok(false); }; active.store(false, Ordering::SeqCst); + #[cfg(target_os = "macos")] + if let Some(stream) = playback.stream.as_ref() { + let _ = stream.sender.send(PocketStreamCommand::Stop); + } Ok(true) } @@ -926,6 +1111,10 @@ fn finish_playback(playback: &std::sync::Mutex, completed: &Arc .is_some_and(|active| Arc::ptr_eq(active, completed)) { playback.active = None; + #[cfg(target_os = "macos")] + { + playback.stream = None; + } } } } @@ -1599,6 +1788,227 @@ async fn download_artifact( Ok(()) } +#[cfg(target_os = "macos")] +fn emit_pocket_stream_event( + app: &AppHandle, + stream_id: &str, + state: PocketStreamEventState, + error: Option, +) { + let _ = app.emit( + POCKET_STREAM_EVENT, + PocketStreamEvent { + stream_id: stream_id.to_string(), + state, + error, + }, + ); +} + +#[cfg(target_os = "macos")] +#[allow(clippy::too_many_arguments)] +fn run_pocket_voice_stream( + app: &AppHandle, + stream_id: &str, + base: &Path, + voice: PocketVoice, + output_device: Option<&str>, + active: Arc, + speed: f32, + receiver: mpsc::Receiver, +) -> Result { + use std::num::NonZero; + + use rodio::cpal::traits::HostTrait; + + let version = base.join(CACHE_VERSION); + let engine = load_text_to_speech( + version + .to_str() + .ok_or_else(|| "Pocket model path is not valid UTF-8".to_string())?, + )?; + let style = load_voice_style(&version.join("voices").join(voice.filename))?; + let sink = if let Some(name) = output_device { + let host = rodio::cpal::default_host(); + let mut matching = None; + for device in host + .output_devices() + .map_err(|error| format!("enumerate audio outputs: {error}"))? + { + if device + .description() + .ok() + .is_some_and(|description| description.name() == name) + { + matching = Some(device); + break; + } + } + let device = matching.ok_or_else(|| format!("audio output not found: {name}"))?; + rodio::DeviceSinkBuilder::from_device(device) + .map_err(|error| format!("configure audio output {name}: {error}"))? + .open_stream() + .map_err(|error| format!("open audio output {name}: {error}"))? + } else { + rodio::DeviceSinkBuilder::open_default_sink() + .map_err(|error| format!("open default audio output: {error}"))? + }; + let channels = + NonZero::new(1_u16).ok_or_else(|| "Pocket channel count invariant failed".to_string())?; + let rate = NonZero::new(SAMPLE_RATE) + .ok_or_else(|| "Pocket sample rate invariant failed".to_string())?; + let player = Player::connect_new(sink.mixer()); + let mut speed_processor = StreamingSpeedProcessor::new(speed, SAMPLE_RATE)?; + let mut pending = String::new(); + let mut first_chunk_pending = true; + let mut playback_started = false; + + loop { + if !active.load(Ordering::SeqCst) { + player.stop(); + return Ok(PocketStreamEventState::Interrupted); + } + match receiver.recv() { + Ok(PocketStreamCommand::Append(text)) => { + pending.push_str(&text); + if !synthesize_pocket_stream_ready( + app, + stream_id, + &engine, + &style, + &active, + &player, + channels, + rate, + &mut speed_processor, + &mut pending, + &mut first_chunk_pending, + &mut playback_started, + false, + )? { + return Ok(PocketStreamEventState::Interrupted); + } + } + Ok(PocketStreamCommand::Finish) => { + if !synthesize_pocket_stream_ready( + app, + stream_id, + &engine, + &style, + &active, + &player, + channels, + rate, + &mut speed_processor, + &mut pending, + &mut first_chunk_pending, + &mut playback_started, + true, + )? { + return Ok(PocketStreamEventState::Interrupted); + } + let tail = speed_processor.finish()?; + if !tail.is_empty() { + player.append(SamplesBuffer::new(channels, rate, tail)); + if !playback_started { + emit_pocket_stream_event( + app, + stream_id, + PocketStreamEventState::Started, + None, + ); + } + } + while !player.empty() { + if !active.load(Ordering::SeqCst) { + player.stop(); + return Ok(PocketStreamEventState::Interrupted); + } + std::thread::sleep(Duration::from_millis(10)); + } + return Ok(PocketStreamEventState::Completed); + } + Ok(PocketStreamCommand::Stop) | Err(_) => { + active.store(false, Ordering::SeqCst); + player.stop(); + return Ok(PocketStreamEventState::Interrupted); + } + } + } +} + +#[cfg(target_os = "macos")] +#[allow(clippy::too_many_arguments)] +fn synthesize_pocket_stream_ready( + app: &AppHandle, + stream_id: &str, + engine: &PocketTts, + style: &VoiceStyle, + active: &Arc, + player: &Player, + channels: std::num::NonZero, + rate: std::num::NonZero, + speed_processor: &mut StreamingSpeedProcessor, + pending: &mut String, + first_chunk_pending: &mut bool, + playback_started: &mut bool, + flush: bool, +) -> Result { + let split = engine.take_streaming_text_chunks(pending, *first_chunk_pending, flush)?; + *pending = split.pending; + *first_chunk_pending = split.first_chunk_pending; + for text in split.ready { + if !active.load(Ordering::SeqCst) { + player.stop(); + return Ok(false); + } + let mut callback_error = None; + let completed = engine.synth_chunk_streaming( + text.trim(), + style, + STREAMING_EMIT_FRAMES, + &mut |samples| { + if !active.load(Ordering::SeqCst) { + return false; + } + if samples.is_empty() { + return true; + } + let delta = match speed_processor.process(&samples) { + Ok(processed) => processed, + Err(error) => { + callback_error = Some(error); + return false; + } + }; + if delta.is_empty() { + return true; + } + player.append(SamplesBuffer::new(channels, rate, delta)); + if !*playback_started { + *playback_started = true; + emit_pocket_stream_event(app, stream_id, PocketStreamEventState::Started, None); + println!("VOICE_CONVERSATION_PLAYBACK_STARTED"); + if let Err(error) = std::io::stdout().flush() { + callback_error = Some(format!("signal Pocket playback start: {error}")); + return false; + } + } + true + }, + )?; + if let Some(error) = callback_error { + player.stop(); + return Err(error); + } + if !completed { + player.stop(); + return Ok(false); + } + } + Ok(true) +} + #[cfg(target_os = "macos")] fn synthesize_and_stream( base: &Path, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index bcb5ffee2..d8179c60b 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -620,6 +620,9 @@ pub fn run() { commands::pocket_voice::set_pocket_playback_speed, commands::pocket_voice::preview_pocket_voice, commands::pocket_voice::speak_pocket_voice, + commands::pocket_voice::start_pocket_voice_stream, + commands::pocket_voice::append_pocket_voice_stream, + commands::pocket_voice::finish_pocket_voice_stream, commands::pocket_voice::stop_pocket_voice, commands::pocket_voice::remove_voice_model, commands::native_voice::get_native_voice_conversation_status, From 9c4bc69ab02275e66529a7e2e48c4f7e60359f0d Mon Sep 17 00:00:00 2001 From: John Tennant Date: Wed, 19 Aug 2026 22:34:38 -0400 Subject: [PATCH 05/10] fix(voice): keep the voice catalog portable --- src-tauri/src/commands/pocket_voice.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src-tauri/src/commands/pocket_voice.rs b/src-tauri/src/commands/pocket_voice.rs index 18329fbad..6d89c0a33 100644 --- a/src-tauri/src/commands/pocket_voice.rs +++ b/src-tauri/src/commands/pocket_voice.rs @@ -93,7 +93,6 @@ const MODEL_ARTIFACTS: &[Artifact] = &[ Artifact { filename: "LICENSE", size: 18_655, sha256: "fe7b4ce83b8381cc5b216bbb4af73c570688d1b819c73bbaed8ca401f4677cd6", url: "https://huggingface.co/KevinAHM/pocket-tts-onnx/resolve/58a6d00cf13d239b6748cb0769f35c580a8f606c/onnx/LICENSE" }, ]; -#[cfg(target_os = "macos")] #[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq)] #[serde(rename_all = "camelCase")] pub struct PocketVoice { From 9734c62df016a2f91907b6af797652a2fdf53846 Mon Sep 17 00:00:00 2001 From: John Tennant Date: Wed, 19 Aug 2026 22:39:41 -0400 Subject: [PATCH 06/10] fix(voice): keep stream commands lint-clean cross-platform --- src-tauri/src/commands/pocket_voice.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src-tauri/src/commands/pocket_voice.rs b/src-tauri/src/commands/pocket_voice.rs index 6d89c0a33..bad80bb3c 100644 --- a/src-tauri/src/commands/pocket_voice.rs +++ b/src-tauri/src/commands/pocket_voice.rs @@ -93,6 +93,7 @@ const MODEL_ARTIFACTS: &[Artifact] = &[ Artifact { filename: "LICENSE", size: 18_655, sha256: "fe7b4ce83b8381cc5b216bbb4af73c570688d1b819c73bbaed8ca401f4677cd6", url: "https://huggingface.co/KevinAHM/pocket-tts-onnx/resolve/58a6d00cf13d239b6748cb0769f35c580a8f606c/onnx/LICENSE" }, ]; +#[cfg(target_os = "macos")] #[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq)] #[serde(rename_all = "camelCase")] pub struct PocketVoice { @@ -700,7 +701,7 @@ pub fn start_pocket_voice_stream( #[cfg(not(target_os = "macos"))] { let _ = (app, state, native_voice, stream_id); - return Err("Pocket voice playback is currently supported on macOS only".to_string()); + Err("Pocket voice playback is currently supported on macOS only".to_string()) } #[cfg(target_os = "macos")] @@ -781,7 +782,7 @@ pub fn append_pocket_voice_stream( #[cfg(not(target_os = "macos"))] { let _ = (state, stream_id, text); - return Err("Pocket voice playback is currently supported on macOS only".to_string()); + Err("Pocket voice playback is currently supported on macOS only".to_string()) } #[cfg(target_os = "macos")] @@ -801,7 +802,7 @@ pub fn finish_pocket_voice_stream( #[cfg(not(target_os = "macos"))] { let _ = (state, stream_id); - return Err("Pocket voice playback is currently supported on macOS only".to_string()); + Err("Pocket voice playback is currently supported on macOS only".to_string()) } #[cfg(target_os = "macos")] From 7b51dd00138dc34ca1320e91934b227d433540c5 Mon Sep 17 00:00:00 2001 From: John Tennant Date: Wed, 19 Aug 2026 22:43:10 -0400 Subject: [PATCH 07/10] fix(voice): gate only stream events on macOS --- src-tauri/src/commands/pocket_voice.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src-tauri/src/commands/pocket_voice.rs b/src-tauri/src/commands/pocket_voice.rs index bad80bb3c..8e622015b 100644 --- a/src-tauri/src/commands/pocket_voice.rs +++ b/src-tauri/src/commands/pocket_voice.rs @@ -93,7 +93,6 @@ const MODEL_ARTIFACTS: &[Artifact] = &[ Artifact { filename: "LICENSE", size: 18_655, sha256: "fe7b4ce83b8381cc5b216bbb4af73c570688d1b819c73bbaed8ca401f4677cd6", url: "https://huggingface.co/KevinAHM/pocket-tts-onnx/resolve/58a6d00cf13d239b6748cb0769f35c580a8f606c/onnx/LICENSE" }, ]; -#[cfg(target_os = "macos")] #[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq)] #[serde(rename_all = "camelCase")] pub struct PocketVoice { @@ -153,6 +152,7 @@ enum PocketStreamCommand { Stop, } +#[cfg(target_os = "macos")] #[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq)] #[serde(rename_all = "camelCase")] enum PocketStreamEventState { From 14b2b7c5fdace4549b830329a2e05c4189aac3ad Mon Sep 17 00:00:00 2001 From: John Tennant Date: Thu, 20 Aug 2026 07:42:43 -0400 Subject: [PATCH 08/10] Remove unused Pocket TTS compatibility APIs --- src-tauri/crates/berd-voice/src/lib.rs | 4 +- src-tauri/crates/berd-voice/src/pocket.rs | 82 +------------------ .../crates/berd-voice/src/pocket_april.rs | 15 +--- 3 files changed, 3 insertions(+), 98 deletions(-) diff --git a/src-tauri/crates/berd-voice/src/lib.rs b/src-tauri/crates/berd-voice/src/lib.rs index 88ddda099..9f25705b9 100644 --- a/src-tauri/crates/berd-voice/src/lib.rs +++ b/src-tauri/crates/berd-voice/src/lib.rs @@ -3,7 +3,5 @@ mod pocket; pub use pocket::{ - april_model_info, load_text_to_speech, load_voice_style, PocketModelArtifact, PocketModelInfo, - PocketTts, StreamingTextChunks, VoiceStyle, APRIL_BUNDLE_ID, APRIL_MODEL_ID, - APRIL_MODEL_REVISION, DEFAULT_VOICE, SAMPLE_RATE, VOICE_FILE_EXT, + load_text_to_speech, load_voice_style, PocketTts, StreamingTextChunks, VoiceStyle, SAMPLE_RATE, }; diff --git a/src-tauri/crates/berd-voice/src/pocket.rs b/src-tauri/crates/berd-voice/src/pocket.rs index c68b2ce0a..229d81a7a 100644 --- a/src-tauri/crates/berd-voice/src/pocket.rs +++ b/src-tauri/crates/berd-voice/src/pocket.rs @@ -15,21 +15,14 @@ //! cached model files. use std::cell::RefCell; -use std::path::{Path, PathBuf}; +use std::path::Path; use std::sync::Mutex; use sherpa_onnx::Wave; #[path = "pocket_april.rs"] mod pocket_april; -#[path = "pocket_models.rs"] -mod pocket_models; - use pocket_april::{prepare_april_prompt, AprilPocketTts}; -pub use pocket_models::{ - april_model_info, PocketModelArtifact, PocketModelInfo, APRIL_BUNDLE_ID, APRIL_MODEL_ID, - APRIL_MODEL_REVISION, -}; /// Pocket TTS emits 24 kHz mono PCM. pub const SAMPLE_RATE: u32 = 24_000; @@ -128,19 +121,6 @@ pub fn load_text_to_speech(model_dir: &str) -> Result { } impl PocketTts { - /// Split text into model-safe synthesis units that satisfy the bundle's - /// exact 50-token input limit, packing sentences whenever they fit. - pub fn split_text_into_chunks(&self, text: &str) -> Result, String> { - self.reject_reentry()?; - 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) - } - /// Drain model-safe units from text that may still be growing. /// /// The first complete sentence is made ready immediately. Later text stays @@ -167,35 +147,6 @@ impl PocketTts { }) } - /// 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, String> { - self.reject_reentry()?; - let Some(prepared) = prepare_april_prompt(text) else { - return Ok(Vec::new()); - }; - let mut engine = self - .inner - .lock() - .map_err(|_| "Pocket TTS engine lock poisoned".to_string())?; - let chunks = engine.split_prompt(&prepared)?; - let mut samples = Vec::new(); - for chunk in chunks { - let prepared = prepare_april_prompt(&chunk) - .ok_or_else(|| "Pocket TTS prompt chunk became empty".to_string())?; - samples.extend(engine.synth_chunk(&prepared, style)?); - } - Ok(samples) - } - /// Stream synthesis as PCM deltas become decoder-safe. `emit_frames` is /// rounded down to a positive multiple of the Mimi decoder's 12-frame /// chunk size. Concatenated non-empty deltas equal one `synth_chunk` @@ -265,21 +216,6 @@ fn callback_allows_audio( mod tests { use super::*; - #[test] - fn desktop_model_is_april_int8_only() { - let info = april_model_info(); - assert_eq!(info.max_token_per_chunk, 50); - assert_eq!(info.sample_rate, SAMPLE_RATE); - assert!(info - .artifacts - .iter() - .any(|artifact| artifact.filename == "flow_lm_main_int8.onnx")); - assert!(!info - .artifacts - .iter() - .any(|artifact| artifact.filename == "flow_lm_main.onnx")); - } - #[test] fn active_engine_reentry_is_rejected() { let _guard = SynthesisCallGuard::enter(42).expect("first call"); @@ -296,20 +232,4 @@ mod tests { ); } - #[test] - #[ignore = "requires BERD_POCKET_TEST_MODEL_DIR"] - fn production_api_emits_non_silent_april_int8_pcm() { - let dir = std::env::var("BERD_POCKET_TEST_MODEL_DIR") - .expect("set BERD_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 samples = engine - .synth_chunk("Bright birds begin beside the bay.", "en", &style, 1) - .expect("synthesize through the production API"); - - assert!(!samples.is_empty()); - assert!(samples.iter().all(|sample| sample.is_finite())); - assert!(samples.iter().any(|sample| sample.abs() > 1.0e-6)); - } } diff --git a/src-tauri/crates/berd-voice/src/pocket_april.rs b/src-tauri/crates/berd-voice/src/pocket_april.rs index 8beaf6aa8..fc47cbdfb 100644 --- a/src-tauri/crates/berd-voice/src/pocket_april.rs +++ b/src-tauri/crates/berd-voice/src/pocket_april.rs @@ -408,19 +408,6 @@ impl AprilPocketTts { ) } - pub(crate) fn synth_chunk( - &mut self, - prepared: &AprilPreparedPrompt, - style: &VoiceStyle, - ) -> Result, String> { - let mut audio = Vec::new(); - self.synth_chunk_streaming(prepared, style, DECODER_CHUNK_FRAMES, &mut |delta| { - audio.extend(delta); - true - })?; - Ok(audio) - } - /// Return a fresh Flow LM state conditioned on the reference voice, /// restoring a cached snapshot when the same voice samples were /// conditioned before. Keyed by voice content, like `cached_voice` — @@ -473,7 +460,7 @@ impl AprilPocketTts { } if token_ids.len() > self.bundle.max_token_per_chunk { return Err(format!( - "Pocket TTS prompt has {} tokens; split_text_into_chunks maximum is {}", + "Pocket TTS prompt has {} tokens; maximum is {}", token_ids.len(), self.bundle.max_token_per_chunk )); From 03f89c0ffabea52ab940c22dc6abfabe2eaf19fc Mon Sep 17 00:00:00 2001 From: John Tennant Date: Thu, 20 Aug 2026 07:43:40 -0400 Subject: [PATCH 09/10] Fix Pocket model path borrowing --- src-tauri/crates/berd-voice/src/pocket.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src-tauri/crates/berd-voice/src/pocket.rs b/src-tauri/crates/berd-voice/src/pocket.rs index 229d81a7a..61abb1a5d 100644 --- a/src-tauri/crates/berd-voice/src/pocket.rs +++ b/src-tauri/crates/berd-voice/src/pocket.rs @@ -116,7 +116,7 @@ pub struct StreamingTextChunks { pub fn load_text_to_speech(model_dir: &str) -> Result { let dir = Path::new(model_dir); Ok(PocketTts { - inner: Mutex::new(AprilPocketTts::load(&dir, tts_num_threads())?), + inner: Mutex::new(AprilPocketTts::load(dir, tts_num_threads())?), }) } From 6e60999add43249e1f3d514f9a38b8a8dd5dff3c Mon Sep 17 00:00:00 2001 From: John Tennant Date: Thu, 20 Aug 2026 08:05:48 -0400 Subject: [PATCH 10/10] Fix streaming crate formatting on Windows --- src-tauri/crates/berd-voice/src/pocket.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src-tauri/crates/berd-voice/src/pocket.rs b/src-tauri/crates/berd-voice/src/pocket.rs index 61abb1a5d..f0e357904 100644 --- a/src-tauri/crates/berd-voice/src/pocket.rs +++ b/src-tauri/crates/berd-voice/src/pocket.rs @@ -231,5 +231,4 @@ mod tests { "Pocket TTS synthesis callback panicked" ); } - }