diff --git a/desktop/src-tauri/examples/pocket_quality_ab.rs b/desktop/src-tauri/examples/pocket_quality_ab.rs new file mode 100644 index 0000000000..0c31f1c910 --- /dev/null +++ b/desktop/src-tauri/examples/pocket_quality_ab.rs @@ -0,0 +1,519 @@ +//! Reproducible blind Pocket TTS quality corpus generator. +//! +//! Renders Buzz's production prompt preparation and post-processing across: +//! INT8/FP32 × per-sentence/grouped generation. The generated filenames are +//! deterministically blinded; keep `key.json` away from listeners until their +//! scoring sheet is complete. +//! +//! Usage: +//! cargo run --release --example pocket_quality_ab -- \ +//! [--idle-minutes N --only ITEM] +//! +//! The optional idle run intentionally creates one engine per condition, warms +//! all four, sleeps once, and then makes each clip the first generation after +//! dormancy. It requires `--only` because only the first synthesis after an +//! uninterrupted idle is a valid post-idle observation. Run each 5/15-minute +//! item as a separate process. + +// Importing the production module also brings in runtime-only helpers that this +// standalone corpus generator deliberately does not call. +#![allow(dead_code)] + +#[path = "../src/huddle/pocket.rs"] +mod production_pocket; +#[path = "../src/huddle/preprocessing.rs"] +mod production_preprocessing; + +use std::collections::HashMap; +use std::fs; +use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant}; + +use serde::Serialize; +use sha2::{Digest, Sha256}; +use sherpa_onnx::{GenerationConfig, OfflineTts, OfflineTtsConfig, Wave}; + +use production_pocket::{prepare_pocket_prompt, SAMPLE_RATE}; +use production_preprocessing::{preprocess_for_tts, split_sentences}; + +const NUM_STEPS: i32 = 1; +const SILENCE_SCALE: f32 = 1.0; +const INTER_SENTENCE_SILENCE_SAMPLES: usize = SAMPLE_RATE as usize / 10; +const LEAD_IN_SAMPLES: usize = SAMPLE_RATE as usize / 50; +const FADE_OUT_SAMPLES: usize = SAMPLE_RATE as usize * 8 / 1000; +const TARGET_RMS_DBFS: f32 = -23.0; +const BLINDING_SEED: &str = "pocket-quality-2026-07-21-v1"; + +const CORPUS: &[CorpusItem] = &[ + CorpusItem { id: "short_one_word", kind: "short", text: "Yep." }, + CorpusItem { id: "short_four_words", kind: "short", text: "Sounds good to me." }, + CorpusItem { + id: "multi_relay_review", + kind: "multi-sentence", + text: "I looked at the relay code this morning. The lease logic is solid. There's one race in the worker claim path, though. I'll write it up and send you a patch.", + }, + CorpusItem { + id: "multi_community_size", + kind: "multi-sentence", + text: "Great question. The answer is it depends on the community size. For small ones, keep it simple.", + }, + CorpusItem { + id: "mixed_agent_message", + kind: "mixed", + text: "That's 42 open PRs right now — mostly small. I'll triage them after lunch.", + }, +]; + +#[derive(Clone, Copy)] +struct CorpusItem { + id: &'static str, + kind: &'static str, + text: &'static str, +} + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +enum Precision { + Int8, + Fp32, +} + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +enum Chunking { + PerSentence, + Grouped, +} + +#[derive(Clone, Copy, Debug)] +struct Condition { + precision: Precision, + chunking: Chunking, +} + +const CONDITIONS: [Condition; 4] = [ + Condition { + precision: Precision::Int8, + chunking: Chunking::PerSentence, + }, + Condition { + precision: Precision::Int8, + chunking: Chunking::Grouped, + }, + Condition { + precision: Precision::Fp32, + chunking: Chunking::PerSentence, + }, + Condition { + precision: Precision::Fp32, + chunking: Chunking::Grouped, + }, +]; + +#[derive(Serialize)] +struct KeyFile { + warning: &'static str, + blinding_seed: &'static str, + target_rms_dbfs: f32, + items: Vec, +} + +#[derive(Serialize)] +struct KeyItem { + id: String, + kind: String, + text: String, + clips: Vec, +} + +#[derive(Serialize)] +struct KeyClip { + file: String, + precision: Precision, + chunking: Chunking, + cold_start: bool, + idle_minutes: Option, + synthesis_ms: u128, + audio_seconds: f32, +} + +struct Voice { + samples: Vec, + sample_rate: i32, +} + +struct Engine { + inner: OfflineTts, + voice: Voice, +} + +fn main() -> Result<(), String> { + let mut args = std::env::args().skip(1); + let int8_dir = required_path(args.next(), "INT8 model directory")?; + let fp32_dir = required_path(args.next(), "FP32 model directory")?; + let output_dir = required_path(args.next(), "output directory")?; + let mut idle_minutes = None; + let mut only_item = None; + while let Some(arg) = args.next() { + match arg.as_str() { + "--idle-minutes" => { + idle_minutes = Some( + args.next() + .ok_or("--idle-minutes requires a value")? + .parse::() + .map_err(|e| format!("invalid idle minutes: {e}"))?, + ); + } + "--only" => only_item = Some(args.next().ok_or("--only requires an item ID")?), + _ => return Err(format!("unknown argument: {arg}")), + } + } + + if idle_minutes.is_some() && only_item.is_none() { + return Err("--idle-minutes requires --only so every clip is first-after-idle".into()); + } + if let Some(ref requested) = only_item { + if !CORPUS.iter().any(|item| item.id == requested) { + return Err(format!("unknown corpus item for --only: {requested}")); + } + } + + validate_model_dir(&int8_dir, Precision::Int8)?; + validate_model_dir(&fp32_dir, Precision::Fp32)?; + fs::create_dir_all(&output_dir).map_err(|e| e.to_string())?; + + let mut engines = Vec::with_capacity(CONDITIONS.len()); + for condition in CONDITIONS { + let dir = match condition.precision { + Precision::Int8 => &int8_dir, + Precision::Fp32 => &fp32_dir, + }; + let engine = load_engine(dir, condition.precision)?; + // Production warms once before serving a real utterance. Cold cases use + // separate fresh engines below and deliberately skip this call. + synth_chunks(&engine, &["warmup".to_string()])?; + engines.push(engine); + } + + if let Some(minutes) = idle_minutes { + eprintln!("All four warmed engines idle for {minutes} minute(s)…"); + std::thread::sleep(Duration::from_secs(minutes * 60)); + } + + let mut key_items = Vec::new(); + for item in CORPUS { + if only_item + .as_deref() + .is_some_and(|requested| requested != item.id) + { + continue; + } + let preprocessed = preprocess_for_tts(item.text); + let per_sentence: Vec = split_sentences(&preprocessed) + .into_iter() + .filter(|s| !s.trim().is_empty()) + .collect(); + // These corpus texts are deliberately below the upstream ~50-token + // grouping target, so grouped mode is one exact generate() call. + let grouped = vec![per_sentence.join(" ")]; + let item_dir = output_dir.join(item.id); + fs::create_dir_all(&item_dir).map_err(|e| e.to_string())?; + let clip_order = blinded_order(item.id); + let mut clips = Vec::new(); + + let mut rendered = Vec::new(); + for (condition_index, engine) in engines.iter().enumerate() { + let condition = CONDITIONS[condition_index]; + let chunks = match condition.chunking { + Chunking::PerSentence => &per_sentence, + Chunking::Grouped => &grouped, + }; + let started = Instant::now(); + let audio = synth_chunks(engine, chunks)?; + rendered.push(( + condition_index, + condition, + audio, + started.elapsed().as_millis(), + )); + } + loudness_match_item(&mut rendered); + for (condition_index, condition, audio, synth_ms) in rendered { + let clip_number = clip_order[condition_index] + 1; + let file_name = format!("clip{clip_number}.wav"); + write_wav(&item_dir.join(&file_name), &audio)?; + clips.push(KeyClip { + file: format!("{}/{file_name}", item.id), + precision: condition.precision, + chunking: condition.chunking, + cold_start: false, + idle_minutes, + synthesis_ms: synth_ms, + audio_seconds: audio.len() as f32 / SAMPLE_RATE as f32, + }); + } + clips.sort_by(|a, b| a.file.cmp(&b.file)); + key_items.push(KeyItem { + id: item.id.to_string(), + kind: item.kind.to_string(), + text: item.text.to_string(), + clips, + }); + } + + // Explicit fresh-engine cold-start clips for the two highest-signal texts. + // Idle runs intentionally omit them: they happen after the post-idle clips + // and add no valid idle observation. + for item in if idle_minutes.is_none() { CORPUS } else { &[] } { + if !matches!(item.id, "short_one_word" | "multi_relay_review") { + continue; + } + if only_item + .as_deref() + .is_some_and(|requested| requested != item.id) + { + continue; + } + let cold_id = format!("cold_{}", item.id); + let preprocessed = preprocess_for_tts(item.text); + let sentences: Vec = split_sentences(&preprocessed) + .into_iter() + .filter(|s| !s.trim().is_empty()) + .collect(); + let grouped = vec![sentences.join(" ")]; + let item_dir = output_dir.join(&cold_id); + fs::create_dir_all(&item_dir).map_err(|e| e.to_string())?; + let clip_order = blinded_order(&cold_id); + let mut clips = Vec::new(); + let mut rendered = Vec::new(); + for (condition_index, condition) in CONDITIONS.iter().copied().enumerate() { + let dir = match condition.precision { + Precision::Int8 => &int8_dir, + Precision::Fp32 => &fp32_dir, + }; + let engine = load_engine(dir, condition.precision)?; + let chunks = match condition.chunking { + Chunking::PerSentence => &sentences, + Chunking::Grouped => &grouped, + }; + let started = Instant::now(); + let audio = synth_chunks(&engine, chunks)?; + rendered.push(( + condition_index, + condition, + audio, + started.elapsed().as_millis(), + )); + } + loudness_match_item(&mut rendered); + for (condition_index, condition, audio, synth_ms) in rendered { + let clip_number = clip_order[condition_index] + 1; + let file_name = format!("clip{clip_number}.wav"); + write_wav(&item_dir.join(&file_name), &audio)?; + clips.push(KeyClip { + file: format!("{cold_id}/{file_name}"), + precision: condition.precision, + chunking: condition.chunking, + cold_start: true, + idle_minutes: None, + synthesis_ms: synth_ms, + audio_seconds: audio.len() as f32 / SAMPLE_RATE as f32, + }); + } + clips.sort_by(|a, b| a.file.cmp(&b.file)); + key_items.push(KeyItem { + id: cold_id, + kind: "cold-start".to_string(), + text: item.text.to_string(), + clips, + }); + } + + let key = KeyFile { + warning: "DO NOT OPEN UNTIL LISTENING SCORES ARE FINAL", + blinding_seed: BLINDING_SEED, + target_rms_dbfs: TARGET_RMS_DBFS, + items: key_items, + }; + fs::write( + output_dir.join("key.json"), + serde_json::to_vec_pretty(&key).map_err(|e| e.to_string())?, + ) + .map_err(|e| e.to_string())?; + write_scoring_sheet(&output_dir, &key)?; + println!("Wrote blind corpus to {}", output_dir.display()); + println!("Give listeners the WAV folders and SCORING.md; withhold key.json."); + Ok(()) +} + +fn required_path(value: Option, label: &str) -> Result { + value + .map(PathBuf::from) + .ok_or_else(|| format!("missing {label}")) +} + +fn model_file(precision: Precision, base: &str) -> String { + match precision { + Precision::Int8 => format!("{base}.int8.onnx"), + Precision::Fp32 => format!("{base}.onnx"), + } +} + +fn validate_model_dir(dir: &Path, precision: Precision) -> Result<(), String> { + for file in [ + model_file(precision, "lm_main"), + model_file(precision, "lm_flow"), + "encoder.onnx".into(), + model_file(precision, "decoder"), + "text_conditioner.onnx".into(), + "vocab.json".into(), + "token_scores.json".into(), + "reference_sample.wav".into(), + ] { + if !dir.join(&file).is_file() { + return Err(format!("missing {}", dir.join(file).display())); + } + } + Ok(()) +} + +fn load_engine(dir: &Path, precision: Precision) -> Result { + let p = |name: &str| dir.join(name).to_string_lossy().into_owned(); + let mut cfg = OfflineTtsConfig::default(); + cfg.model.pocket.lm_main = Some(p(&model_file(precision, "lm_main"))); + cfg.model.pocket.lm_flow = Some(p(&model_file(precision, "lm_flow"))); + cfg.model.pocket.encoder = Some(p("encoder.onnx")); + cfg.model.pocket.decoder = Some(p(&model_file(precision, "decoder"))); + cfg.model.pocket.text_conditioner = Some(p("text_conditioner.onnx")); + cfg.model.pocket.vocab_json = Some(p("vocab.json")); + cfg.model.pocket.token_scores_json = Some(p("token_scores.json")); + cfg.model.pocket.voice_embedding_cache_capacity = 16; + cfg.model.num_threads = 1; + cfg.model.debug = false; + let inner = + OfflineTts::create(&cfg).ok_or_else(|| format!("failed to create {precision:?} engine"))?; + let wave = + Wave::read(&p("reference_sample.wav")).ok_or("failed to read reference_sample.wav")?; + Ok(Engine { + inner, + voice: Voice { + samples: wave.samples().to_vec(), + sample_rate: wave.sample_rate(), + }, + }) +} + +fn synth_chunks(engine: &Engine, chunks: &[String]) -> Result, String> { + let mut out = Vec::new(); + for chunk in chunks { + let prepared = prepare_pocket_prompt(chunk).ok_or("empty prepared prompt")?; + let extra = prepared.max_frames.map(|max_frames| { + HashMap::from([( + "max_frames".to_string(), + serde_json::Value::from(max_frames), + )]) + }); + let cfg = GenerationConfig { + num_steps: NUM_STEPS, + silence_scale: SILENCE_SCALE, + reference_audio: Some(engine.voice.samples.clone()), + reference_sample_rate: engine.voice.sample_rate, + extra, + ..Default::default() + }; + let audio = engine + .inner + .generate_with_config(&prepared.text, &cfg, None:: bool>) + .ok_or_else(|| format!("synthesis failed for {chunk:?}"))?; + let mut samples: Vec = audio.samples().iter().map(|s| s.clamp(-1.0, 1.0)).collect(); + apply_fade_out(&mut samples); + out.extend(std::iter::repeat_n(0.0, LEAD_IN_SAMPLES)); + out.extend(samples); + out.extend(std::iter::repeat_n( + 0.0, + INTER_SENTENCE_SILENCE_SAMPLES - LEAD_IN_SAMPLES, + )); + } + Ok(out) +} + +fn apply_fade_out(samples: &mut [f32]) { + let fade = FADE_OUT_SAMPLES.min(samples.len() / 2); + for i in 0..fade { + samples[samples.len() - 1 - i] *= i as f32 / fade as f32; + } +} + +fn active_rms(samples: &[f32]) -> Option { + let (sum_squares, count) = samples + .iter() + .filter(|sample| sample.abs() > 1.0e-4) + .fold((0.0_f32, 0_usize), |(sum, count), sample| { + (sum + sample * sample, count + 1) + }); + (count > 0).then(|| (sum_squares / count as f32).sqrt()) +} + +/// Attenuate every clip in one comparison set to the quietest active-speech RMS. +/// This removes the louder-is-better confound without normalizing dynamics or +/// claiming standards-compliant integrated LUFS. The dBFS value is a ceiling. +fn loudness_match_item(rendered: &mut [(usize, Condition, Vec, u128)]) { + let ceiling = 10.0_f32.powf(TARGET_RMS_DBFS / 20.0); + let target = rendered + .iter() + .filter_map(|(_, _, samples, _)| active_rms(samples)) + .fold(ceiling, f32::min); + for (_, _, samples, _) in rendered { + let Some(rms) = active_rms(samples) else { + continue; + }; + let gain = (target / rms).min(1.0); + for sample in samples { + *sample *= gain; + } + } +} + +fn blinded_order(item_id: &str) -> [usize; 4] { + let mut keyed: Vec<(usize, Vec)> = (0..4) + .map(|index| { + let digest = Sha256::digest(format!("{BLINDING_SEED}:{item_id}:{index}")); + (index, digest.to_vec()) + }) + .collect(); + keyed.sort_by(|a, b| a.1.cmp(&b.1)); + let mut condition_to_clip = [0; 4]; + for (clip, (condition, _)) in keyed.into_iter().enumerate() { + condition_to_clip[condition] = clip; + } + condition_to_clip +} + +fn write_wav(path: &Path, samples: &[f32]) -> Result<(), String> { + let path = path + .to_str() + .ok_or_else(|| format!("non-UTF8 path: {}", path.display()))?; + if sherpa_onnx::write(path, samples, SAMPLE_RATE as i32) { + Ok(()) + } else { + Err(format!("failed to write {path}")) + } +} + +fn write_scoring_sheet(output_dir: &Path, key: &KeyFile) -> Result<(), String> { + let mut sheet = String::from("# Pocket TTS blind listening sheet\n\nDo not open `key.json` until this sheet is complete. Rank best to worst; ties are allowed.\n\n"); + for item in &key.items { + sheet.push_str(&format!( + "## {} ({})\n\n> {}\n\n", + item.id, item.kind, item.text + )); + sheet.push_str("Rank: `____ > ____ > ____ > ____`\n\n| Clip | seam | onset | garble | robotic | timbre | truncate | note |\n|---|---|---|---|---|---|---|---|\n"); + for clip in 1..=4 { + sheet.push_str(&format!( + "| clip{clip} | [ ] | [ ] | [ ] | [ ] | [ ] | [ ] | |\n" + )); + } + sheet.push('\n'); + } + fs::write(output_dir.join("SCORING.md"), sheet).map_err(|e| e.to_string()) +} diff --git a/desktop/src-tauri/src/huddle/models.rs b/desktop/src-tauri/src/huddle/models.rs index 8571c40436..169ddf66c0 100644 --- a/desktop/src-tauri/src/huddle/models.rs +++ b/desktop/src-tauri/src/huddle/models.rs @@ -38,12 +38,19 @@ use sha2::{Digest, Sha256}; /// Computed from a known-good download. Update when upgrading model versions. const STT_ARCHIVE_SHA256: &str = "17f945007b52ccd8b7200ffc7c5652e9e8e961dfdf479cefcabd06cf5703630b"; -/// HuggingFace base URL for the sherpa-onnx Pocket TTS int8 repackage. +/// HuggingFace base URL for the sherpa-onnx Pocket TTS fp32 repackage. /// -/// Pinned to commit e715955cf50d18d919d37231513c0e914b83661a +/// Pinned to commit 96d1e53ce3311ca6c2c6a35e2062d36b4cec6fa3 /// (2026-02-10) for reproducible downloads. +/// +/// fp32 (not int8): a direct same-runtime A/B (k2-fsa/sherpa-onnx#3172) +/// found the ONNX int8 quantization audibly degraded Pocket TTS output and +/// that fp32 "significantly improved quality even at 1 step". The runtime +/// bundle grows from ~189 MB to ~473 MB; encoder, text conditioner, both +/// JSON tables, and LICENSE are byte-identical between the two repos — only +/// the three quantized sessions (lm_main, lm_flow, decoder) change. const POCKET_HF_BASE: &str = - "https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-int8-2026-01-26/resolve/e715955cf50d18d919d37231513c0e914b83661a"; + "https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-2026-01-26/resolve/96d1e53ce3311ca6c2c6a35e2062d36b4cec6fa3"; /// Reference voice WAV: "Mary (f, conversation)" from the Kyutai TTS demo /// voice set — VCTK speaker p333, ai-coustics-enhanced. Pinned to @@ -61,10 +68,10 @@ const POCKET_REFERENCE_WAV_URL: &str = /// Computed from known-good pinned downloads. Update when upgrading model versions. #[rustfmt::skip] const TTS_FILE_HASHES: &[(&str, &str)] = &[ - ("decoder.int8.onnx", "12b0857402d31aead94df19d6783b4350d1f740e811f3a3202c70ad89ae11eea"), + ("decoder.onnx", "f267880fde6c58b17b0a8f3647eaf8dcfad321f833f32d583ebc2fb2d1a15f10"), ("encoder.onnx", "e8f2f6d301ffb96e398b138a7dc6d3038622d236044636b73d920bab85890260"), - ("lm_flow.int8.onnx", "8d627d235c44a597da908e1085ebe241cbbe358964c502c5a5063d18851a5529"), - ("lm_main.int8.onnx", "bfc0c7e7e3d72864fa3bb2ee499f62f21ddc1474b885f5f3ca570f8be73e787e"), + ("lm_flow.onnx", "79c013a554a54e63319c33c0cc8830cbbedc9b7e448ae7e26f7923ae11f9873e"), + ("lm_main.onnx", "255d1a9263c5abdf36034abfc19c11d21cc5f40f0f87d8361288e972cbd5c578"), ("text_conditioner.onnx", "0b84e837d7bfaf2c896627b03e3f080320309f37f4fc7df7698c644f7ba5e6b1"), ("vocab.json", "6fb646346cf931016f70c4921aab0900ce7a304b893cb02135c74e294abfea01"), ("token_scores.json", "5be2f278caf9b9800741f0fd82bff677f4943ec764c356f907213434b622d958"), @@ -91,7 +98,9 @@ const STT_MODEL_VERSION: &str = "2"; /// from kyutai/tts-voices. The hash mismatch on `reference_sample.wav` would /// fail readiness on its own, but the manifest bump makes the re-download /// reason explicit and skips the failing-then-re-fetching transient state. -const TTS_MODEL_VERSION: &str = "2"; +/// Bumped "2" → "3" for the int8 → fp32 model swap (see `POCKET_HF_BASE`): +/// existing int8 installs must re-download the suffixless fp32 sessions. +const TTS_MODEL_VERSION: &str = "3"; /// Filename for the version manifest written alongside model files. const MANIFEST_FILENAME: &str = ".buzz-model-manifest"; @@ -101,8 +110,9 @@ const MANIFEST_FILENAME: &str = ".buzz-model-manifest"; /// Maximum expected STT archive size (200 MB — actual is ~100 MB). const MAX_STT_DOWNLOAD_BYTES: u64 = 200 * 1024 * 1024; -/// Maximum expected Pocket TTS file size (200 MB per file — largest is ~73 MB). -const MAX_TTS_FILE_BYTES: u64 = 200 * 1024 * 1024; +/// Maximum expected Pocket TTS file size (400 MB per file — largest is +/// `lm_main.onnx` at ~303 MB fp32). +const MAX_TTS_FILE_BYTES: u64 = 400 * 1024 * 1024; /// NVIDIA Parakeet TDT-CTC 110M (English, int8) — packaged for sherpa-onnx by /// k2-fsa. Single ONNX file (CTC head) + tokens.txt. Avg WER ~7.5% across @@ -173,7 +183,7 @@ Mimi neural codec by Kyutai is bundled as part of the model. ONNX export by KevinAHM: https://huggingface.co/KevinAHM/pocket-tts-onnx Sherpa-onnx repackage by csukuangfj / k2-fsa: -https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-int8-2026-01-26 +https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-2026-01-26 Bundled reference voice (reference_sample.wav): \"Mary (f, conversation)\" preset from the Kyutai TTS demo voice catalogue @@ -193,10 +203,10 @@ license text for full warranty disclaimer. /// All files that must be present for Pocket TTS to be considered ready. const TTS_EXPECTED_FILES: &[&str] = &[ - "decoder.int8.onnx", + "decoder.onnx", "encoder.onnx", - "lm_flow.int8.onnx", - "lm_main.int8.onnx", + "lm_flow.onnx", + "lm_main.onnx", "text_conditioner.onnx", "vocab.json", "token_scores.json", @@ -759,10 +769,10 @@ impl ModelManager { fresh_temp_dir(&temp_dir).await?; let model_files = [ - "decoder.int8.onnx", + "decoder.onnx", "encoder.onnx", - "lm_flow.int8.onnx", - "lm_main.int8.onnx", + "lm_flow.onnx", + "lm_main.onnx", "text_conditioner.onnx", "vocab.json", "token_scores.json", diff --git a/desktop/src-tauri/src/huddle/pocket.rs b/desktop/src-tauri/src/huddle/pocket.rs index 9ae0465036..ee1faf928a 100644 --- a/desktop/src-tauri/src/huddle/pocket.rs +++ b/desktop/src-tauri/src/huddle/pocket.rs @@ -1,10 +1,15 @@ //! Pocket TTS engine wrapper around sherpa-onnx's `OfflineTts`. //! -//! Pocket TTS is a small (~189 MB int8 ONNX) zero-shot voice-cloning TTS +//! Pocket TTS is a small (~473 MB fp32 ONNX) zero-shot voice-cloning TTS //! model from Kyutai. It runs quickly on CPU via sherpa-onnx, replacing the //! previous Kokoro-82M engine that also required an espeak-free but //! lexicon-heavy G2P pipeline (Misaki + CMUdict). //! +//! Full-precision fp32 sessions, not the ~189 MB int8 quantization we +//! originally shipped: a direct same-runtime A/B (k2-fsa/sherpa-onnx#3172) +//! found the int8 ONNX export audibly degraded output quality, and fp32 +//! "significantly improved quality even at 1 step". +//! //! ## Attribution //! //! - **Model**: Kyutai *Pocket TTS* — Charles, Roebel, et al., 2026. @@ -14,7 +19,7 @@ //! - **ONNX export**: KevinAHM — //! . CC-BY-4.0. //! - **sherpa-onnx repackage**: csukuangfj / k2-fsa — -//! . +//! . //! Repackages KevinAHM's export with the file layout sherpa-onnx's //! `OfflineTtsPocketModelConfig` expects. CC-BY-4.0. //! - **Reference voice WAV** (`reference_sample.wav`): the "Mary @@ -37,12 +42,16 @@ //! - `VOICE_FILE_EXT: &str` — extension for per-voice files on disk. //! - `load_text_to_speech(model_dir)` → `Result` //! - `load_voice_style(path)` → `Result` -//! - `Engine::synth_chunk(&self, text, lang, &VoiceStyle, steps, speed)` +//! - `Engine::synth_chunk(&self, text, lang, &VoiceStyle, steps)` //! → `Result, String>` //! //! `lang` and `steps` are accepted for API compatibility with the previous //! Kokoro engine but are unused — Pocket TTS does its own language ID from //! the input text and is not a diffusion model (consistency LM, one step). +//! There is no speed knob: sherpa-onnx's `GenerationConfig.speed` is only +//! read by some model families (vits), never by the Pocket impl +//! (`offline-tts-pocket-impl.h` — zero references), and upstream pocket-tts +//! has no speed parameter either. use std::collections::HashMap; use std::path::{Path, PathBuf}; @@ -143,10 +152,10 @@ const SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT: i32 = 3; // ── ONNX file names (five Pocket TTS sessions plus two JSON tables) ─────────── -const FILE_LM_MAIN: &str = "lm_main.int8.onnx"; -const FILE_LM_FLOW: &str = "lm_flow.int8.onnx"; +const FILE_LM_MAIN: &str = "lm_main.onnx"; +const FILE_LM_FLOW: &str = "lm_flow.onnx"; const FILE_ENCODER: &str = "encoder.onnx"; -const FILE_DECODER: &str = "decoder.int8.onnx"; +const FILE_DECODER: &str = "decoder.onnx"; const FILE_TEXT_COND: &str = "text_conditioner.onnx"; const FILE_VOCAB: &str = "vocab.json"; const FILE_TOKEN_SCORES: &str = "token_scores.json"; @@ -385,7 +394,6 @@ impl PocketTts { _lang: &str, style: &VoiceStyle, _steps: usize, - speed: f32, ) -> Result, String> { // Mirror upstream pocket-tts prompt prep — without this short or // unpunctuated inputs can cause the LM's EOS logit to never trip, @@ -405,12 +413,13 @@ impl PocketTts { let extra = build_generation_extra(&prepared); let cfg = GenerationConfig { - speed, num_steps: SYNTH_NUM_STEPS, silence_scale: SYNTH_SILENCE_SCALE, reference_audio: Some(style.samples.clone()), reference_sample_rate: style.sample_rate, extra, + // `speed` stays at its default: the Pocket impl never reads it + // (see the engine-contract note in the module docs). ..Default::default() }; diff --git a/desktop/src-tauri/src/huddle/tts.rs b/desktop/src-tauri/src/huddle/tts.rs index 436f455be1..63a435cd8e 100644 --- a/desktop/src-tauri/src/huddle/tts.rs +++ b/desktop/src-tauri/src/huddle/tts.rs @@ -69,9 +69,6 @@ const MONITOR_TICK: Duration = Duration::from_millis(10); /// Pocket TTS is a one-step consistency model, not diffusion. Kept for API compat. const SYNTH_STEPS: usize = 1; -/// Synthesis speed multiplier. Slightly faster than natural speech. -const SYNTH_SPEED: f32 = 1.05; - /// Fade-out length in samples (8 ms at 24 kHz ≈ 192 samples). /// /// Applied only at the *end* of each synthesised sentence to eliminate the @@ -94,10 +91,24 @@ const FADE_OUT_SAMPLES: usize = (SAMPLE_RATE as f64 * 0.008) as usize; /// existing inter-sentence pause, so it does not lengthen multi-sentence gaps. const SENTENCE_LEAD_IN_SAMPLES: usize = (SAMPLE_RATE as f64 * 0.020) as usize; -/// Sentence-by-sentence synthesis — keeps first-sentence latency low and lets -/// playback of sentence N overlap with synthesis of sentence N+1 (see the -/// lookahead pipelining note in the module doc-comment above). -const BATCH_SIZE: usize = 1; +/// Approximate character budget for one synthesis chunk. +/// +/// Upstream pocket-tts groups sentences into chunks of up to +/// `MAX_TOKEN_PER_CHUNK = 50` tokenizer tokens (`default_parameters.py`) — +/// typically multi-sentence chunks — because every `generate()` call is an +/// independent generation with a cold FlowLM start, and each chunk boundary +/// is an exposed prosody seam (kyutai-labs/pocket-tts #151; the Kyutai team +/// names chunk stitching as the reliability lever). Our previous +/// sentence-per-call path created ~2–4× more seams than upstream. +/// +/// We don't ship the SentencePiece tokenizer, so 50 tokens is approximated +/// with a character budget. The bundled 4k-entry vocab averages ~4 chars per +/// token, but usage-weighted English text leans on short common tokens, so +/// the effective ratio is ~2–4 chars/token and 200 chars ≈ 60–100 tokens — +/// modestly above upstream's 50, deliberately: erring large means fewer +/// seams, and even ~100 tokens is far below the model's 500-LM-step (~40 s) +/// ceiling. Do not shrink this budget to chase an exact 50-token match. +const MAX_CHUNK_CHARS: usize = 200; /// Silence inserted between sentences by the TTS pipeline (seconds). /// Injected as a silent buffer between each synthesized sentence chunk. @@ -275,7 +286,7 @@ fn tts_worker( // and discard the output so the first real utterance runs at warm-session speed. { let t = std::time::Instant::now(); - match engine.synth_chunk("warmup", "en", &style, SYNTH_STEPS, SYNTH_SPEED) { + match engine.synth_chunk("warmup", "en", &style, SYNTH_STEPS) { Ok(_) => eprintln!( "buzz-desktop: TTS warmup completed in {:.0}ms", t.elapsed().as_millis() @@ -480,15 +491,20 @@ fn tts_worker( continue; } - // Split into sentences. Each sentence is synthesized individually and - // appended to the Player immediately — synthesis of sentence N+1 overlaps - // with playback of sentence N (lookahead pipelining). + // Split into sentences, then group into synthesis chunks: the first + // 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. let sentences: Vec = split_sentences(&text) .into_iter() .filter(|s| !s.trim().is_empty()) .collect(); + let chunks = group_sentences_into_chunks(&sentences, MAX_CHUNK_CHARS); - for sentence in &sentences { + for chunk in &chunks { if handle_cancel_or_shutdown( &cancel, &shutdown, @@ -500,12 +516,12 @@ fn tts_worker( break; } - let text = sentence.trim(); + let text = chunk.trim(); if text.is_empty() { continue; } - match engine.synth_chunk(text, "en", &style, SYNTH_STEPS, SYNTH_SPEED) { + match engine.synth_chunk(text, "en", &style, SYNTH_STEPS) { Ok(samples) if !samples.is_empty() => { let mut audio = clamp_to_full_scale(samples); // Fade-out only — fading-in would attenuate the consonant @@ -711,13 +727,50 @@ fn build_sentence_append_buffer( buf } +/// Group sentences into synthesis chunks. +/// +/// The first sentence always stands alone — it is what the listener hears +/// first, and synthesizing it by itself keeps time-to-first-audio at the +/// single-sentence cost. Subsequent sentences pack greedily: a sentence +/// joins the current chunk while the combined length stays within +/// `max_chars`; otherwise it starts a new chunk. A single sentence longer +/// than `max_chars` becomes its own chunk unsplit — Pocket TTS handles long +/// single sentences fine (the ceiling is the 500-LM-step default), it's the +/// *seams* we're minimizing. +/// +/// Sentences within a chunk are joined with a single space; sentence-ending +/// punctuation is preserved by `split_sentences`, so the model sees natural +/// multi-sentence prose — the same shape upstream's ~50-token chunker feeds it. +fn group_sentences_into_chunks(sentences: &[String], max_chars: usize) -> Vec { + let mut chunks: Vec = Vec::new(); + for (i, sentence) in sentences.iter().enumerate() { + let sentence = sentence.trim(); + if sentence.is_empty() { + continue; + } + if i == 0 || chunks.is_empty() { + chunks.push(sentence.to_string()); + continue; + } + // Never merge into the first chunk — it's the latency-critical one. + let can_merge = chunks.len() > 1 + && chunks + .last() + .is_some_and(|c| c.len() + 1 + sentence.len() <= max_chars); + if can_merge { + let last = chunks.last_mut().expect("non-empty checked above"); + last.push(' '); + last.push_str(sentence); + } else { + chunks.push(sentence.to_string()); + } + } + chunks +} + // drain_until_shutdown lives in super (huddle/mod.rs) — shared with stt.rs. use super::drain_until_shutdown; -// BATCH_SIZE is used implicitly (one sentence per iteration). Suppress dead_code -// lint since it documents the design intent. -const _: () = assert!(BATCH_SIZE == 1); - // ── Tests ───────────────────────────────────────────────────────────────────── #[cfg(test)] diff --git a/desktop/src-tauri/src/huddle/tts_tests.rs b/desktop/src-tauri/src/huddle/tts_tests.rs index 5668172a5a..7887f8bbdb 100644 --- a/desktop/src-tauri/src/huddle/tts_tests.rs +++ b/desktop/src-tauri/src/huddle/tts_tests.rs @@ -915,3 +915,80 @@ fn clamp_to_full_scale_empty_buffer() { let out = clamp_to_full_scale(Vec::new()); assert!(out.is_empty()); } + +// ── group_sentences_into_chunks tests ───────────────────────────────────── + +fn s(v: &[&str]) -> Vec { + v.iter().map(|x| x.to_string()).collect() +} + +/// The first sentence always stands alone — it bounds time-to-first-audio. +/// Even when the whole message would fit in one chunk, sentence one must +/// not wait on synthesis of the rest. +#[test] +fn chunk_grouping_first_sentence_is_always_alone() { + let chunks = group_sentences_into_chunks(&s(&["Hi there.", "Short.", "Tiny."]), 200); + assert_eq!(chunks[0], "Hi there."); + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[1], "Short. Tiny."); +} + +/// Sentences after the first pack greedily up to the char budget, then +/// spill into a new chunk. Fewer generate() calls = fewer prosody seams. +#[test] +fn chunk_grouping_packs_up_to_budget_then_spills() { + let a = "A".repeat(50) + "."; + let b = "B".repeat(50) + "."; + let c = "C".repeat(50) + "."; + let d = "D".repeat(50) + "."; + // Budget of 110: b+c fits (51+1+51 = 103), adding d (103+1+51) does not. + let chunks = group_sentences_into_chunks(&s(&[&a, &b, &c, &d]), 110); + assert_eq!(chunks.len(), 3, "chunks: {chunks:?}"); + assert_eq!(chunks[0], a); + assert_eq!(chunks[1], format!("{b} {c}")); + assert_eq!(chunks[2], d); +} + +/// A single sentence longer than the budget is passed through unsplit — +/// long single sentences are fine (the LM cap bounds runaway); only seams +/// are being minimized. +#[test] +fn chunk_grouping_oversized_sentence_passes_through() { + let long = "word ".repeat(60).trim_end().to_string() + "."; + assert!(long.len() > 200); + let chunks = group_sentences_into_chunks(&s(&["First.", &long]), 200); + assert_eq!(chunks, vec!["First.".to_string(), long]); +} + +/// Single-sentence messages — the common huddle case, since agents are +/// prompted to send one sentence per message — are unaffected by grouping. +#[test] +fn chunk_grouping_single_sentence_unchanged() { + let chunks = group_sentences_into_chunks(&s(&["Just one sentence here."]), 200); + assert_eq!(chunks, vec!["Just one sentence here.".to_string()]); +} + +/// Empty and whitespace-only entries are dropped, and never produce +/// empty chunks (which would synthesize as garbage). +#[test] +fn chunk_grouping_skips_blank_sentences() { + let chunks = group_sentences_into_chunks(&s(&["", " ", "Real sentence.", " ", "Two."]), 200); + assert_eq!(chunks[0], "Real sentence."); + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[1], "Two."); +} + +/// Empty input produces no chunks (the worker loop then synthesizes nothing). +#[test] +fn chunk_grouping_empty_input() { + assert!(group_sentences_into_chunks(&[], 200).is_empty()); +} + +/// Chunks joined with a single space preserve each sentence's terminal +/// punctuation — the model sees natural multi-sentence prose, matching the +/// shape upstream's ~50-token chunker produces. +#[test] +fn chunk_grouping_preserves_punctuation_at_joins() { + let chunks = group_sentences_into_chunks(&s(&["Lead.", "Really?", "Yes!", "Good."]), 200); + assert_eq!(chunks[1], "Really? Yes! Good."); +}