fix(desktop): preserve Pocket token split continuity

Signed-off-by: John Tennant <jtennant@block.xyz>
This commit is contained in:
John Tennant
2026-07-28 20:02:45 -04:00
parent 67fd208e55
commit 2192d6cc52
5 changed files with 14 additions and 438 deletions
-116
View File
@@ -1,116 +0,0 @@
//! Cold-vs-warm latency bench for Pocket TTS.
//!
//! This duplicates the small config-building snippet from `huddle::pocket` so it
//! doesn't depend on changing module visibility for a one-off dev tool.
//! Keep in sync with `huddle::pocket::load_text_to_speech`.
//!
//! Run with the model files in a directory (defaults to /tmp/pocket-tts-bench):
//! cargo run --release --example pocket_bench
//! cargo run --release --example pocket_bench /path/to/pocket-tts
use std::path::PathBuf;
use std::time::Instant;
use sherpa_onnx::{
self, GenerationConfig, OfflineTts, OfflineTtsConfig, OfflineTtsModelConfig,
OfflineTtsPocketModelConfig, Wave,
};
const SAMPLE_RATE: u32 = 24_000;
const TEST_TEXT: &str =
"Hello, this is a test of the new Pocket TTS engine running on sherpa-onnx.";
fn main() {
let model_dir = std::env::args()
.nth(1)
.unwrap_or_else(|| "/tmp/pocket-tts-bench".to_string());
println!("Model dir: {model_dir}");
let dir = PathBuf::from(&model_dir);
let p = |name: &str| dir.join(name).to_string_lossy().into_owned();
let t0 = Instant::now();
let cfg = OfflineTtsConfig {
model: OfflineTtsModelConfig {
pocket: OfflineTtsPocketModelConfig {
lm_main: Some(p("lm_main.int8.onnx")),
lm_flow: Some(p("lm_flow.int8.onnx")),
encoder: Some(p("encoder.onnx")),
decoder: Some(p("decoder.int8.onnx")),
text_conditioner: Some(p("text_conditioner.onnx")),
vocab_json: Some(p("vocab.json")),
token_scores_json: Some(p("token_scores.json")),
voice_embedding_cache_capacity: 16,
},
num_threads: 1,
debug: false,
..Default::default()
},
..Default::default()
};
let engine = OfflineTts::create(&cfg).expect("engine create");
let load_ms = t0.elapsed().as_secs_f32() * 1000.0;
println!("Engine load: {load_ms:.1} ms");
let t0 = Instant::now();
let voice_path = dir.join("reference_sample.wav");
let wave = Wave::read(voice_path.to_str().unwrap()).expect("voice WAV");
let samples = wave.samples().to_vec();
let sr = wave.sample_rate();
let voice_ms = t0.elapsed().as_secs_f32() * 1000.0;
println!("Voice load: {voice_ms:.1} ms");
let gen = || GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 1.0, // production setting (huddle::pocket::SYNTH_SILENCE_SCALE)
reference_audio: Some(samples.clone()),
reference_sample_rate: sr,
..Default::default()
};
let t0 = Instant::now();
let cold = engine
.generate_with_config(TEST_TEXT, &gen(), None::<fn(&[f32], f32) -> bool>)
.expect("cold synth");
let cold_ms = t0.elapsed().as_secs_f32() * 1000.0;
let cold_audio_ms = (cold.samples().len() as f32 / SAMPLE_RATE as f32) * 1000.0;
let cold_rtf_x = cold_audio_ms / cold_ms;
println!(
"Cold synth: {cold_ms:.1} ms → {cold_audio_ms:.1} ms audio → {cold_rtf_x:.2}× realtime"
);
let t0 = Instant::now();
let warm = engine
.generate_with_config(TEST_TEXT, &gen(), None::<fn(&[f32], f32) -> bool>)
.expect("warm synth");
let warm_ms = t0.elapsed().as_secs_f32() * 1000.0;
let warm_audio_ms = (warm.samples().len() as f32 / SAMPLE_RATE as f32) * 1000.0;
let warm_rtf_x = warm_audio_ms / warm_ms;
println!(
"Warm synth: {warm_ms:.1} ms → {warm_audio_ms:.1} ms audio → {warm_rtf_x:.2}× realtime"
);
let out_path = "/tmp/pocket_bench_out.wav";
let ok = sherpa_onnx::write(out_path, warm.samples(), SAMPLE_RATE as i32);
println!(
"Wrote {} ({} samples, ok={ok})",
out_path,
warm.samples().len()
);
let delta_ms = cold_ms - warm_ms;
let delta_pct = (delta_ms / warm_ms) * 100.0;
println!();
println!("Cold/warm delta: {delta_ms:+.1} ms ({delta_pct:+.1}%)");
println!(
"Decision: warmup {}.",
if delta_ms > 200.0 {
"RECOMMENDED — significant cold-call penalty"
} else if delta_ms > 50.0 {
"OPTIONAL — small cold-call penalty"
} else {
"UNNECESSARY — cold and warm essentially equal"
}
);
}
@@ -1,122 +0,0 @@
//! Clipping probe for any fixed playback gain applied after Pocket TTS synth.
//!
//! Synthesises a spread of sentences (short/long, calm/energetic) and reports
//! the raw peak of each, the post-gain peak, and the fraction of samples that
//! would hit a ±1.0 clamp — i.e. how much a fixed gain would flat-top the
//! waveform ("blown out" distortion).
//!
//! History: the production pipeline briefly shipped a fixed 9.3× gain
//! calibrated on a single bench utterance that peaked at 0.076. This probe
//! showed real output peaks at 0.4–0.97, so that gain clipped 13–34% of all
//! samples (the 2026-06-12 "blown out" report). Production now applies no
//! gain — run this probe before reintroducing one.
//!
//! Run with model files in ~/.buzz/models/pocket-tts (override with arg 1):
//! cargo run --release --example pocket_clip_probe
use std::path::PathBuf;
use sherpa_onnx::{
self, GenerationConfig, OfflineTts, OfflineTtsConfig, OfflineTtsModelConfig,
OfflineTtsPocketModelConfig, Wave,
};
/// Candidate gain under test (the regressed production value).
const GAIN: f32 = 9.3;
const PROMPTS: &[&str] = &[
"Hello, this is a test of the new Pocket TTS engine running on sherpa-onnx.",
"Yep, I can hear you.",
"Absolutely! That sounds fantastic, let's do it right now!",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I found three problems in the code: a race condition, a memory leak, and an off-by-one error in the loop bounds.",
"No.",
"Warning! The build failed because seventeen tests crashed unexpectedly!",
"Sure, I can walk you through the whole pipeline step by step whenever you're ready.",
];
fn main() {
let model_dir = std::env::args().nth(1).unwrap_or_else(|| {
dirs::home_dir()
.expect("home dir")
.join(".buzz/models/pocket-tts")
.to_string_lossy()
.into_owned()
});
eprintln!("Model dir: {model_dir}");
let dir = PathBuf::from(&model_dir);
let p = |name: &str| dir.join(name).to_string_lossy().into_owned();
let cfg = OfflineTtsConfig {
model: OfflineTtsModelConfig {
pocket: OfflineTtsPocketModelConfig {
lm_main: Some(p("lm_main.int8.onnx")),
lm_flow: Some(p("lm_flow.int8.onnx")),
encoder: Some(p("encoder.onnx")),
decoder: Some(p("decoder.int8.onnx")),
text_conditioner: Some(p("text_conditioner.onnx")),
vocab_json: Some(p("vocab.json")),
token_scores_json: Some(p("token_scores.json")),
voice_embedding_cache_capacity: 16,
},
num_threads: 1,
debug: false,
..Default::default()
},
..Default::default()
};
let engine = OfflineTts::create(&cfg).expect("engine create");
let voice_path = dir.join("reference_sample.wav");
let wave = Wave::read(voice_path.to_str().unwrap()).expect("voice WAV");
let voice_samples = wave.samples().to_vec();
let voice_sr = wave.sample_rate();
let gen = || GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 1.0,
reference_audio: Some(voice_samples.clone()),
reference_sample_rate: voice_sr,
..Default::default()
};
let _ = engine.generate_with_config("warmup.", &gen(), None::<fn(&[f32], f32) -> bool>);
println!(
"{:<46} | {:>8} | {:>9} | {:>9} | {:>10}",
"prompt", "raw peak", "raw RMS", "post-gain", "% clipped"
);
println!("{}", "-".repeat(95));
let mut worst_clip = 0.0f32;
for prompt in PROMPTS {
let out = engine
.generate_with_config(prompt, &gen(), None::<fn(&[f32], f32) -> bool>)
.expect("synth");
let samples = out.samples();
let peak = samples.iter().fold(0.0f32, |m, s| m.max(s.abs()));
let rms = (samples.iter().map(|s| s * s).sum::<f32>() / samples.len() as f32).sqrt();
let post = peak * GAIN;
let clipped = samples.iter().filter(|s| s.abs() * GAIN > 1.0).count();
let clip_pct = 100.0 * clipped as f32 / samples.len() as f32;
worst_clip = worst_clip.max(clip_pct);
let label: String = prompt.chars().take(44).collect();
println!("{label:<46} | {peak:>8.4} | {rms:>9.4} | {post:>9.3} | {clip_pct:>9.3}%");
}
println!();
println!(
"Verdict: worst-case clipped fraction {worst_clip:.3}% — {}",
if worst_clip > 0.1 {
"AUDIBLE DISTORTION LIKELY (gain too hot)"
} else if worst_clip > 0.0 {
"marginal — occasional transient clipping"
} else {
"no clipping at this gain"
}
);
}
@@ -1,149 +0,0 @@
//! Onset-attenuation probe for Pocket TTS.
//!
//! Synthesises a handful of short sentences and dumps per-sentence onset
//! statistics (samples[0], 1ms/5ms/20ms peak + RMS) so we can decide whether
//! the production `apply_fades` 8 ms fade-in is masking real audio.
//!
//! Also writes the raw (un-faded, un-normalised) audio of each sentence to
//! /tmp so they can be inspected in Audacity / aplay without rodio in the
//! loop.
//!
//! Run with model files in /tmp/pocket-tts-bench (override with arg 1):
//! cargo run --release --example pocket_onset_probe
//! cargo run --release --example pocket_onset_probe /path/to/pocket-tts
use std::path::PathBuf;
use sherpa_onnx::{
self, GenerationConfig, OfflineTts, OfflineTtsConfig, OfflineTtsModelConfig,
OfflineTtsPocketModelConfig, Wave,
};
const SAMPLE_RATE: u32 = 24_000;
/// Test prompts chosen to span different onsets:
/// - palatal glide 'Y' (soft onset)
/// - voiceless fricative 'H' (very soft onset)
/// - labio-velar glide 'W' (medium onset)
/// - voiceless stop 'T' (hard onset)
const PROMPTS: &[&str] = &[
"Yep, I can hear you.",
"Hello there friend.",
"What can I help with?",
"Try this experiment now.",
];
fn main() {
let model_dir = std::env::args()
.nth(1)
.unwrap_or_else(|| "/tmp/pocket-tts-bench".to_string());
eprintln!("Model dir: {model_dir}");
let dir = PathBuf::from(&model_dir);
let p = |name: &str| dir.join(name).to_string_lossy().into_owned();
let cfg = OfflineTtsConfig {
model: OfflineTtsModelConfig {
pocket: OfflineTtsPocketModelConfig {
lm_main: Some(p("lm_main.int8.onnx")),
lm_flow: Some(p("lm_flow.int8.onnx")),
encoder: Some(p("encoder.onnx")),
decoder: Some(p("decoder.int8.onnx")),
text_conditioner: Some(p("text_conditioner.onnx")),
vocab_json: Some(p("vocab.json")),
token_scores_json: Some(p("token_scores.json")),
voice_embedding_cache_capacity: 16,
},
num_threads: 1,
debug: false,
..Default::default()
},
..Default::default()
};
let engine = OfflineTts::create(&cfg).expect("engine create");
let voice_path = dir.join("reference_sample.wav");
let wave = Wave::read(voice_path.to_str().unwrap()).expect("voice WAV");
let voice_samples = wave.samples().to_vec();
let voice_sr = wave.sample_rate();
// Warmup so we're not measuring cold-call jitter.
{
let cfg = GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 1.0, // production setting (huddle::pocket::SYNTH_SILENCE_SCALE)
reference_audio: Some(voice_samples.clone()),
reference_sample_rate: voice_sr,
..Default::default()
};
let _ = engine.generate_with_config("warmup.", &cfg, None::<fn(&[f32], f32) -> bool>);
}
println!(
"{:<28} | {:>10} | {:>10} {:>10} | {:>10} {:>10} | {:>10} {:>10}",
"prompt",
"samples[0]",
"peak@1ms",
"rms@1ms",
"peak@5ms",
"rms@5ms",
"peak@20ms",
"rms@20ms"
);
println!("{}", "-".repeat(120));
for prompt in PROMPTS {
// Mirror the production prompt-prep (capitalise + terminal punctuation).
// These prompts already have it, so this is just to match what
// sherpa-onnx sees in production.
let cfg = GenerationConfig {
speed: 1.05,
num_steps: 1,
silence_scale: 1.0, // production setting (huddle::pocket::SYNTH_SILENCE_SCALE)
reference_audio: Some(voice_samples.clone()),
reference_sample_rate: voice_sr,
..Default::default()
};
let out = engine
.generate_with_config(prompt, &cfg, None::<fn(&[f32], f32) -> bool>)
.expect("synth");
let samples = out.samples();
let n_1ms = (SAMPLE_RATE as f32 * 0.001) as usize;
let n_5ms = (SAMPLE_RATE as f32 * 0.005) as usize;
let n_20ms = (SAMPLE_RATE as f32 * 0.020) as usize;
let stats = |range: &[f32]| -> (f32, f32) {
if range.is_empty() {
return (0.0, 0.0);
}
let peak = range.iter().fold(0.0_f32, |a, &x| a.max(x.abs()));
let sumsq: f32 = range.iter().map(|x| x * x).sum();
let rms = (sumsq / range.len() as f32).sqrt();
(peak, rms)
};
let first = samples.first().copied().unwrap_or(0.0);
let (p1, r1) = stats(&samples[..n_1ms.min(samples.len())]);
let (p5, r5) = stats(&samples[..n_5ms.min(samples.len())]);
let (p20, r20) = stats(&samples[..n_20ms.min(samples.len())]);
println!(
"{:<28} | {:>10.6} | {:>10.6} {:>10.6} | {:>10.6} {:>10.6} | {:>10.6} {:>10.6}",
prompt, first, p1, r1, p5, r5, p20, r20
);
let safe: String = prompt
.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
.collect();
let out_path = format!("/tmp/pocket_onset_{}.wav", &safe[..safe.len().min(24)]);
let _ = sherpa_onnx::write(&out_path, samples, SAMPLE_RATE as i32);
eprintln!(
" → wrote {out_path} ({} samples = {:.3} s)",
samples.len(),
samples.len() as f32 / SAMPLE_RATE as f32
);
}
}
-12
View File
@@ -87,18 +87,6 @@ pub fn load_text_to_speech(model_dir: &str) -> Result<PocketTts, String> {
}
impl PocketTts {
/// Split text into synthesis units that satisfy the bundle's exact
/// 50-token input limit.
pub fn split_text_into_chunks(&self, text: &str) -> Result<Vec<String>, String> {
let Some(prepared) = prepare_april_prompt(text) else {
return Ok(Vec::new());
};
self.inner
.lock()
.map_err(|_| "Pocket TTS engine lock poisoned".to_string())?
.split_prompt(&prepared)
}
/// Synthesize text with the supplied reference voice.
///
/// Pocket detects language from text and this model uses one synthesis
+14 -39
View File
@@ -73,9 +73,8 @@ const SYNTH_STEPS: usize = 1;
///
/// Applied only at the *end* of each synthesised sentence to eliminate the
/// click that would otherwise occur when a non-zero waveform terminates
/// abruptly. **No fade-in is applied** — see `apply_fade_out` for the
/// rationale and `examples/pocket_onset_probe.rs` for the measurement that
/// motivated removing the leading fade.
/// abruptly. **No fade-in is applied** — see `apply_fade_out` for why preserving
/// the leading waveform is important.
const FADE_OUT_SAMPLES: usize = (SAMPLE_RATE as f64 * 0.008) as usize;
/// Length of the zero-sample cushion prepended before each synthesized
@@ -491,29 +490,14 @@ fn tts_worker(
// sentence stays alone (fast time-to-first-audio), the rest pack
// greedily up to MAX_CHUNK_CHARS. Each chunk is one `generate()`
// call; playback of chunk N overlaps synthesis of chunk N+1
// (lookahead pipelining). Grouping matches upstream's ~50-token
// chunking and halves the exposed prosody seams on multi-sentence
// replies — see MAX_CHUNK_CHARS.
// (lookahead pipelining). The Pocket engine applies its exact 50-token
// limit internally; keeping those internal units within one playback
// chunk avoids adding fades and pauses at token-only boundaries.
let sentences: Vec<String> = split_sentences(&text)
.into_iter()
.filter(|s| !s.trim().is_empty())
.collect();
let grouped_chunks = group_sentences_into_chunks(&sentences, MAX_CHUNK_CHARS);
let mut chunks = Vec::new();
let mut split_failed = false;
for chunk in grouped_chunks {
match engine.split_text_into_chunks(&chunk) {
Ok(model_chunks) => chunks.extend(model_chunks),
Err(error) => {
eprintln!("buzz-desktop: TTS chunking failed: {error}");
split_failed = true;
break;
}
}
}
if split_failed {
continue;
}
let chunks = group_sentences_into_chunks(&sentences, MAX_CHUNK_CHARS);
for chunk in &chunks {
if handle_cancel_or_shutdown(
@@ -657,15 +641,10 @@ fn lock_player_ops(ops: &Mutex<()>) -> MutexGuard<'_, ()> {
/// Hard-clamp samples to ±1.0 full scale.
///
/// No gain is applied: Pocket TTS already emits speech-level audio
/// (peaks 0.4–0.97, RMS ≈ −20 dBFS across varied sentences — measured by
/// `examples/pocket_clip_probe`), matching the kyutai reference pipeline,
/// which applies no output scaling. Two earlier gain stages were both
/// regressions against that baseline: per-sentence peak normalization caused
/// level pumping between sentences, and the fixed 9.3× gain that replaced it
/// was calibrated on a single anomalously-quiet bench utterance (peak 0.076)
/// and clipped 13–34% of samples on real speech ("blown out", 2026-06-12).
/// The clamp alone remains as the safety net against outlier transients.
/// No gain is applied because Pocket TTS already emits speech-level audio and
/// the reference pipeline applies no output scaling. Normalizing each sentence
/// would cause level pumping between chunks. The clamp remains only as a safety
/// net against outlier transients.
fn clamp_to_full_scale(samples: Vec<f32>) -> Vec<f32> {
samples.into_iter().map(|s| s.clamp(-1.0, 1.0)).collect()
}
@@ -678,14 +657,10 @@ fn clamp_to_full_scale(samples: Vec<f32>) -> Vec<f32> {
///
/// # Why no fade-in
///
/// An earlier revision (pre 2026-05) symmetrically faded *in* over the same
/// 8 ms window. That swallowed the leading consonant attack on every
/// sentence — Pocket TTS produces real audio energy inside the first
/// millisecond (RMS ≈ 0.02, peak ≈ 0.03 measured across four prompts in
/// `examples/pocket_onset_probe.rs`), and a linear 0→1 ramp over 192 samples
/// scales those onset samples by ≤50 % for the first ~4 ms. The result was
/// the "first little sound or two is missing" regression heard on
/// 2026-05-18.
/// A symmetric fade-in would attenuate the leading consonant attack because
/// Pocket TTS produces real audio energy inside the first millisecond. A
/// linear 0→1 ramp over 192 samples scales those onset samples by ≤50% for the
/// first ~4 ms, which can make the first phoneme sound clipped.
///
/// The first sample of Pocket output measures ≈ 0.0018 (≈ −54 dBFS) — well
/// below the threshold at which a DC-jump would be audible as a click — so