perf(huddle): stream Pocket TTS decoder chunks

Signed-off-by: John Tennant <jtennant@squareup.com>
This commit is contained in:
John Tennant
2026-07-29 10:29:04 -04:00
parent 052174a148
commit b2385ee3bc
16 changed files with 931 additions and 472 deletions

View File

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

View File

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

View File

@@ -108,6 +108,11 @@ pub(crate) struct AprilPocketTts {
cached_voice: Option<CachedVoice>,
}
pub(crate) enum AprilSynthesisOutcome {
Complete,
Interrupted,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct AprilPreparedPrompt {
pub(crate) text: String,
@@ -304,11 +309,15 @@ impl AprilPocketTts {
.collect()
}
pub(crate) fn synth_chunk(
pub(crate) fn synth_chunk_streaming<F>(
&mut self,
prepared: &AprilPreparedPrompt,
style: &VoiceStyle,
) -> Result<Vec<f32>, String> {
callback: F,
) -> Result<AprilSynthesisOutcome, String>
where
F: FnMut(&[f32], f32) -> bool,
{
let voice_embeddings = self.voice_embeddings(style)?;
let mut flow_state = self.condition_voice(&voice_embeddings)?;
let token_ids = self
@@ -321,7 +330,7 @@ impl AprilPocketTts {
.map(i64::from)
.collect::<Vec<_>>();
if token_ids.is_empty() {
return Ok(Vec::new());
return Ok(AprilSynthesisOutcome::Complete);
}
if token_ids.len() > self.bundle.max_token_per_chunk {
return Err(format!(
@@ -337,7 +346,7 @@ impl AprilPocketTts {
let max_frames = estimate_max_frames(token_count, self.bundle.frame_rate);
let latents =
self.generate_latents(max_frames, prepared.frames_after_eos, &mut flow_state)?;
self.decode_latents(&latents)
self.decode_latents(&latents, callback)
}
fn prepared_token_count(&self, text: &str) -> Result<usize, String> {
@@ -597,9 +606,16 @@ impl AprilPocketTts {
Ok(latents)
}
fn decode_latents(&mut self, latents: &[f32]) -> Result<Vec<f32>, String> {
fn decode_latents<F>(
&mut self,
latents: &[f32],
mut callback: F,
) -> Result<AprilSynthesisOutcome, String>
where
F: FnMut(&[f32], f32) -> bool,
{
if latents.is_empty() {
return Ok(Vec::new());
return Ok(AprilSynthesisOutcome::Complete);
}
if !latents.len().is_multiple_of(self.bundle.latent_dim) {
return Err(format!(
@@ -610,7 +626,6 @@ 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);
@@ -630,11 +645,14 @@ impl AprilPocketTts {
let samples = outputs[0]
.try_extract_tensor::<f32>()
.map_err(ort_error("extract Mimi audio"))?
.1;
audio.extend_from_slice(samples);
.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(audio)
Ok(AprilSynthesisOutcome::Complete)
}
}

View File

@@ -17,8 +17,9 @@
//! `tts_starting` sentinels prevent TOCTOU races during the ~200ms window.
//! 3. **`agent_pubkeys` has its own inner `Arc<Mutex>`** — the transcription task
//! clones the `Arc` and reads at post time without the outer lock.
//! 4. **Atomics for cross-thread signaling** — `tts_active`, `tts_cancel`,
//! `ptt_active`, `session_generation` are shared with pipeline worker threads.
//! 4. **Atomics for cross-thread signaling** — `tts_active`,
//! `tts_synthesizing`, `tts_cancel`, `ptt_active`, `session_generation` are
//! shared with pipeline worker threads.
//! 5. **Pipeline teardown extracts handles before dropping** — `teardown_huddle`
//! takes `stt_pipeline`/`tts_pipeline` out of the lock, then calls `shutdown()`
//! and drops them outside the lock (thread joins can block ~200ms).

View File

@@ -281,6 +281,7 @@ pub(crate) async fn maybe_start_stt_pipeline(
// the worker thread (~200ms) and must not block under the mutex.
let (
tts_active,
tts_synthesizing,
tts_cancel,
agent_pubkeys_arc,
session_gen,
@@ -312,6 +313,7 @@ pub(crate) async fn maybe_start_stt_pipeline(
};
(
Arc::clone(&hs.tts_active),
Arc::clone(&hs.tts_synthesizing),
Some(Arc::clone(&hs.tts_cancel)),
Arc::clone(&hs.agent_pubkeys),
Arc::clone(&hs.session_generation),
@@ -325,7 +327,13 @@ pub(crate) async fn maybe_start_stt_pipeline(
drop(old_stt);
let constructed = tokio::task::spawn_blocking(move || {
stt::SttPipeline::new(model_dir, tts_active, tts_cancel, ptt_active_for_stt)
stt::SttPipeline::new(
model_dir,
tts_active,
tts_synthesizing,
tts_cancel,
ptt_active_for_stt,
)
})
.await;
let (pipeline, text_rx) = match constructed {
@@ -429,7 +437,7 @@ pub(crate) async fn maybe_start_tts_pipeline(state: &AppState) -> Result<bool, S
// Atomically check preconditions and claim the construction slot.
// The sentinel prevents a second caller from starting construction
// while we're building outside the lock.
let (tts_active, tts_cancel, tts_starting) = {
let (tts_active, tts_synthesizing, tts_cancel, tts_starting) = {
let hs = state.huddle()?;
if hs.tts_pipeline.is_some() {
return Ok(false);
@@ -442,6 +450,7 @@ pub(crate) async fn maybe_start_tts_pipeline(state: &AppState) -> Result<bool, S
}
(
Arc::clone(&hs.tts_active),
Arc::clone(&hs.tts_synthesizing),
Arc::clone(&hs.tts_cancel),
Arc::clone(&hs.tts_starting),
)
@@ -455,6 +464,7 @@ pub(crate) async fn maybe_start_tts_pipeline(state: &AppState) -> Result<bool, S
tts::TtsPipeline::new_with_voice(
model_dir,
tts_active,
tts_synthesizing,
tts_cancel,
&initial_voice,
output_device,

View File

@@ -137,6 +137,7 @@ pub(crate) async fn run_playout_recv_loop(
app_handle: Option<tauri::AppHandle>,
initial_peers: Vec<(u8, String)>,
tts_active: Arc<AtomicBool>,
tts_synthesizing: Arc<AtomicBool>,
tts_cancel: Arc<AtomicBool>,
) {
use rodio::buffer::SamplesBuffer;
@@ -256,7 +257,10 @@ pub(crate) async fn run_playout_recv_loop(
}
// TTS interrupt frame counter — reset on TTS rising edge.
let tts_now = tts_active.load(Ordering::Acquire);
let tts_now = super::tts::is_tts_interruptible(
tts_active.as_ref(),
tts_synthesizing.as_ref(),
);
if tts_now && !tts_was_active {
frame_counts.clear();
last_frame_reset = tokio::time::Instant::now();

View File

@@ -58,9 +58,13 @@ pub(crate) async fn connect_audio_relay(
let keys = state.keys.lock().map_err(|e| e.to_string())?.clone();
// TTS interrupt flags — recv task cancels TTS when remote humans speak.
let (tts_cancel, tts_active) = {
let (tts_cancel, tts_active, tts_synthesizing) = {
let hs = state.huddle()?;
(Arc::clone(&hs.tts_cancel), Arc::clone(&hs.tts_active))
(
Arc::clone(&hs.tts_cancel),
Arc::clone(&hs.tts_active),
Arc::clone(&hs.tts_synthesizing),
)
};
let app_handle = state.app_handle.lock().ok().and_then(|g| g.clone());
@@ -180,6 +184,7 @@ pub(crate) async fn connect_audio_relay(
initial_peers,
tts_cancel,
tts_active,
tts_synthesizing,
output_device_name,
})
.await
@@ -214,6 +219,7 @@ struct AudioRelayPipelineArgs {
initial_peers: Vec<(u8, String)>,
tts_cancel: Arc<AtomicBool>,
tts_active: Arc<AtomicBool>,
tts_synthesizing: Arc<AtomicBool>,
output_device_name: Option<String>,
}
@@ -227,6 +233,7 @@ async fn audio_relay_pipeline(args: AudioRelayPipelineArgs) -> Result<(), String
initial_peers,
tts_cancel,
tts_active,
tts_synthesizing,
output_device_name,
} = args;
@@ -335,6 +342,7 @@ async fn audio_relay_pipeline(args: AudioRelayPipelineArgs) -> Result<(), String
app_handle,
initial_peers,
tts_active,
tts_synthesizing,
tts_cancel,
));

View File

@@ -93,6 +93,12 @@ pub struct HuddleState {
/// Shared with the STT pipeline for barge-in / echo gating.
#[serde(skip)]
pub tts_active: Arc<AtomicBool>,
/// Shared flag: true while Pocket synthesis is in flight.
///
/// Unlike `tts_active`, this does not gate microphone input. It only keeps
/// barge-in armed while a streamed decoder gap has no queued audio.
#[serde(skip)]
pub tts_synthesizing: Arc<AtomicBool>,
/// Shared barge-in cancel flag. Set by STT when it detects speech during TTS.
/// Read by TTS to stop playback. Lives in HuddleState so it survives pipeline
/// restarts — both STT and TTS reference the same flag for the entire huddle.
@@ -172,6 +178,7 @@ impl Clone for HuddleState {
transcription_enabled: self.transcription_enabled,
transcription_user_controlled: self.transcription_user_controlled,
tts_active: Arc::clone(&self.tts_active),
tts_synthesizing: Arc::clone(&self.tts_synthesizing),
tts_cancel: Arc::clone(&self.tts_cancel),
tts_starting: Arc::clone(&self.tts_starting),
stt_starting: Arc::clone(&self.stt_starting),
@@ -201,6 +208,7 @@ impl Default for HuddleState {
transcription_enabled: false,
transcription_user_controlled: false,
tts_active: Arc::new(AtomicBool::new(false)),
tts_synthesizing: Arc::new(AtomicBool::new(false)),
tts_cancel: Arc::new(AtomicBool::new(false)),
tts_starting: Arc::new(AtomicBool::new(false)),
stt_starting: Arc::new(AtomicBool::new(false)),

View File

@@ -61,12 +61,15 @@ pub struct SttPipeline {
impl SttPipeline {
/// Spawn the pipeline thread.
///
/// `tts_active` is a shared flag set by the TTS pipeline while audio is
/// playing. The STT worker uses it to:
/// `tts_active` is set while TTS audio is queued or playing. The STT worker
/// uses it to:
/// - discard accumulated speech (echo prevention / barge-in gating)
/// - apply a 200 ms cooldown after TTS stops before re-enabling STT
/// - apply a short cooldown after TTS stops before re-enabling STT
/// - detect barge-in: speech onset during TTS → set `tts_cancel`
///
/// `tts_synthesizing` keeps barge-in armed during silent gaps between
/// streamed decoder blocks without suppressing microphone input.
///
/// `tts_cancel` (optional) is the TTS pipeline's cancel flag. When the STT
/// worker detects speech onset while TTS is active, it sets this flag to
/// stop playback immediately (barge-in). Pass `None` if TTS is unavailable.
@@ -86,6 +89,7 @@ impl SttPipeline {
pub fn new(
model_dir: PathBuf,
tts_active: Arc<AtomicBool>,
tts_synthesizing: Arc<AtomicBool>,
tts_cancel: Option<Arc<AtomicBool>>,
ptt_active: Option<Arc<AtomicBool>>,
) -> Result<(Self, tokio_mpsc::Receiver<String>), String> {
@@ -96,6 +100,10 @@ impl SttPipeline {
let shutdown_worker = Arc::clone(&shutdown);
let tts_cancel_worker = tts_cancel.as_ref().map(Arc::clone);
let ptt_active_worker = ptt_active.as_ref().map(Arc::clone);
let tts_activity = TtsActivity {
active: tts_active,
synthesizing: tts_synthesizing,
};
let handle = thread::Builder::new()
.name("stt-worker".into())
.spawn(move || {
@@ -104,7 +112,7 @@ impl SttPipeline {
audio_rx,
text_tx,
shutdown_worker,
tts_active,
tts_activity,
tts_cancel_worker,
ptt_active_worker,
)
@@ -201,12 +209,21 @@ const TTS_COOLDOWN: Duration = Duration::from_millis(50);
/// shows it's safe on the minimum-spec target.
const STT_NUM_THREADS: i32 = 1;
fn tts_playback_ended(was_active: bool, is_active: bool, is_synthesizing: bool) -> bool {
was_active && !is_active && !is_synthesizing
}
struct TtsActivity {
active: Arc<AtomicBool>,
synthesizing: Arc<AtomicBool>,
}
fn stt_worker(
model_dir: PathBuf,
audio_rx: Receiver<Vec<u8>>,
text_tx: tokio_mpsc::Sender<String>,
shutdown: Arc<AtomicBool>,
tts_active: Arc<AtomicBool>,
tts_activity: TtsActivity,
tts_cancel: Option<Arc<AtomicBool>>,
ptt_active: Option<Arc<AtomicBool>>,
) {
@@ -291,8 +308,9 @@ fn stt_worker(
}
// Track TTS transitions to set the cooldown timer.
let tts_now = tts_active.load(Ordering::Acquire);
if tts_was_active && !tts_now {
let tts_now = tts_activity.active.load(Ordering::Acquire);
let tts_is_synthesizing = tts_activity.synthesizing.load(Ordering::Acquire);
if tts_playback_ended(tts_was_active, tts_now, tts_is_synthesizing) {
// TTS just stopped — record the timestamp for the cooldown window.
tts_stopped_at = Some(std::time::Instant::now());
}
@@ -345,7 +363,8 @@ fn stt_worker(
&mut barge_in_frames,
&recognizer,
&text_tx,
&tts_active,
&tts_activity.active,
&tts_activity.synthesizing,
tts_cancel.as_deref(),
&mut tts_stopped_at,
ptt_active.as_ref(),
@@ -408,6 +427,7 @@ fn process_16k_samples(
recognizer: &sherpa_onnx::OfflineRecognizer,
text_tx: &tokio_mpsc::Sender<String>,
tts_active: &Arc<AtomicBool>,
tts_synthesizing: &Arc<AtomicBool>,
tts_cancel: Option<&AtomicBool>,
tts_stopped_at: &mut Option<std::time::Instant>,
ptt_active: Option<&Arc<AtomicBool>>,
@@ -432,6 +452,8 @@ fn process_16k_samples(
};
let tts_playing = tts_active.load(Ordering::Acquire);
let tts_interruptible =
super::tts::is_tts_interruptible(tts_active.as_ref(), tts_synthesizing.as_ref());
// While TTS is playing: skip accumulation (echo prevention).
if tts_playing {
@@ -468,6 +490,14 @@ fn process_16k_samples(
continue;
}
// A decoder gap is silent, so microphone accumulation remains enabled,
// but real speech still cancels the in-flight utterance.
if tts_interruptible && is_speech {
if let Some(cancel) = tts_cancel {
cancel.store(true, Ordering::Release);
}
}
// TTS not playing — check cooldown window.
if let Some(stopped) = *tts_stopped_at {
if stopped.elapsed() < TTS_COOLDOWN {
@@ -565,3 +595,18 @@ fn bytes_to_f32(bytes: &[u8]) -> Vec<f32> {
// drain_until_shutdown lives in super (huddle/mod.rs) — shared with tts.rs.
use super::drain_until_shutdown;
#[cfg(test)]
mod tests {
use super::tts_playback_ended;
#[test]
fn decoder_gap_is_not_a_playback_end() {
assert!(!tts_playback_ended(true, false, true));
}
#[test]
fn drained_completed_stream_is_a_playback_end() {
assert!(tts_playback_ended(true, false, false));
}
}

View File

@@ -8,8 +8,8 @@
//! → tts_worker thread (owns 1 Pocket TTS engine + 1 persistent Player)
//! 1. Preprocess text
//! 2. Split into sentences
//! 3. Synthesize each sentence individually → f32 PCM
//! 4. Clamp to full scale + fade out each sentence
//! 3. Synthesize each sentence while decoder blocks become available
//! 4. Queue monotonic PCM deltas while retaining the final fade suffix
//! 5. Append each buffer to the persistent rodio Player (gapless)
//! 6. While audio is draining, keep pulling queued text items and
//! synthesizing ahead — playback of item N overlaps synthesis of
@@ -48,9 +48,15 @@ use std::{
};
use super::pocket::{
load_text_to_speech, load_voice_style, DEFAULT_VOICE, SAMPLE_RATE, VOICE_FILE_EXT,
load_text_to_speech, load_voice_style, SynthesisOutcome, DEFAULT_VOICE, SAMPLE_RATE,
VOICE_FILE_EXT,
};
use super::preprocessing::{preprocess_for_tts, split_sentences};
#[path = "tts_streaming.rs"]
mod streaming;
#[cfg(test)]
use streaming::STREAM_TAIL_SAMPLES;
use streaming::{CumulativePcmDelta, PocketStreamAssembler};
#[path = "tts_voice_transition.rs"]
mod voice_transition;
@@ -75,10 +81,25 @@ const RECV_TIMEOUT: Duration = Duration::from_millis(100);
/// latency: a cancel is noticed within one tick, and rodio's internal
/// `periodic_access` wrapper stops the in-flight source within a further
/// ~5 ms — so playing audio dies ~15 ms after the flag is set, even while
/// the worker is blocked inside `synth_chunk`.
/// the worker is between decoder callbacks.
const MONITOR_TICK: Duration = Duration::from_millis(10);
const AUDIO_PRIME_TIMEOUT: Duration = Duration::from_secs(2);
struct SynthesisActivityGuard<'a>(&'a AtomicBool);
impl<'a> SynthesisActivityGuard<'a> {
fn new(active: &'a AtomicBool) -> Self {
active.store(true, Ordering::Release);
Self(active)
}
}
impl Drop for SynthesisActivityGuard<'_> {
fn drop(&mut self) {
self.0.store(false, Ordering::Release);
}
}
/// Pocket TTS is a one-step consistency model, not diffusion. Kept for API compat.
const SYNTH_STEPS: usize = 1;
@@ -122,7 +143,12 @@ const MAX_CHUNK_CHARS: usize = 200;
/// Injected as a silent buffer between each synthesized sentence chunk.
const INTER_SENTENCE_SILENCE: f32 = 0.1;
type WorkerControlState = (Arc<AtomicBool>, Arc<AtomicBool>, WorkerCancelSignals);
type WorkerControlState = (
Arc<AtomicBool>,
Arc<AtomicBool>,
Arc<AtomicBool>,
WorkerCancelSignals,
);
// ── Public pipeline handle ────────────────────────────────────────────────────
@@ -164,6 +190,7 @@ impl TtsPipeline {
pub fn new_with_voice(
model_dir: PathBuf,
tts_active: Arc<AtomicBool>,
tts_synthesizing: Arc<AtomicBool>,
cancel: Arc<AtomicBool>,
voice: &str,
output_device: Option<String>,
@@ -177,6 +204,7 @@ impl TtsPipeline {
let voice_cancel = Arc::new(AtomicBool::new(false));
let worker_voice_cancel = Arc::clone(&voice_cancel);
let tts_active_worker = Arc::clone(&tts_active);
let tts_synthesizing_worker = Arc::clone(&tts_synthesizing);
let voice = Arc::new(Mutex::new(voice.to_string()));
let voice_worker = Arc::clone(&voice);
let voice_generation = Arc::new(AtomicU64::new(1));
@@ -199,6 +227,7 @@ impl TtsPipeline {
text_rx,
(
tts_active_worker,
tts_synthesizing_worker,
shutdown_worker,
(cancel_worker, worker_voice_cancel),
),
@@ -312,7 +341,7 @@ fn tts_worker(
startup_tx: mpsc::SyncSender<Result<(), String>>,
) {
let (selected_voice, voice_generation, voice_change_ack) = voice_state;
let (tts_active, shutdown, cancel_signals) = control_state;
let (tts_active, tts_synthesizing, shutdown, cancel_signals) = control_state;
let (cancel, voice_cancel) = cancel_signals;
// ── 1. Initialise TTS engine ──────────────────────────────────────────────
let model_dir_str = model_dir.to_string_lossy().to_string();
@@ -434,15 +463,13 @@ fn tts_worker(
// ── 3b. Barge-in monitor thread ───────────────────────────────────────────
//
// The worker loop only observes `cancel` between sentences — while it is
// blocked inside `synth_chunk` (hundreds of ms for a long sentence),
// nothing would silence the audio that is already playing. The monitor
// closes that gap: every MONITOR_TICK it checks the flag and, while set,
// silences the player and releases the mic gate. It does NOT consume the
// flag — the worker still owns that (drain queue, reset lead-in), so the
// monitor keeps re-clearing until the worker catches up, which also
// covers a sentence appended in the race window after the worker's own
// post-synthesis cancel check.
// The decoder callback observes `cancel`, but audio queued by an earlier
// block must stop immediately rather than waiting for the next callback.
// Every MONITOR_TICK the monitor checks the flag and, while set, silences
// the player and releases the mic gate. It does NOT consume the flag — the
// worker still owns that (drain queue, reset lead-in), so the monitor keeps
// re-clearing until the worker catches up. This also covers a block
// appended in the race window after the worker's own cancel check.
//
// `player_ops` closes the converse race (found in review): the monitor
// loads `cancel == true`, is preempted, the worker consumes the cancel
@@ -460,6 +487,7 @@ fn tts_worker(
let cancel = Arc::clone(&cancel);
let voice_cancel = Arc::clone(&voice_cancel);
let tts_active = Arc::clone(&tts_active);
let tts_synthesizing = Arc::clone(&tts_synthesizing);
let stop = Arc::clone(&monitor_stop);
let player_ops = Arc::clone(&player_ops);
thread::Builder::new()
@@ -480,6 +508,14 @@ fn tts_worker(
player.play();
tts_active.store(false, Ordering::Release);
}
} else if tts_synthesizing.load(Ordering::Acquire)
&& tts_active.load(Ordering::Acquire)
{
let _ops = lock_player_ops(&player_ops);
if !cancel.load(Ordering::Acquire) && !voice_cancel.load(Ordering::Acquire)
{
release_tts_active_if_drained(player.empty(), tts_active.as_ref());
}
}
thread::sleep(MONITOR_TICK);
}
@@ -501,39 +537,12 @@ fn tts_worker(
// arm or on item receipt before synthesis begins.
let silence_buf_len = (INTER_SENTENCE_SILENCE * SAMPLE_RATE as f32) as usize;
// `first_append` = "no audio queued since the player last went idle".
// Flipped by `build_sentence_append_buffer` on the first real append; the
// idle branch below uses it to decide when to drop `tts_active` and to
// arm a fresh lead-in cushion for the next utterance.
// Flipped after the first streamed append; the idle branch below uses it
// to decide when to drop `tts_active` and arm a fresh lead-in cushion for
// the next utterance.
let mut first_append = true;
let mut last_route_id = 0;
let mut deferred_text = VecDeque::new();
let append_audio = |prepared: PreparedModelAudio, route_id: u64| {
let _ops = lock_player_ops(&player_ops);
if cancel.load(Ordering::Acquire)
|| voice_cancel.load(Ordering::Acquire)
|| shutdown.load(Ordering::Acquire)
{
let reason = if shutdown.load(Ordering::Acquire) {
"shutdown"
} else if cancel.load(Ordering::Acquire) {
"barge_in"
} else {
"voice_switch"
};
eprintln!(
"buzz-desktop: tts stage=synthesis status=cancelled reason={reason} route_id={route_id}"
);
return false;
}
player.append(SamplesBuffer::new(channels, rate, prepared.buffer));
eprintln!(
"buzz-desktop: tts stage=player status=append_accepted route_id={route_id} chunk_index={} sample_count={}",
prepared.chunk_index, prepared.sample_count
);
// Set this only after append so STT remains open during synthesis.
tts_active.store(true, Ordering::Release);
true
};
loop {
let mut no_current_text = None;
@@ -672,8 +681,7 @@ fn tts_worker(
let mut synthesis_outcome = "completed";
let mut appended_audio = false;
let mut model_unit_index = 0_usize;
'playback_chunks: for chunk in &chunks {
'playback_chunks: for (chunk_index, chunk) in chunks.iter().enumerate() {
let mut no_current_text = None;
if handle_cancel_or_shutdown(
(&cancel, &voice_cancel),
@@ -694,110 +702,143 @@ fn tts_worker(
continue;
}
let model_chunks = match engine.split_text_into_chunks(text) {
Ok(model_chunks) => model_chunks,
Err(_) => {
eprintln!(
"buzz-desktop: tts stage=synthesis status=failed reason=chunking route_id={route_id}"
);
synthesis_outcome = "failed";
break 'playback_chunks;
}
};
if model_chunks.is_empty() {
eprintln!(
"buzz-desktop: tts stage=synthesis status=empty reason=no_chunks route_id={route_id}"
);
continue;
}
let mut playback_audio = PlaybackChunkAudio::new();
for model_chunk in &model_chunks {
let chunk_index = model_unit_index;
model_unit_index += 1;
let mut no_current_text = None;
if handle_cancel_or_shutdown(
(&cancel, &voice_cancel),
&shutdown,
&tts_active,
(&text_rx, &mut deferred_text, &mut no_current_text),
&voice_change_ack,
Some(route_id),
Some((&player, &player_ops)),
) {
first_append = true;
synthesis_outcome = "cancelled";
break 'playback_chunks;
}
let synthesis = engine.synth_chunk(model_chunk, "en", &style, SYNTH_STEPS);
let mut stream = PocketStreamAssembler::default();
let mut cumulative_delta = CumulativePcmDelta::default();
let mut callback_error = None::<String>;
let callback = |samples: &[f32], _progress: f32| {
if cancel.load(Ordering::Acquire)
|| voice_cancel.load(Ordering::Acquire)
|| shutdown.load(Ordering::Acquire)
{
let reason = if shutdown.load(Ordering::Acquire) {
"shutdown"
} else if cancel.load(Ordering::Acquire) {
"barge_in"
} else {
"voice_switch"
};
return false;
}
let delta = match cumulative_delta.next(samples) {
Ok(Some(delta)) => delta,
Ok(None) => return true,
Err(error) => {
callback_error = Some(error);
return false;
}
};
let sample_count = delta.len();
let _ops = lock_player_ops(&player_ops);
let playback_idle = player.empty();
let result = stream.push(delta, playback_idle, |buffer| {
if cancel.load(Ordering::Acquire)
|| voice_cancel.load(Ordering::Acquire)
|| shutdown.load(Ordering::Acquire)
{
return Err("Pocket TTS streaming cancelled".to_string());
}
player.append(SamplesBuffer::new(channels, rate, buffer));
eprintln!(
"buzz-desktop: tts stage=synthesis status=cancelled reason={reason} route_id={route_id}"
"buzz-desktop: tts stage=player status=append_accepted route_id={route_id} chunk_index={chunk_index} sample_count={sample_count}"
);
// The monitor already stopped any queued playback. Discard
// synthesis that completed after cancellation so stale audio
// never reaches the player, while keeping buzz-voice's
// extracted April engine API unchanged.
first_append = true;
synthesis_outcome = "cancelled";
break 'playback_chunks;
}
match synthesis {
Ok(samples) if !samples.is_empty() => {
if let Some(prepared) = playback_audio.push(
samples,
chunk_index,
&mut first_append,
silence_buf_len,
player.empty(),
) {
if !append_audio(prepared, route_id) {
first_append = true;
synthesis_outcome = "cancelled";
break 'playback_chunks;
}
appended_audio = true;
last_route_id = route_id;
}
}
Ok(_) => {
eprintln!(
"buzz-desktop: tts stage=synthesis status=empty route_id={route_id} chunk_index={chunk_index}"
);
}
Err(_) => {
eprintln!(
"buzz-desktop: tts stage=synthesis status=failed reason=inference route_id={route_id} chunk_index={chunk_index}"
);
synthesis_outcome = "failed";
break;
}
}
}
if let Some(prepared) =
playback_audio.finish(&mut first_append, silence_buf_len, player.empty())
{
if !append_audio(prepared, route_id) {
first_append = true;
synthesis_outcome = "cancelled";
break 'playback_chunks;
tts_active.store(true, Ordering::Release);
Ok(())
});
if let Err(error) = result {
callback_error = Some(error);
return false;
}
true
};
let _synthesis_activity = SynthesisActivityGuard::new(tts_synthesizing.as_ref());
let synth_result =
engine.synth_chunk_streaming(text, "en", &style, SYNTH_STEPS, callback);
if stream.queued_samples > 0 {
first_append = false;
appended_audio = true;
last_route_id = route_id;
}
if synthesis_outcome == "failed" {
if let Some(error) = callback_error {
if !cancel.load(Ordering::Acquire)
&& !voice_cancel.load(Ordering::Acquire)
&& !shutdown.load(Ordering::Acquire)
{
eprintln!(
"buzz-desktop: tts stage=player status=failed reason=queue route_id={route_id} error={error}"
);
synthesis_outcome = "failed";
} else {
synthesis_outcome = "cancelled";
}
break 'playback_chunks;
}
match synth_result {
Ok(SynthesisOutcome::Complete(samples)) if !samples.is_empty() => {
if cancel.load(Ordering::Acquire)
|| voice_cancel.load(Ordering::Acquire)
|| shutdown.load(Ordering::Acquire)
{
first_append = true;
synthesis_outcome = "cancelled";
break 'playback_chunks;
}
let _ops = lock_player_ops(&player_ops);
let playback_idle = player.empty();
let finish_result =
stream.finish(&samples, silence_buf_len, playback_idle, |buffer| {
let sample_count = buffer.len();
if cancel.load(Ordering::Acquire)
|| voice_cancel.load(Ordering::Acquire)
|| shutdown.load(Ordering::Acquire)
{
return Err("Pocket TTS streaming cancelled".to_string());
}
player.append(SamplesBuffer::new(channels, rate, buffer));
eprintln!(
"buzz-desktop: tts stage=player status=append_accepted route_id={route_id} chunk_index={chunk_index} sample_count={sample_count}"
);
tts_active.store(true, Ordering::Release);
Ok(())
});
match finish_result {
Ok(()) => {
first_append = false;
appended_audio = true;
last_route_id = route_id;
}
Err(_)
if cancel.load(Ordering::Acquire)
|| voice_cancel.load(Ordering::Acquire)
|| shutdown.load(Ordering::Acquire) =>
{
first_append = true;
synthesis_outcome = "cancelled";
break 'playback_chunks;
}
Err(error) => {
eprintln!(
"buzz-desktop: tts stage=player status=failed reason=queue route_id={route_id} error={error}"
);
synthesis_outcome = "failed";
break 'playback_chunks;
}
}
}
Ok(SynthesisOutcome::Complete(_)) => {
eprintln!(
"buzz-desktop: tts stage=synthesis status=empty route_id={route_id} chunk_index={chunk_index}"
);
}
Ok(SynthesisOutcome::Interrupted) => {
synthesis_outcome = "cancelled";
break 'playback_chunks;
}
Err(_) => {
eprintln!(
"buzz-desktop: tts stage=synthesis status=failed reason=inference route_id={route_id} chunk_index={chunk_index}"
);
synthesis_outcome = "failed";
break 'playback_chunks;
}
}
}
if synthesis_outcome == "completed" && appended_audio {
eprintln!("buzz-desktop: tts stage=synthesis status=completed route_id={route_id}");
@@ -816,6 +857,7 @@ fn tts_worker(
}
finish_voice_change_ack(&voice_change_ack);
tts_synthesizing.store(false, Ordering::Release);
tts_active.store(false, Ordering::Release);
}
@@ -903,8 +945,21 @@ fn lock_player_ops(ops: &Mutex<()>) -> MutexGuard<'_, ()> {
ops.lock().unwrap_or_else(PoisonError::into_inner)
}
fn release_tts_active_if_drained(player_empty: bool, tts_active: &AtomicBool) {
if player_empty {
tts_active.store(false, Ordering::Release);
}
}
pub(super) fn is_tts_interruptible(tts_active: &AtomicBool, tts_synthesizing: &AtomicBool) -> bool {
tts_active.load(Ordering::Acquire) || tts_synthesizing.load(Ordering::Acquire)
}
// ── Tests ─────────────────────────────────────────────────────────────────────
#[cfg(test)]
#[path = "tts_streaming_tests.rs"]
mod streaming_tests;
#[cfg(test)]
#[path = "tts_tests.rs"]
mod tests;

View File

@@ -1,90 +1,4 @@
use super::{FADE_OUT_SAMPLES, SENTENCE_LEAD_IN_SAMPLES};
pub(super) struct PreparedModelAudio {
pub(super) buffer: Vec<f32>,
pub(super) sample_count: usize,
pub(super) chunk_index: usize,
}
/// Holds one synthesized model unit so playback-boundary decoration is based
/// on the first and last unit that actually produced audio.
pub(super) struct PlaybackChunkAudio {
pending: Option<(Vec<f32>, usize)>,
appended: bool,
}
impl PlaybackChunkAudio {
pub(super) fn new() -> Self {
Self {
pending: None,
appended: false,
}
}
pub(super) fn push(
&mut self,
samples: Vec<f32>,
chunk_index: usize,
first_append: &mut bool,
silence_buf_len: usize,
playback_idle: bool,
) -> Option<PreparedModelAudio> {
if samples.is_empty() {
return None;
}
let previous = self.pending.replace((samples, chunk_index))?;
let prepared = prepare_model_audio(
previous,
first_append,
silence_buf_len,
!self.appended || playback_idle,
false,
);
self.appended = true;
Some(prepared)
}
pub(super) fn finish(
&mut self,
first_append: &mut bool,
silence_buf_len: usize,
playback_idle: bool,
) -> Option<PreparedModelAudio> {
let pending = self.pending.take()?;
Some(prepare_model_audio(
pending,
first_append,
silence_buf_len,
!self.appended || playback_idle,
true,
))
}
}
fn prepare_model_audio(
(samples, chunk_index): (Vec<f32>, usize),
first_append: &mut bool,
silence_buf_len: usize,
starts_playback_chunk: bool,
ends_playback_chunk: bool,
) -> PreparedModelAudio {
let sample_count = samples.len();
let mut audio = clamp_to_full_scale(samples);
if ends_playback_chunk {
apply_fade_out(&mut audio);
}
PreparedModelAudio {
buffer: build_sentence_append_buffer(
first_append,
audio,
silence_buf_len,
starts_playback_chunk,
ends_playback_chunk,
),
sample_count,
chunk_index,
}
}
use super::FADE_OUT_SAMPLES;
/// Hard-clamp samples to ±1.0 full scale.
pub(super) fn clamp_to_full_scale(samples: Vec<f32>) -> Vec<f32> {
@@ -100,34 +14,6 @@ pub(super) fn apply_fade_out(samples: &mut [f32]) {
}
}
pub(super) fn build_sentence_append_buffer(
first_append: &mut bool,
audio: Vec<f32>,
silence_buf_len: usize,
starts_playback_chunk: bool,
ends_playback_chunk: bool,
) -> Vec<f32> {
if *first_append {
*first_append = false;
}
let lead_in_len = if starts_playback_chunk {
SENTENCE_LEAD_IN_SAMPLES
} else {
0
};
let trailing_silence_len = if ends_playback_chunk {
silence_buf_len.saturating_sub(SENTENCE_LEAD_IN_SAMPLES)
} else {
0
};
let mut buffer = Vec::with_capacity(lead_in_len + audio.len() + trailing_silence_len);
buffer.extend(std::iter::repeat_n(0.0_f32, lead_in_len));
buffer.extend(audio);
buffer.extend(std::iter::repeat_n(0.0_f32, trailing_silence_len));
buffer
}
pub(super) fn group_sentences_into_chunks(sentences: &[String], max_chars: usize) -> Vec<String> {
let mut chunks: Vec<String> = Vec::new();
for (index, sentence) in sentences.iter().enumerate() {
@@ -160,76 +46,11 @@ mod tests {
use super::*;
#[test]
fn multi_unit_audio_decorates_only_outer_playback_boundaries() {
let mut chunk = PlaybackChunkAudio::new();
let mut first_append = true;
let silence = SENTENCE_LEAD_IN_SAMPLES + 100;
assert!(chunk
.push(vec![0.4; 16], 0, &mut first_append, silence, false)
.is_none());
let first = chunk
.push(vec![0.5; 16], 1, &mut first_append, silence, false)
.expect("first ready model unit");
assert_eq!(first.buffer.len(), SENTENCE_LEAD_IN_SAMPLES + 16);
assert!(first.buffer[..SENTENCE_LEAD_IN_SAMPLES]
.iter()
.all(|sample| *sample == 0.0));
assert_eq!(first.buffer[SENTENCE_LEAD_IN_SAMPLES], 0.4);
let last = chunk
.finish(&mut first_append, silence, false)
.expect("last ready model unit");
assert_eq!(last.buffer.len(), 16 + 100);
assert_eq!(last.buffer.last(), Some(&0.0));
}
#[test]
fn empty_edge_units_do_not_steal_lead_in_or_trailing_boundary() {
let mut chunk = PlaybackChunkAudio::new();
let mut first_append = true;
let silence = SENTENCE_LEAD_IN_SAMPLES + 100;
assert!(chunk
.push(Vec::new(), 0, &mut first_append, silence, false)
.is_none());
assert!(chunk
.push(vec![0.5; 16], 1, &mut first_append, silence, false)
.is_none());
assert!(chunk
.push(Vec::new(), 2, &mut first_append, silence, false)
.is_none());
let only = chunk
.finish(&mut first_append, silence, false)
.expect("only audible model unit");
assert_eq!(only.buffer.len(), SENTENCE_LEAD_IN_SAMPLES + 16 + 100);
assert!(only.buffer[..SENTENCE_LEAD_IN_SAMPLES]
.iter()
.all(|sample| *sample == 0.0));
assert_eq!(only.buffer.last(), Some(&0.0));
}
#[test]
fn playback_underrun_rearms_the_onset_cushion() {
let mut chunk = PlaybackChunkAudio::new();
let mut first_append = true;
let silence = SENTENCE_LEAD_IN_SAMPLES + 100;
assert!(chunk
.push(vec![0.4; 16], 0, &mut first_append, silence, false)
.is_none());
let first = chunk
.push(vec![0.5; 16], 1, &mut first_append, silence, false)
.expect("first model unit");
assert_eq!(first.buffer.len(), SENTENCE_LEAD_IN_SAMPLES + 16);
let after_underrun = chunk
.push(vec![0.6; 16], 2, &mut first_append, silence, true)
.expect("model unit after underrun");
assert_eq!(after_underrun.buffer.len(), SENTENCE_LEAD_IN_SAMPLES + 16);
assert!(after_underrun.buffer[..SENTENCE_LEAD_IN_SAMPLES]
.iter()
.all(|sample| *sample == 0.0));
fn grouped_sentences_keep_existing_latency_policy() {
let sentences = vec!["First.".to_string(), "Second.".to_string()];
assert_eq!(
group_sentences_into_chunks(&sentences, 200),
["First.", "Second."]
);
}
}

View File

@@ -617,10 +617,12 @@ pub async fn preview_pocket_voice(
let voice_name = pocket_voice_reference(&app, std::slice::from_ref(&voice_key))?;
tokio::task::spawn_blocking(move || {
let active = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let synthesizing = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let cancel = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let pipeline = super::tts::TtsPipeline::new_with_voice(
model_dir,
active.clone(),
synthesizing.clone(),
cancel,
&voice_name,
output_device,
@@ -630,8 +632,9 @@ pub async fn preview_pocket_voice(
let mut heard_audio = false;
while started.elapsed() < std::time::Duration::from_secs(30) {
let is_active = active.load(std::sync::atomic::Ordering::Acquire);
let is_synthesizing = synthesizing.load(std::sync::atomic::Ordering::Acquire);
heard_audio |= is_active;
if heard_audio && !is_active {
if voice_preview_finished(heard_audio, is_active, is_synthesizing) {
return Ok(());
}
std::thread::sleep(std::time::Duration::from_millis(25));
@@ -705,6 +708,10 @@ pub async fn delete_pocket_voice(
})
}
fn voice_preview_finished(heard_audio: bool, is_active: bool, is_synthesizing: bool) -> bool {
heard_audio && !is_active && !is_synthesizing
}
#[cfg(test)]
mod tests {
use super::*;
@@ -730,6 +737,12 @@ mod tests {
assert!(!cancel.load(std::sync::atomic::Ordering::Acquire));
}
#[test]
fn voice_preview_waits_across_streaming_gaps() {
assert!(!voice_preview_finished(true, false, true));
assert!(voice_preview_finished(true, false, false));
}
#[test]
fn defaults_are_backwards_compatible_and_use_mary() {
assert_eq!(

View File

@@ -0,0 +1,182 @@
//! Pocket callback-to-playback PCM assembly.
use super::{
apply_fade_out, clamp_to_full_scale, FADE_OUT_SAMPLES, SAMPLE_RATE, SENTENCE_LEAD_IN_SAMPLES,
};
/// Number of samples retained until the next Pocket callback.
///
/// Retaining a short suffix lets us determine which decoder block is final
/// and apply the utterance fade exactly once without delaying the preceding
/// block.
pub(super) const STREAM_TAIL_SAMPLES: usize = FADE_OUT_SAMPLES * 2;
#[derive(Debug, Default)]
pub(super) struct CumulativePcmDelta {
previous_len: usize,
}
impl CumulativePcmDelta {
/// Return only PCM that was not present in the preceding cumulative batch.
///
/// Pocket may repeat an equal-length cumulative buffer while advancing to
/// another model-safe text chunk. A shorter buffer violates the callback
/// contract and must stop playback rather than duplicate or reorder PCM.
pub(super) fn next<'a>(&mut self, cumulative: &'a [f32]) -> Result<Option<&'a [f32]>, String> {
if cumulative.len() < self.previous_len {
return Err(format!(
"Pocket TTS cumulative callback decreased from {} to {} samples",
self.previous_len,
cumulative.len()
));
}
if cumulative.len() == self.previous_len {
return Ok(None);
}
let delta = &cumulative[self.previous_len..];
self.previous_len = cumulative.len();
Ok(Some(delta))
}
}
#[derive(Debug, Default)]
pub(super) struct PocketStreamAssembler {
pending: Vec<f32>,
callback_samples: usize,
pub(super) callback_count: usize,
pub(super) queued_samples: usize,
}
impl PocketStreamAssembler {
/// Copy and enqueue one independent decoder-block delta.
///
/// `enqueue` must only queue/copy the buffer; playback itself is expected
/// to happen asynchronously.
pub(super) fn push<F>(
&mut self,
samples: &[f32],
playback_idle: bool,
mut enqueue: F,
) -> Result<(), String>
where
F: FnMut(Vec<f32>) -> Result<(), String>,
{
if samples.is_empty() {
return Ok(());
}
self.callback_count += 1;
self.callback_samples += samples.len();
self.pending.extend_from_slice(samples);
let emit_end = self.pending.len().saturating_sub(STREAM_TAIL_SAMPLES);
if emit_end == 0 {
return Ok(());
}
let emit_start = if self.queued_samples == 0 {
match leading_audio_start(&self.pending[..emit_end]) {
Some(start) => start,
// Do not mark silent prefix PCM as playback. Keep it until a
// later callback contains a usable onset.
None => return Ok(()),
}
} else {
0
};
let needs_lead_in = self.queued_samples == 0 || playback_idle;
let mut buffer = Vec::with_capacity(
usize::from(needs_lead_in) * SENTENCE_LEAD_IN_SAMPLES
+ emit_end.saturating_sub(emit_start),
);
if needs_lead_in {
buffer.extend(std::iter::repeat_n(0.0_f32, SENTENCE_LEAD_IN_SAMPLES));
}
buffer.extend(
self.pending[emit_start..emit_end]
.iter()
.map(|sample| sample.clamp(-1.0, 1.0)),
);
enqueue(buffer)?;
self.queued_samples += emit_end - emit_start;
self.pending.drain(..emit_end);
Ok(())
}
/// Queue the retained final suffix, or the complete waveform when the
/// synthesis backend produced no progress callbacks.
pub(super) fn finish<F>(
&mut self,
complete_samples: &[f32],
silence_buf_len: usize,
playback_idle: bool,
mut enqueue: F,
) -> Result<(), String>
where
F: FnMut(Vec<f32>) -> Result<(), String>,
{
if self.callback_count == 0 {
self.pending.extend_from_slice(complete_samples);
self.callback_samples = complete_samples.len();
} else if self.callback_samples != complete_samples.len() {
return Err(format!(
"Pocket TTS callback contract mismatch: callbacks yielded {} samples, final audio contained {}",
self.callback_samples,
complete_samples.len()
));
}
let start = if self.queued_samples == 0 {
leading_audio_start(&self.pending).unwrap_or(0)
} else {
0
};
let mut audio = clamp_to_full_scale(self.pending[start..].to_vec());
apply_fade_out(&mut audio);
if !audio.is_empty() {
let trailing_silence_len = silence_buf_len.saturating_sub(SENTENCE_LEAD_IN_SAMPLES);
let needs_lead_in = self.queued_samples == 0 || playback_idle;
let mut buffer = Vec::with_capacity(
usize::from(needs_lead_in) * SENTENCE_LEAD_IN_SAMPLES
+ audio.len()
+ trailing_silence_len,
);
if needs_lead_in {
buffer.extend(std::iter::repeat_n(0.0_f32, SENTENCE_LEAD_IN_SAMPLES));
}
buffer.extend(audio);
buffer.extend(std::iter::repeat_n(0.0_f32, trailing_silence_len));
enqueue(buffer)?;
self.queued_samples += self.pending.len().saturating_sub(start);
}
self.pending.clear();
Ok(())
}
}
/// Find the first sustained audio in a Pocket progress batch.
fn leading_audio_start(samples: &[f32]) -> Option<usize> {
const WINDOW: usize = SAMPLE_RATE as usize / 100;
const RETAIN: usize = SAMPLE_RATE as usize / 20;
const THRESHOLD: f32 = 0.002;
let threshold_squared = THRESHOLD * THRESHOLD;
let mut energy = 0.0_f32;
for (index, sample) in samples.iter().enumerate() {
energy += sample * sample;
if index >= WINDOW {
let expired = samples[index - WINDOW];
energy -= expired * expired;
}
if index + 1 >= WINDOW && energy / WINDOW as f32 >= threshold_squared {
return Some((index + 1 - WINDOW).saturating_sub(RETAIN));
}
}
None
}

View File

@@ -0,0 +1,223 @@
//! Focused tests for Pocket TTS callback streaming and playback queueing.
use super::*;
use std::sync::mpsc;
fn streaming_test_chunk(value: f32, len: usize) -> Vec<f32> {
vec![value; len]
}
#[test]
fn cumulative_pcm_returns_only_new_samples() {
let mut delta = CumulativePcmDelta::default();
let first = vec![0.25; 1000];
let mut second = first.clone();
second.extend(vec![0.5; 800]);
assert_eq!(delta.next(&first).expect("first batch"), Some(&first[..]));
assert_eq!(
delta.next(&second).expect("second batch"),
Some(&second[1000..])
);
}
#[test]
fn cumulative_pcm_ignores_equal_length_repeats_between_model_chunks() {
let mut delta = CumulativePcmDelta::default();
let samples = vec![0.25; 1000];
assert!(delta.next(&samples).expect("first batch").is_some());
assert_eq!(delta.next(&samples).expect("repeated batch"), None);
}
#[test]
fn cumulative_pcm_rejects_decreasing_lengths() {
let mut delta = CumulativePcmDelta::default();
delta.next(&vec![0.25; 1000]).expect("first batch");
let error = delta
.next(&vec![0.25; 800])
.expect_err("decreasing cumulative length must fail");
assert!(error.contains("decreased from 1000 to 800"));
}
#[test]
fn cumulative_multi_chunk_callbacks_reconstruct_final_pcm_once() {
let first = streaming_test_chunk(0.25, 1000);
let second = streaming_test_chunk(0.5, 1000);
let complete = [first.as_slice(), second.as_slice()].concat();
let cumulative = [first.clone(), first.clone(), complete.clone()];
let mut delta = CumulativePcmDelta::default();
let mut assembler = PocketStreamAssembler::default();
let mut queued = Vec::<Vec<f32>>::new();
for callback in cumulative {
if let Some(samples) = delta.next(&callback).expect("valid cumulative callback") {
assembler
.push(samples, queued.is_empty(), |buffer| {
queued.push(buffer);
Ok(())
})
.expect("queue callback delta");
}
}
assembler
.finish(&complete, 2400, queued.is_empty(), |buffer| {
queued.push(buffer);
Ok(())
})
.expect("finish stream");
assert_eq!(assembler.callback_count, 2);
assert_eq!(assembler.queued_samples, complete.len());
let output = queued.concat();
let speech = &output[SENTENCE_LEAD_IN_SAMPLES..SENTENCE_LEAD_IN_SAMPLES + complete.len()];
assert!(speech[..1000].iter().all(|sample| *sample == 0.25));
assert!(speech[1000..1000 + (1000 - FADE_OUT_SAMPLES)]
.iter()
.all(|sample| *sample == 0.5));
}
#[test]
fn pocket_streaming_queues_before_generation_finishes() {
let mut assembler = PocketStreamAssembler::default();
let (sender, receiver) = mpsc::channel();
assembler
.push(&streaming_test_chunk(0.25, 1000), true, |buffer| {
sender
.send(buffer)
.map_err(|error| format!("send test PCM: {error}"))
})
.expect("stream first callback");
let queued = receiver
.try_recv()
.expect("playback queue receives PCM before finish");
assert_eq!(
queued.len(),
SENTENCE_LEAD_IN_SAMPLES + 1000 - STREAM_TAIL_SAMPLES
);
}
#[test]
fn pocket_streaming_preserves_chunk_lead_in_while_playback_is_active() {
let mut assembler = PocketStreamAssembler::default();
let mut queued = Vec::<Vec<f32>>::new();
assembler
.push(&streaming_test_chunk(0.25, 1000), false, |buffer| {
queued.push(buffer);
Ok(())
})
.expect("stream first callback behind active playback");
assert_eq!(queued.len(), 1);
assert!(queued[0][..SENTENCE_LEAD_IN_SAMPLES]
.iter()
.all(|sample| *sample == 0.0));
assert_eq!(queued[0][SENTENCE_LEAD_IN_SAMPLES], 0.25);
}
#[test]
fn pocket_streaming_preserves_quiet_speech_after_leading_silence() {
let mut assembler = PocketStreamAssembler::default();
let mut samples = vec![0.0; 5000];
samples.extend(streaming_test_chunk(0.005, 500));
samples.extend(streaming_test_chunk(0.5, 1000));
let mut queued = Vec::<Vec<f32>>::new();
assembler
.push(&samples, true, |buffer| {
queued.push(buffer);
Ok(())
})
.expect("stream quiet onset");
assembler
.finish(&samples, 2400, queued.is_empty(), |buffer| {
queued.push(buffer);
Ok(())
})
.expect("finish stream");
let output = queued.concat();
let speech = &output[SENTENCE_LEAD_IN_SAMPLES..];
let quiet_start = speech
.iter()
.position(|sample| *sample == 0.005)
.expect("quiet speech onset must be preserved");
assert!(speech[quiet_start..quiet_start + 500]
.iter()
.all(|sample| *sample == 0.005));
assert!(speech[quiet_start + 500..].contains(&0.5));
}
#[test]
fn pocket_streaming_cancellation_stops_queueing() {
let mut assembler = PocketStreamAssembler::default();
let queued = Vec::<Vec<f32>>::new();
let error = assembler
.push(&streaming_test_chunk(0.25, 1000), true, |buffer| {
let _ = buffer;
Err("Pocket TTS streaming cancelled".to_string())
})
.expect_err("cancelled append must stop generation");
assert_eq!(error, "Pocket TTS streaming cancelled");
assert_eq!(assembler.queued_samples, 0);
assert!(queued.is_empty());
}
#[test]
fn pocket_streaming_surfaces_playback_queue_errors() {
let mut assembler = PocketStreamAssembler::default();
let error = assembler
.push(&streaming_test_chunk(0.25, 1000), true, |_| {
Err("audio device disconnected".to_string())
})
.expect_err("queue failure must be returned");
assert_eq!(error, "audio device disconnected");
}
#[test]
fn pocket_streaming_rejects_callback_contract_mismatch() {
let mut assembler = PocketStreamAssembler::default();
assembler
.push(&streaming_test_chunk(0.25, 1000), true, |_| Ok(()))
.expect("callback");
let error = assembler
.finish(&streaming_test_chunk(0.25, 1200), 2400, false, |_| Ok(()))
.expect_err("mismatched callback samples must not be duplicated");
assert!(error.contains("callback contract mismatch"));
}
#[test]
fn pocket_streaming_rearms_lead_in_after_playback_underrun() {
let mut assembler = PocketStreamAssembler::default();
let mut queued = Vec::<Vec<f32>>::new();
assembler
.push(&streaming_test_chunk(0.25, 1000), true, |buffer| {
queued.push(buffer);
Ok(())
})
.expect("queue first decoder block");
assembler
.push(&streaming_test_chunk(0.5, 1000), true, |buffer| {
queued.push(buffer);
Ok(())
})
.expect("queue decoder block after simulated drain");
assert_eq!(queued.len(), 2);
assert!(queued[1][..SENTENCE_LEAD_IN_SAMPLES]
.iter()
.all(|sample| *sample == 0.0));
assert_eq!(queued[1][SENTENCE_LEAD_IN_SAMPLES], 0.25);
}

View File

@@ -541,6 +541,60 @@ fn monitor_tick_noop_without_cancel() {
);
}
#[test]
fn monitor_releases_mic_gate_when_streaming_playback_drains() {
let tts_active = AtomicBool::new(true);
release_tts_active_if_drained(true, &tts_active);
assert!(
!tts_active.load(Ordering::Acquire),
"silent synthesis gaps must not suppress user speech",
);
}
#[test]
fn monitor_keeps_mic_gate_while_streaming_playback_is_queued() {
let tts_active = AtomicBool::new(true);
release_tts_active_if_drained(false, &tts_active);
assert!(tts_active.load(Ordering::Acquire));
}
#[test]
fn streaming_gap_keeps_barge_in_armed_after_releasing_mic_gate() {
let tts_active = AtomicBool::new(true);
let tts_synthesizing = AtomicBool::new(true);
release_tts_active_if_drained(true, &tts_active);
assert!(!tts_active.load(Ordering::Acquire));
assert!(is_tts_interruptible(&tts_active, &tts_synthesizing));
}
#[test]
fn completed_stream_disarms_barge_in_after_playback_drains() {
let tts_active = AtomicBool::new(true);
let tts_synthesizing = AtomicBool::new(false);
release_tts_active_if_drained(true, &tts_active);
assert!(!is_tts_interruptible(&tts_active, &tts_synthesizing));
}
#[test]
fn synthesis_activity_stays_armed_until_final_queueing_finishes() {
let tts_synthesizing = AtomicBool::new(false);
{
let _activity = SynthesisActivityGuard::new(&tts_synthesizing);
assert!(tts_synthesizing.load(Ordering::Acquire));
}
assert!(!tts_synthesizing.load(Ordering::Acquire));
}
/// Stale-branch race (PR #997 review blocker): monitor observes
/// `cancel == true`, then the worker — under `player_ops` — consumes the
/// cancel and appends a fresh post-cancel utterance before the monitor
@@ -788,103 +842,6 @@ fn apply_fade_out_single_sample() {
assert_eq!(samples[0], 1.0);
}
// ── build_sentence_append_buffer tests ───────────────────────────────────
/// REGRESSION: every chunk needs an onset cushion; synthesized chunks
/// can start with speech energy within the first millisecond.
#[test]
fn lead_in_pad_is_present_for_every_sentence_chunk() {
const SENTENCE_AUDIO_LEN: usize = 1000;
const SILENCE_BUF_LEN: usize = 2400; // 100 ms at 24 kHz, like production
const N_SENTENCES: usize = 5;
let mut first = true;
for _ in 0..N_SENTENCES {
let buf = build_sentence_append_buffer(
&mut first,
vec![0.5_f32; SENTENCE_AUDIO_LEN],
SILENCE_BUF_LEN,
true,
true,
);
assert_eq!(buf.len(), SENTENCE_AUDIO_LEN + SILENCE_BUF_LEN);
assert!(
buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0),
"lead-in pad must be pure silence"
);
assert!(
buf[SENTENCE_LEAD_IN_SAMPLES..SENTENCE_LEAD_IN_SAMPLES + SENTENCE_AUDIO_LEN]
.iter()
.all(|&s| s == 0.5),
"sentence audio must immediately follow the lead-in"
);
assert!(
buf[SENTENCE_LEAD_IN_SAMPLES + SENTENCE_AUDIO_LEN..]
.iter()
.all(|&s| s == 0.0),
"trailing gap must be pure silence"
);
}
assert!(!first, "first_append flag must be cleared after first call");
}
/// `first_append` still flips on the first call for `tts_active` gating.
#[test]
fn build_sentence_append_buffer_flips_first_append() {
let mut first = true;
let _ = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400, true, true);
assert!(!first, "first call must flip the flag");
// Subsequent call: still has a per-sentence lead-in, flag stays false.
let buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400, true, true);
assert!(buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0));
assert!(!first);
}
/// Leading silence is exactly the lead-in; no pre-audio gap is double-counted.
#[test]
fn first_sentence_leading_silence_is_exactly_lead_in() {
let mut first = true;
let buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400, true, true);
assert!(buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0));
assert_eq!(buf[SENTENCE_LEAD_IN_SAMPLES], 0.5);
}
/// Tail silence plus the next lead-in preserves the 100 ms sentence gap.
#[test]
fn sentence_gap_budget_is_preserved() {
let mut first = true;
let silence_buf_len = 2400;
let first_buf =
build_sentence_append_buffer(&mut first, vec![0.5; 100], silence_buf_len, true, true);
let second_buf =
build_sentence_append_buffer(&mut first, vec![0.5; 100], silence_buf_len, true, true);
let first_tail = &first_buf[SENTENCE_LEAD_IN_SAMPLES + 100..];
let second_lead = &second_buf[..SENTENCE_LEAD_IN_SAMPLES];
assert_eq!(first_tail.len(), silence_buf_len - SENTENCE_LEAD_IN_SAMPLES);
assert_eq!(second_lead.len(), SENTENCE_LEAD_IN_SAMPLES);
assert_eq!(first_tail.len() + second_lead.len(), silence_buf_len);
}
/// Regression guard: one contiguous rodio source per synthesized sentence.
#[test]
fn sentence_append_buffer_is_one_contiguous_source() {
let mut first = true;
let buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], 2400, true, true);
assert_eq!(buf.len(), 2400 + 100);
assert!(buf[..SENTENCE_LEAD_IN_SAMPLES].iter().all(|&s| s == 0.0));
assert!(
buf[SENTENCE_LEAD_IN_SAMPLES..SENTENCE_LEAD_IN_SAMPLES + 100]
.iter()
.all(|&s| s == 0.5)
);
}
// ── clamp_to_full_scale tests ─────────────────────────────────────────────
/// In-range speech audio passes through bit-exact — no gain is applied.

View File

@@ -5,20 +5,3 @@ use super::*;
fn sentence_lead_in_is_sane() {
assert_eq!(SENTENCE_LEAD_IN_SAMPLES, 480, "20 ms × 24 kHz");
}
/// Model-token splits remain contiguous: only the playback chunk as a whole
/// receives its onset cushion and trailing sentence gap.
#[test]
fn token_split_units_do_not_add_sentence_boundary_padding() {
let mut first = true;
let silence_buf_len = 2400;
let first_unit =
build_sentence_append_buffer(&mut first, vec![0.5; 100], silence_buf_len, true, false);
let last_unit =
build_sentence_append_buffer(&mut first, vec![0.25; 100], silence_buf_len, false, true);
assert_eq!(first_unit.len(), SENTENCE_LEAD_IN_SAMPLES + 100);
assert_eq!(first_unit.last(), Some(&0.5));
assert_eq!(last_unit.first(), Some(&0.25));
assert_eq!(first_unit.len() + last_unit.len(), 200 + silence_buf_len);
}