diff --git a/src-tauri/crates/berd-voice/src/lib.rs b/src-tauri/crates/berd-voice/src/lib.rs index 651f0fd59..9f25705b9 100644 --- a/src-tauri/crates/berd-voice/src/lib.rs +++ b/src-tauri/crates/berd-voice/src/lib.rs @@ -3,5 +3,5 @@ mod pocket; pub use pocket::{ - load_text_to_speech, load_voice_style, PocketTts, SynthesisOutcome, VoiceStyle, SAMPLE_RATE, + 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 c7f8b8207..f0e357904 100644 --- a/src-tauri/crates/berd-voice/src/pocket.rs +++ b/src-tauri/crates/berd-voice/src/pocket.rs @@ -22,7 +22,7 @@ use sherpa_onnx::Wave; #[path = "pocket_april.rs"] mod pocket_april; -use pocket_april::{prepare_april_prompt, AprilPocketTts, AprilSynthesisOutcome}; +use pocket_april::{prepare_april_prompt, AprilPocketTts}; /// Pocket TTS emits 24 kHz mono PCM. pub const SAMPLE_RATE: u32 = 24_000; @@ -48,6 +48,10 @@ impl SynthesisCallGuard { 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 { @@ -61,6 +65,16 @@ impl Drop for SynthesisCallGuard { } } +/// 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() + .and_then(|v| v.parse::().ok()) + .filter(|&n| n >= 1) + .unwrap_or(TTS_NUM_THREADS) +} + /// Loaded reference voice samples and their original sample rate. #[derive(Debug, Clone)] pub struct VoiceStyle { @@ -90,112 +104,111 @@ 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, +/// 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); 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. + /// Drain model-safe units from text that may still be growing. /// - /// 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( + /// 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, + }) + } + + /// 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, style: &VoiceStyle, - mut callback: F, - ) -> Result - where - F: FnMut(&[f32], f32) -> bool, - { + 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(SynthesisOutcome::Complete(Vec::new())); + return Ok(true); }; 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 { + 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())?; 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 + 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 matches!(outcome, AprilSynthesisOutcome::Interrupted) { - return Ok(SynthesisOutcome::Interrupted); + if !completed { + return Ok(false); } } - Ok(SynthesisOutcome::Complete(samples)) + Ok(true) } -} -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 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_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))) +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()) } @@ -203,111 +216,19 @@ where 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), - ] - ); - } - - #[test] - fn equal_length_pre_decoder_callback_can_cancel() { - let mut callback = |samples: &[f32], progress: f32| { - assert!(samples.is_empty()); - assert_eq!(progress, 0.25); - false - }; - let mut samples = Vec::new(); - - assert!(!append_and_callback(&mut samples, &[], &mut callback, 0.25) - .expect("pre-decoder cancellation callback")); - } - - #[test] - fn callback_panic_is_reported_without_unwinding() { - let mut callback = |_: &[f32], _: f32| -> bool { - panic!("callback failure"); - }; - assert_eq!( - callback_allows_progress(&mut callback, &[], 0.0).unwrap_err(), - "Pocket TTS synthesis callback panicked" - ); - } - #[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] - #[ignore = "requires BERD_POCKET_TEST_MODEL_DIR"] - fn production_streaming_callbacks_are_cumulative_across_model_chunks() { - 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(), + fn callback_panic_is_reported_without_unwinding() { + let mut callback = |_: Vec| -> bool { panic!("callback failure") }; + assert_eq!( + callback_allows_audio(&mut callback, Vec::new()).unwrap_err(), + "Pocket TTS synthesis callback panicked" ); - - assert!(saw_equal_repeat); - assert_eq!(reconstructed, samples); - 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..fc47cbdfb 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, @@ -96,13 +101,128 @@ 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, } +/// 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), + 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 +233,9 @@ pub(crate) struct AprilPocketTts { flow: Session, mimi_decoder: Session, cached_voice: Option, -} - -pub(crate) enum AprilSynthesisOutcome { - Complete, - Interrupted, + /// Post-`condition_voice` Flow LM state cached per reference voice. + /// Restoring it avoids repeating conditioning after the first chunk. + cached_conditioning: Option, } #[derive(Debug, Clone, PartialEq)] @@ -251,6 +369,7 @@ impl AprilPocketTts { tokenizer, bos_embedding, cached_voice: None, + cached_conditioning: None, }) } @@ -269,17 +388,64 @@ impl AprilPocketTts { ) } - pub(crate) fn synth_chunk_streaming( + 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, + first_chunk_pending, + flush, + |candidate| { + let Some(prepared) = prepare_april_prompt(candidate) else { + return Ok(0); + }; + self.prepared_token_count(&prepared.text) + }, + ) + } + + /// 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 state = self.condition_voice(&voice_embeddings)?; + self.cached_conditioning = Some(CachedConditioning { + key, + state: snapshot_state(&state)?, + }); + Ok(state) + } + + /// 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, style: &VoiceStyle, - mut callback: F, - ) -> Result - where - F: FnMut(&[f32], f32) -> bool, - { - let voice_embeddings = self.voice_embeddings(style)?; - let mut flow_state = self.condition_voice(&voice_embeddings)?; + 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 .encode(prepared.text.as_str(), false) @@ -290,11 +456,11 @@ 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!( - "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 )); @@ -304,18 +470,168 @@ 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 = 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); + 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 { + 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(), + )) + .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) + } + + /// 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 +650,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 +693,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) @@ -474,166 +784,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], - callback: &mut F, - ) -> Result<(Vec, bool), String> - where - F: FnMut(&[f32], f32) -> bool, - { - let mut current = vec![f32::NAN; self.bundle.latent_dim]; - let mut latents = Vec::with_capacity(max_frames * self.bundle.latent_dim); - let mut eos_step = None; - let mut rng = rand::rng(); - - for step in 0..max_frames { - // Preserve the mobile callback's pre-PCM cancellation point while - // reserving the second half of progress for decoder-block output. - // The empty block becomes an equal-length cumulative callback in - // `PocketTts::synth_chunk_streaming`. - if !callback(&[], step as f32 / max_frames as f32 * 0.5) { - return Ok((latents, true)); - } - let sequence = Tensor::from_array(( - vec![1_i64, 1, self.bundle.latent_dim as i64], - current.clone().into_boxed_slice(), - )) - .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, false)) - } - - fn decode_latents( - &mut self, - latents: &[f32], - mut callback: F, - ) -> Result - where - F: FnMut(&[f32], f32) -> bool, - { - if latents.is_empty() { - return Ok(AprilSynthesisOutcome::Complete); - } - 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)?; - - 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 - .to_vec(); - 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) - } } fn split_at_natural_boundaries( @@ -680,9 +830,15 @@ 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 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; + } word_end = Some(end); match natural_boundary(&text[start..end], end == text.len()) { @@ -742,6 +898,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; @@ -773,11 +1020,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 { @@ -1008,7 +1260,7 @@ mod tests { } #[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(); assert_eq!(chunks, ["One two. ", "Three four. Five six."]); @@ -1023,6 +1275,102 @@ mod tests { 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_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_at_natural_boundaries(chunk.trim(), 50, false, whitespace_token_count) + .unwrap() + }) + .collect(); + 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."; @@ -1086,6 +1434,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 +1480,140 @@ mod tests { assert_eq!(estimate_max_frames(300, 12.5), 1_275); } + #[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 { + 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 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..8e622015b 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, SynthesisOutcome}; +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,11 +33,15 @@ 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); 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, @@ -127,6 +133,42 @@ 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, +} + +#[cfg(target_os = "macos")] +#[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)] @@ -649,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); + 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); + 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); + 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) @@ -663,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) } @@ -924,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; + } } } } @@ -1597,29 +1788,225 @@ 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() - )); +#[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); + } + } } - if samples.len() == *previous_len { - return Ok(None); +} + +#[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); + } } - let delta = &samples[*previous_len..]; - *previous_len = samples.len(); - Ok(Some(delta)) + Ok(true) } #[cfg(target_os = "macos")] @@ -1676,7 +2063,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 +2072,43 @@ 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; + } + if samples.is_empty() { + return true; + } + 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 +2129,7 @@ fn synthesize_and_stream( player.stop(); return Err(error); } - if matches!(outcome, SynthesisOutcome::Interrupted) { + if !completed { player.stop(); return Ok(()); } @@ -1956,32 +2332,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"); 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,