From da9c5699d28d089645afcaaaa27ee17d3a41fd2a Mon Sep 17 00:00:00 2001 From: John Tennant Date: Tue, 28 Jul 2026 14:07:36 -0400 Subject: [PATCH 1/6] Upgrade Desktop Pocket TTS to April INT8 Signed-off-by: John Tennant --- desktop/src-tauri/Cargo.lock | 472 ++++++++- desktop/src-tauri/Cargo.toml | 5 + .../src-tauri/examples/pocket_quality_ab.rs | 519 ---------- desktop/src-tauri/src/huddle/models.rs | 241 +++-- desktop/src-tauri/src/huddle/models_tests.rs | 148 +++ desktop/src-tauri/src/huddle/pocket.rs | 672 ++----------- desktop/src-tauri/src/huddle/pocket_april.rs | 940 ++++++++++++++++++ desktop/src-tauri/src/huddle/pocket_models.rs | 130 +++ desktop/src-tauri/src/huddle/tts.rs | 32 +- desktop/src-tauri/src/huddle/tts_tests.rs | 5 +- 10 files changed, 1931 insertions(+), 1233 deletions(-) delete mode 100644 desktop/src-tauri/examples/pocket_quality_ab.rs create mode 100644 desktop/src-tauri/src/huddle/models_tests.rs create mode 100644 desktop/src-tauri/src/huddle/pocket_april.rs create mode 100644 desktop/src-tauri/src/huddle/pocket_models.rs diff --git a/desktop/src-tauri/Cargo.lock b/desktop/src-tauri/Cargo.lock index 66553ef595..c470484f36 100644 --- a/desktop/src-tauri/Cargo.lock +++ b/desktop/src-tauri/Cargo.lock @@ -75,6 +75,20 @@ dependencies = [ "subtle", ] +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if 1.0.4", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -694,6 +708,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fd307490d624467aa6f74b0eabb77633d1f758a7b25f12bceb0b22e08d9726f6" +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + [[package]] name = "base64" version = "0.21.7" @@ -718,6 +738,12 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32637268377fc7b10a8c6d51de3e7fba1ce5dd371a96e342b34e6078db558e7f" +[[package]] +name = "beef" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a8241f3ebb85c056b509d4327ad0358fbbba6ffb340bf388f26350aeda225b1" + [[package]] name = "bip39" version = "2.2.2" @@ -1049,8 +1075,11 @@ dependencies = [ "notify-rust", "objc2-app-kit", "opus", + "ort", + "ort-sys", "plist", "png 0.18.1", + "rand 0.10.2", "regex", "reqwest 0.13.4", "rodio", @@ -1058,6 +1087,7 @@ dependencies = [ "rusqlite", "rustls", "security-framework 3.7.0", + "sentencepiece-model", "serde", "serde_json", "serde_yaml", @@ -1077,6 +1107,7 @@ dependencies = [ "tauri-plugin-updater", "tauri-plugin-window-state", "tempfile", + "tokenizers", "tokio", "tokio-tungstenite 0.29.0", "tokio-util", @@ -1545,6 +1576,7 @@ dependencies = [ "itoa", "rustversion", "ryu", + "serde", "static_assertions", ] @@ -1796,6 +1828,16 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + [[package]] name = "crossbeam-epoch" version = "0.9.20" @@ -2120,6 +2162,15 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + [[package]] name = "dasp_sample" version = "0.11.0" @@ -2617,6 +2668,12 @@ version = "3.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59" +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + [[package]] name = "euclid" version = "0.22.14" @@ -2675,6 +2732,17 @@ dependencies = [ "regex", ] +[[package]] +name = "fancy-regex" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e24cb5a94bcae1e5408b0effca5cd7172ea3c5755049c5f3af4cd283a165298" +dependencies = [ + "bit-set 0.8.0", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fast-srgb8" version = "1.0.0" @@ -4746,6 +4814,39 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +[[package]] +name = "logos" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7251356ef8cb7aec833ddf598c6cb24d17b689d20b993f9d11a3d764e34e6458" +dependencies = [ + "logos-derive", +] + +[[package]] +name = "logos-codegen" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59f80069600c0d66734f5ff52cc42f2dabd6b29d205f333d61fd7832e9e9963f" +dependencies = [ + "beef", + "fnv", + "lazy_static", + "proc-macro2", + "quote", + "regex-syntax", + "syn 2.0.118", +] + +[[package]] +name = "logos-derive" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24fb722b06a9dc12adb0963ed585f19fc61dc5413e6a9be9422ef92c091e731d" +dependencies = [ + "logos-codegen", +] + [[package]] name = "loom" version = "0.7.2" @@ -4855,6 +4956,22 @@ dependencies = [ "libc", ] +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + [[package]] name = "markup5ever" version = "0.38.0" @@ -4881,6 +4998,16 @@ version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "maybe-async" version = "0.2.11" @@ -4980,7 +5107,7 @@ dependencies = [ "mesh-llm-types", "model-artifact", "nostr-sdk", - "prost", + "prost 0.14.4", "rand 0.10.2", "rustls", "serde", @@ -5119,7 +5246,7 @@ dependencies = [ "opentelemetry", "opentelemetry-otlp", "opentelemetry_sdk", - "prost", + "prost 0.14.4", "rand 0.10.2", "regex-lite", "reqwest 0.12.28", @@ -5207,8 +5334,8 @@ source = "git+https://github.com/Mesh-LLM/mesh-llm.git?rev=f455d493a2ae82baf2a32 dependencies = [ "anyhow", "async-trait", - "prost", - "prost-build", + "prost 0.14.4", + "prost-build 0.14.4", "protoc-bin-vendored", "rmcp", "schemars 1.2.1", @@ -5243,7 +5370,7 @@ dependencies = [ "anyhow", "hex", "iroh", - "prost", + "prost 0.14.4", "serde_json", "sha2 0.10.9", ] @@ -5358,6 +5485,28 @@ dependencies = [ "tracing", ] +[[package]] +name = "miette" +version = "7.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f98efec8807c63c752b5bd61f862c165c115b0a35685bdcfd9238c7aeb592b7" +dependencies = [ + "cfg-if 1.0.4", + "miette-derive", + "unicode-width 0.1.14", +] + +[[package]] +name = "miette-derive" +version = "7.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db5b29714e950dbb20d5e6f74f9dcec4edbcc1067bb7f8ed198c097b8c1a818b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "mime" version = "0.3.17" @@ -5503,6 +5652,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "more-asserts" version = "0.3.1" @@ -5630,6 +5801,21 @@ dependencies = [ "tempfile", ] +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "ndk" version = "0.9.0" @@ -6611,7 +6797,7 @@ dependencies = [ "opentelemetry-http", "opentelemetry-proto", "opentelemetry_sdk", - "prost", + "prost 0.14.4", "reqwest 0.12.28", "thiserror 2.0.18", ] @@ -6626,7 +6812,7 @@ dependencies = [ "const-hex", "opentelemetry", "opentelemetry_sdk", - "prost", + "prost 0.14.4", "serde", "serde_json", "tonic", @@ -6692,6 +6878,24 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "ort" +version = "2.0.0-rc.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7de3af33d24a745ffb8fab904b13478438d1cd52868e6f17735ef6e1f8bf133" +dependencies = [ + "ndarray", + "ort-sys", + "smallvec", + "tracing", +] + +[[package]] +name = "ort-sys" +version = "2.0.0-rc.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90" + [[package]] name = "os_pipe" version = "1.2.3" @@ -6919,6 +7123,16 @@ dependencies = [ "pest", ] +[[package]] +name = "petgraph" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" +dependencies = [ + "fixedbitset 0.5.7", + "indexmap 2.14.0", +] + [[package]] name = "petgraph" version = "0.8.3" @@ -7183,6 +7397,15 @@ dependencies = [ "serde", ] +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + [[package]] name = "portmapper" version = "0.19.1" @@ -7393,6 +7616,16 @@ dependencies = [ "unarray", ] +[[package]] +name = "prost" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +dependencies = [ + "bytes", + "prost-derive 0.13.5", +] + [[package]] name = "prost" version = "0.14.4" @@ -7400,7 +7633,27 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" dependencies = [ "bytes", - "prost-derive", + "prost-derive 0.14.4", +] + +[[package]] +name = "prost-build" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" +dependencies = [ + "heck 0.4.1", + "itertools", + "log", + "multimap", + "once_cell", + "petgraph 0.7.1", + "prettyplease", + "prost 0.13.5", + "prost-types 0.13.5", + "regex", + "syn 2.0.118", + "tempfile", ] [[package]] @@ -7413,15 +7666,28 @@ dependencies = [ "itertools", "log", "multimap", - "petgraph", + "petgraph 0.8.3", "prettyplease", - "prost", - "prost-types", + "prost 0.14.4", + "prost-types 0.14.4", "regex", "syn 2.0.118", "tempfile", ] +[[package]] +name = "prost-derive" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +dependencies = [ + "anyhow", + "itertools", + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "prost-derive" version = "0.14.4" @@ -7435,13 +7701,35 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "prost-reflect" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5edd582b62f5cde844716e66d92565d7faf7ab1445c8cebce6e00fba83ddb2" +dependencies = [ + "logos", + "miette", + "once_cell", + "prost 0.13.5", + "prost-types 0.13.5", +] + +[[package]] +name = "prost-types" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16" +dependencies = [ + "prost 0.13.5", +] + [[package]] name = "prost-types" version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" dependencies = [ - "prost", + "prost 0.14.4", ] [[package]] @@ -7508,6 +7796,33 @@ version = "3.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95067976aca6421a523e491fce939a3e65249bac4b977adee0ee9771568e8aa3" +[[package]] +name = "protox" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f352af331bf637b8ecc720f7c87bf903d2571fa2e14a66e9b2558846864b54a" +dependencies = [ + "bytes", + "miette", + "prost 0.13.5", + "prost-reflect", + "prost-types 0.13.5", + "protox-parse", + "thiserror 1.0.69", +] + +[[package]] +name = "protox-parse" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3a462d115462c080ae000c29a47f0b3985737e5d3a995fcdbcaa5c782068dde" +dependencies = [ + "logos", + "miette", + "prost-types 0.13.5", + "thiserror 1.0.69", +] + [[package]] name = "pxfm" version = "0.1.30" @@ -7756,7 +8071,7 @@ dependencies = [ "thiserror 2.0.18", "unicode-segmentation", "unicode-truncate", - "unicode-width", + "unicode-width 0.2.2", ] [[package]] @@ -7819,7 +8134,7 @@ dependencies = [ "strum", "time", "unicode-segmentation", - "unicode-width", + "unicode-width 0.2.2", ] [[package]] @@ -7828,6 +8143,43 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "realfft" version = "3.5.0" @@ -8640,6 +8992,18 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cd0b0ec5f1c1ca621c432a25813d8d60c88abe6d3e08a3eb9cf37d97a0fe3d73" +[[package]] +name = "sentencepiece-model" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40b87bf750a8322c3236d7aa63c1f4a6862187d00d2d8b038e1dfe263bfe43ec" +dependencies = [ + "miette", + "prost 0.13.5", + "prost-build 0.13.5", + "protox", +] + [[package]] name = "serde" version = "1.0.228" @@ -9072,8 +9436,8 @@ name = "skippy-protocol" version = "0.73.1" source = "git+https://github.com/Mesh-LLM/mesh-llm.git?rev=f455d493a2ae82baf2a326e2d0fda351433b4b30#f455d493a2ae82baf2a326e2d0fda351433b4b30" dependencies = [ - "prost", - "prost-build", + "prost 0.14.4", + "prost-build 0.14.4", "protoc-bin-vendored", "serde", ] @@ -9252,6 +9616,18 @@ dependencies = [ "der", ] +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom 7.1.3", + "serde", + "unicode-segmentation", +] + [[package]] name = "sse-stream" version = "0.2.4" @@ -9633,7 +10009,7 @@ version = "1.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fce91f2f0ec87dff7e6bcbbeb267439aa1188703003c6055193c821487400432" dependencies = [ - "unicode-width", + "unicode-width 0.2.2", ] [[package]] @@ -10208,7 +10584,7 @@ dependencies = [ "anyhow", "base64 0.22.1", "bitflags 2.13.0", - "fancy-regex", + "fancy-regex 0.11.0", "filedescriptor", "finl_unicode", "fixedbitset 0.4.2", @@ -10371,6 +10747,39 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tokenizers" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str 0.9.1", + "dary_heap", + "derive_builder", + "esaxx-rs", + "fancy-regex 0.14.0", + "getrandom 0.3.4", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "paste", + "rand 0.9.4", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.18", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokio" version = "1.52.3" @@ -10704,7 +11113,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" dependencies = [ "bytes", - "prost", + "prost 0.14.4", "tonic", ] @@ -10885,7 +11294,7 @@ checksum = "b8765b90061cba6c22b5831f675da109ae5561588290f9fa2317adab2714d5a6" dependencies = [ "memchr", "nom 8.0.0", - "petgraph", + "petgraph 0.8.3", ] [[package]] @@ -11056,6 +11465,15 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + [[package]] name = "unicode-segmentation" version = "1.13.3" @@ -11070,9 +11488,15 @@ checksum = "16b380a1238663e5f8a691f9039c73e1cdae598a30e9855f541d29b08b53e9a5" dependencies = [ "itertools", "unicode-segmentation", - "unicode-width", + "unicode-width 0.2.2", ] +[[package]] +name = "unicode-width" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" + [[package]] name = "unicode-width" version = "0.2.2" @@ -11085,6 +11509,12 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "universal-hash" version = "0.5.1" diff --git a/desktop/src-tauri/Cargo.toml b/desktop/src-tauri/Cargo.toml index 324218a49d..1616f7ed67 100644 --- a/desktop/src-tauri/Cargo.toml +++ b/desktop/src-tauri/Cargo.toml @@ -76,6 +76,9 @@ bytes = "1" futures-util = "0.3" opus = "0.3" neteq = { version = "0.8", default-features = false } +ort = { version = "=2.0.0-rc.12", default-features = false, features = ["api-24", "ndarray", "std"] } +ort-sys = { version = "=2.0.0-rc.12", features = ["disable-linking"] } +rand = "0.10" serde = { version = "1", features = ["derive"] } serde_json = "1" serde_yaml = "0.9" @@ -116,6 +119,7 @@ image = { version = "0.25", default-features = false, features = ["jpeg", "png", zip = "8" flate2 = "1" sherpa-onnx = "1.12" +sentencepiece-model = "0.1" regex = "1" rusqlite = { version = "0.37", features = ["bundled"] } axum = "0.8" @@ -126,6 +130,7 @@ audioadapter-buffers = "3.0" tempfile = "3" strip-ansi-escapes = "0.2" tracing = "0.1" +tokenizers = { version = "0.22", default-features = false, features = ["fancy-regex"] } [dev-dependencies] # `test-util` enables tokio's paused-clock (`start_paused`) so the relay diff --git a/desktop/src-tauri/examples/pocket_quality_ab.rs b/desktop/src-tauri/examples/pocket_quality_ab.rs deleted file mode 100644 index 0c31f1c910..0000000000 --- a/desktop/src-tauri/examples/pocket_quality_ab.rs +++ /dev/null @@ -1,519 +0,0 @@ -//! 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 169ddf66c0..11c9ee4c7d 100644 --- a/desktop/src-tauri/src/huddle/models.rs +++ b/desktop/src-tauri/src/huddle/models.rs @@ -24,6 +24,10 @@ use std::sync::{Arc, Mutex, OnceLock}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; +use super::pocket::{ + april_model_info, PocketModelArtifact, APRIL_BUNDLE_ID, APRIL_MODEL_ID, APRIL_MODEL_REVISION, +}; + // ── Integrity verification ──────────────────────────────────────────────────── // // All model artifacts are verified against pinned SHA-256 hashes before @@ -38,19 +42,15 @@ 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 fp32 repackage. -/// -/// 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-2026-01-26/resolve/96d1e53ce3311ca6c2c6a35e2062d36b4cec6fa3"; +fn pocket_artifact_url(filename: &str) -> String { + format!( + "https://huggingface.co/{APRIL_MODEL_ID}/resolve/{APRIL_MODEL_REVISION}/onnx/{APRIL_BUNDLE_ID}/{filename}" + ) +} + +fn pocket_license_url() -> String { + format!("https://huggingface.co/{APRIL_MODEL_ID}/resolve/{APRIL_MODEL_REVISION}/onnx/LICENSE") +} /// Reference voice WAV: "Mary (f, conversation)" from the Kyutai TTS demo /// voice set — VCTK speaker p333, ai-coustics-enhanced. Pinned to @@ -64,20 +64,19 @@ const POCKET_HF_BASE: &str = const POCKET_REFERENCE_WAV_URL: &str = "https://huggingface.co/kyutai/tts-voices/resolve/323332d33f997de8394f24a193e1a76df720e01a/vctk/p333_023_enhanced.wav"; -/// SHA-256 hashes for individual Pocket TTS model files. -/// Computed from known-good pinned downloads. Update when upgrading model versions. -#[rustfmt::skip] -const TTS_FILE_HASHES: &[(&str, &str)] = &[ - ("decoder.onnx", "f267880fde6c58b17b0a8f3647eaf8dcfad321f833f32d583ebc2fb2d1a15f10"), - ("encoder.onnx", "e8f2f6d301ffb96e398b138a7dc6d3038622d236044636b73d920bab85890260"), - ("lm_flow.onnx", "79c013a554a54e63319c33c0cc8830cbbedc9b7e448ae7e26f7923ae11f9873e"), - ("lm_main.onnx", "255d1a9263c5abdf36034abfc19c11d21cc5f40f0f87d8361288e972cbd5c578"), - ("text_conditioner.onnx", "0b84e837d7bfaf2c896627b03e3f080320309f37f4fc7df7698c644f7ba5e6b1"), - ("vocab.json", "6fb646346cf931016f70c4921aab0900ce7a304b893cb02135c74e294abfea01"), - ("token_scores.json", "5be2f278caf9b9800741f0fd82bff677f4943ec764c356f907213434b622d958"), - ("LICENSE", "fe7b4ce83b8381cc5b216bbb4af73c570688d1b819c73bbaed8ca401f4677cd6"), - ("reference_sample.wav", "a35b0468382218e9f37a9a7494d1e4b74deaf18d7ced22265b4e325bb55c183f"), -]; +const TTS_LICENSE_ARTIFACT: PocketModelArtifact = PocketModelArtifact { + filename: "LICENSE", + sha256: "fe7b4ce83b8381cc5b216bbb4af73c570688d1b819c73bbaed8ca401f4677cd6", + size_bytes: 18_655, + quantized: false, +}; + +const TTS_REFERENCE_ARTIFACT: PocketModelArtifact = PocketModelArtifact { + filename: "reference_sample.wav", + sha256: "a35b0468382218e9f37a9a7494d1e4b74deaf18d7ced22265b4e325bb55c183f", + size_bytes: 639_084, + quantized: false, +}; // ── Model versioning ────────────────────────────────────────────────────────── // @@ -92,15 +91,8 @@ const TTS_FILE_HASHES: &[(&str, &str)] = &[ /// honest (each version tag identifies one specific set of model bytes). const STT_MODEL_VERSION: &str = "2"; -/// Model manifest version for Pocket TTS. Increment when upgrading model files. -/// Bumped "1" → "2" when the bundled reference voice changed from KevinAHM's -/// anonymous 16 kHz sample to Mary (VCTK p333, 32 kHz, ai-coustics-enhanced) -/// 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. -/// 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"; +/// Identifies the exact April INT8 asset set expected by readiness checks. +const TTS_MODEL_VERSION: &str = "4"; /// Filename for the version manifest written alongside model files. const MANIFEST_FILENAME: &str = ".buzz-model-manifest"; @@ -110,9 +102,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 (400 MB per file — largest is -/// `lm_main.onnx` at ~303 MB fp32). -const MAX_TTS_FILE_BYTES: u64 = 400 * 1024 * 1024; +/// Maximum expected Pocket TTS file size. The largest pinned INT8 artifact is +/// `flow_lm_main_int8.onnx` at 76,341,079 bytes. +const MAX_TTS_FILE_BYTES: u64 = 100 * 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 @@ -181,9 +173,9 @@ Original model by Kyutai: https://huggingface.co/kyutai/pocket-tts Paper: Charles, Roebel, et al., Pocket TTS (arXiv:2509.06926). 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-2026-01-26 +April 2026 ONNX export by KevinAHM: +https://huggingface.co/KevinAHM/pocket-tts-onnx +Pinned revision: 58a6d00cf13d239b6748cb0769f35c580a8f606c Bundled reference voice (reference_sample.wav): \"Mary (f, conversation)\" preset from the Kyutai TTS demo voice catalogue @@ -203,13 +195,14 @@ 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.onnx", - "encoder.onnx", - "lm_flow.onnx", - "lm_main.onnx", + "bundle.json", + "bos_before_voice.npy", + "flow_lm_main_int8.onnx", + "flow_lm_flow_int8.onnx", + "mimi_decoder_int8.onnx", + "mimi_encoder.onnx", "text_conditioner.onnx", - "vocab.json", - "token_scores.json", + "tokenizer.model", "LICENSE", "reference_sample.wav", TTS_LICENSE_FILE_NAME, @@ -404,6 +397,7 @@ struct ModelSlot { dir_name: &'static str, // subdir under ~/.buzz/models/ expected_files: &'static [&'static str], // files required for "ready" version: &'static str, // manifest version; increment to force re-download + expected_size: fn(&str) -> Option, status: Arc>, just_ready: Arc, // fires once when download completes } @@ -418,11 +412,17 @@ impl ModelSlot { dir_name, expected_files, version, + expected_size: |_| None, status: Arc::new(Mutex::new(ModelStatus::NotDownloaded)), just_ready: Arc::new(AtomicBool::new(false)), } } + fn with_expected_sizes(mut self, expected_size: fn(&str) -> Option) -> Self { + self.expected_size = expected_size; + self + } + fn model_dir(&self, models_dir: &Path) -> PathBuf { models_dir.join(self.dir_name) } @@ -432,7 +432,17 @@ impl ModelSlot { std::fs::read_to_string(dir.join(MANIFEST_FILENAME)) .map(|v| v.trim() == self.version) .unwrap_or(false) - && self.expected_files.iter().all(|f| dir.join(f).is_file()) + && self.expected_files.iter().all(|filename| { + let path = dir.join(filename); + path.is_file() + && (self.expected_size)(filename) + .map(|expected| { + path.metadata() + .map(|metadata| metadata.len() == expected) + .unwrap_or(false) + }) + .unwrap_or(true) + }) } fn dir_if_ready(&self, models_dir: &Path) -> Option { @@ -453,6 +463,39 @@ impl ModelSlot { self.just_ready.swap(false, Ordering::AcqRel) } + /// Recover or clean up the backup left by an interrupted atomic install. + fn recover_interrupted_install(&self, models_dir: &Path) { + let final_dir = self.model_dir(models_dir); + let backup_dir = final_dir.with_extension("old"); + if !backup_dir.exists() { + return; + } + if self.is_ready(models_dir) { + if let Err(error) = std::fs::remove_dir_all(&backup_dir) { + eprintln!( + "buzz-desktop: could not remove stale {} backup: {error}", + self.dir_name + ); + } + return; + } + if final_dir.exists() { + if let Err(error) = std::fs::remove_dir_all(&final_dir) { + eprintln!( + "buzz-desktop: could not remove incomplete {} install: {error}", + self.dir_name + ); + return; + } + } + if let Err(error) = std::fs::rename(&backup_dir, &final_dir) { + eprintln!( + "buzz-desktop: could not restore interrupted {} install: {error}", + self.dir_name + ); + } + } + /// Spawn a background download task if not already ready or downloading. fn start_download( &self, @@ -511,6 +554,9 @@ impl ModelSlot { )); } + std::fs::write(source_dir.join(MANIFEST_FILENAME), self.version) + .map_err(|e| format!("write model manifest: {e}"))?; + let final_dir = self.model_dir(models_dir); let backup_dir = final_dir.with_extension("old"); @@ -529,8 +575,6 @@ impl ModelSlot { return Err(format!("install new model: {e}")); } - std::fs::write(final_dir.join(MANIFEST_FILENAME), self.version) - .map_err(|e| format!("write model manifest: {e}"))?; let _ = tokio::fs::remove_dir_all(&backup_dir).await; if let Some(extra) = temp_cleanup { let _ = tokio::fs::remove_dir_all(extra).await; @@ -542,6 +586,25 @@ impl ModelSlot { } } +fn tts_expected_size(filename: &str) -> Option { + april_model_info() + .artifacts + .iter() + .find(|artifact| artifact.filename == filename) + .map(|artifact| artifact.size_bytes) + .or_else(|| { + [TTS_LICENSE_ARTIFACT, TTS_REFERENCE_ARTIFACT] + .iter() + .find(|artifact| artifact.filename == filename) + .map(|artifact| artifact.size_bytes) + }) +} + +fn tts_model_slot() -> ModelSlot { + ModelSlot::new(TTS_MODEL_DIR_NAME, TTS_EXPECTED_FILES, TTS_MODEL_VERSION) + .with_expected_sizes(tts_expected_size) +} + // ── ModelManager ────────────────────────────────────────────────────────────── /// Manages download and location of STT/TTS model files. @@ -561,11 +624,13 @@ impl ModelManager { /// Returns `None` if the home directory cannot be resolved. pub fn new() -> Option { let models_dir = dirs::home_dir()?.join(".buzz").join("models"); - Some(Self { + let manager = Self { models_dir, stt: ModelSlot::new(STT_MODEL_DIR_NAME, STT_EXPECTED_FILES, STT_MODEL_VERSION), - tts: ModelSlot::new(TTS_MODEL_DIR_NAME, TTS_EXPECTED_FILES, TTS_MODEL_VERSION), - }) + tts: tts_model_slot(), + }; + manager.tts.recover_interrupted_install(&manager.models_dir); + Some(manager) } // ── STT accessors ──────────────────────────────────────────────────────── @@ -638,7 +703,7 @@ impl ModelManager { } } - /// Start a background Pocket TTS download (~189 MB). No-op if already ready or downloading. + /// Start a background Pocket TTS download. No-op if already ready or downloading. pub fn start_tts_download(&self, http_client: reqwest::Client) { let manager = self.clone(); self.tts.start_download( @@ -754,8 +819,8 @@ impl ModelManager { /// Download and verify the Pocket TTS model files from HuggingFace. /// /// Downloads files into `~/.buzz/models/pocket-tts/`: - /// - five ONNX sessions (Pocket TTS + Mimi codec) - /// - `vocab.json` / `token_scores.json` for sherpa-onnx text conditioning + /// - five ONNX sessions selected by the April INT8 bundle + /// - bundle metadata, SentencePiece tokenizer, and learned voice BOS /// - upstream `LICENSE` plus Buzz's `MODEL_LICENSE.txt` attribution sidecar /// - `reference_sample.wav` as the bundled default voice /// @@ -768,24 +833,18 @@ impl ModelManager { let temp_dir = self.models_dir.join("pocket-tts.tmp"); fresh_temp_dir(&temp_dir).await?; - let model_files = [ - "decoder.onnx", - "encoder.onnx", - "lm_flow.onnx", - "lm_main.onnx", - "text_conditioner.onnx", - "vocab.json", - "token_scores.json", - "LICENSE", - ]; - let mut downloads: Vec<(String, &'static str)> = model_files + let mut downloads: Vec<(String, PocketModelArtifact)> = april_model_info() + .artifacts .iter() - .map(|filename| (format!("{POCKET_HF_BASE}/{filename}"), *filename)) + .copied() + .map(|artifact| (pocket_artifact_url(artifact.filename), artifact)) .collect(); - downloads.push((POCKET_REFERENCE_WAV_URL.to_string(), "reference_sample.wav")); + downloads.push((pocket_license_url(), TTS_LICENSE_ARTIFACT)); + downloads.push((POCKET_REFERENCE_WAV_URL.to_string(), TTS_REFERENCE_ARTIFACT)); let total_files = downloads.len() as u32; - for (i, (url, filename)) in downloads.iter().enumerate() { + for (i, (url, artifact)) in downloads.iter().enumerate() { + let filename = artifact.filename; eprintln!("buzz-desktop: downloading Pocket TTS {filename} from {url}"); let response = fetch_url(&http_client, url, filename) @@ -822,16 +881,19 @@ impl ModelManager { })?; eprintln!("buzz-desktop: downloaded {bytes} bytes ({filename}), wrote to disk"); - let expected = TTS_FILE_HASHES - .iter() - .find(|(n, _)| *n == *filename) - .map(|(_, hash)| *hash) - .ok_or_else(|| format!("missing expected hash for Pocket TTS file: {filename}"))?; + if bytes != artifact.size_bytes { + let _ = tokio::fs::remove_dir_all(&temp_dir).await; + return Err(format!( + "Pocket TTS {filename} size check failed: expected {} bytes, got {bytes}", + artifact.size_bytes + )); + } let actual = sha256_file(&dest).await?; - if actual != expected { + if actual != artifact.sha256 { let _ = tokio::fs::remove_dir_all(&temp_dir).await; return Err(format!( - "Pocket TTS {filename} integrity check failed: expected {expected}, got {actual}" + "Pocket TTS {filename} integrity check failed: expected {}, got {actual}", + artifact.sha256 )); } @@ -931,24 +993,5 @@ pub fn is_tts_ready() -> bool { } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn tts_readiness_requires_license_sidecar() { - let temp = tempfile::tempdir().expect("tempdir"); - let slot = ModelSlot::new(TTS_MODEL_DIR_NAME, TTS_EXPECTED_FILES, TTS_MODEL_VERSION); - let model_dir = temp.path().join(TTS_MODEL_DIR_NAME); - std::fs::create_dir_all(&model_dir).expect("create model dir"); - - for file in TTS_EXPECTED_FILES { - std::fs::write(model_dir.join(file), b"test").expect("write expected file"); - } - std::fs::write(model_dir.join(MANIFEST_FILENAME), TTS_MODEL_VERSION).expect("manifest"); - - assert!(slot.is_ready(temp.path())); - - std::fs::remove_file(model_dir.join(TTS_LICENSE_FILE_NAME)).expect("remove sidecar"); - assert!(!slot.is_ready(temp.path())); - } -} +#[path = "models_tests.rs"] +mod tests; diff --git a/desktop/src-tauri/src/huddle/models_tests.rs b/desktop/src-tauri/src/huddle/models_tests.rs new file mode 100644 index 0000000000..4bcb4081e0 --- /dev/null +++ b/desktop/src-tauri/src/huddle/models_tests.rs @@ -0,0 +1,148 @@ +use super::*; + +fn create_ready_model_dir(root: &Path) -> PathBuf { + let model_dir = root.join(TTS_MODEL_DIR_NAME); + std::fs::create_dir_all(&model_dir).expect("create model dir"); + for file in TTS_EXPECTED_FILES { + let path = model_dir.join(file); + let handle = std::fs::File::create(path).expect("create expected file"); + if let Some(size) = tts_expected_size(file) { + handle.set_len(size).expect("size expected file"); + } else { + std::fs::write(model_dir.join(file), b"test").expect("write expected file"); + } + } + std::fs::write(model_dir.join(MANIFEST_FILENAME), TTS_MODEL_VERSION).expect("manifest"); + model_dir +} + +#[test] +fn expected_files_match_april_int8_metadata() { + let mut expected = april_model_info() + .artifacts + .iter() + .map(|artifact| artifact.filename) + .chain([ + TTS_LICENSE_ARTIFACT.filename, + TTS_REFERENCE_ARTIFACT.filename, + TTS_LICENSE_FILE_NAME, + ]) + .collect::>(); + expected.sort_unstable(); + let mut actual = TTS_EXPECTED_FILES.to_vec(); + actual.sort_unstable(); + + assert_eq!(actual, expected); + assert!(!actual.contains(&"flow_lm_main.onnx")); + assert!(!actual.contains(&"flow_lm_flow.onnx")); + assert!(!actual.contains(&"mimi_decoder.onnx")); +} + +#[test] +fn tts_readiness_requires_license_sidecar() { + let temp = tempfile::tempdir().expect("tempdir"); + let slot = tts_model_slot(); + let model_dir = create_ready_model_dir(temp.path()); + + assert!(slot.is_ready(temp.path())); + + std::fs::remove_file(model_dir.join(TTS_LICENSE_FILE_NAME)).expect("remove sidecar"); + assert!(!slot.is_ready(temp.path())); +} + +#[test] +fn tts_readiness_rejects_truncated_pinned_artifact() { + let temp = tempfile::tempdir().expect("tempdir"); + let slot = tts_model_slot(); + let model_dir = create_ready_model_dir(temp.path()); + let artifact = april_model_info().artifacts[0]; + + std::fs::OpenOptions::new() + .write(true) + .open(model_dir.join(artifact.filename)) + .expect("open artifact") + .set_len(artifact.size_bytes - 1) + .expect("truncate artifact"); + + assert!(!slot.is_ready(temp.path())); +} + +#[test] +fn january_cache_is_not_ready_for_april_int8() { + let temp = tempfile::tempdir().expect("tempdir"); + let slot = tts_model_slot(); + let model_dir = temp.path().join(TTS_MODEL_DIR_NAME); + std::fs::create_dir_all(&model_dir).expect("create model dir"); + for file in [ + "decoder.onnx", + "encoder.onnx", + "lm_flow.onnx", + "lm_main.onnx", + "text_conditioner.onnx", + "vocab.json", + "token_scores.json", + "LICENSE", + "reference_sample.wav", + TTS_LICENSE_FILE_NAME, + ] { + std::fs::write(model_dir.join(file), b"january").expect("write January file"); + } + std::fs::write(model_dir.join(MANIFEST_FILENAME), "3").expect("manifest"); + + assert!(!slot.is_ready(temp.path())); +} + +#[test] +fn interrupted_install_restores_backup_when_destination_is_missing() { + let temp = tempfile::tempdir().expect("tempdir"); + let slot = tts_model_slot(); + let backup_dir = temp.path().join("pocket-tts.old"); + std::fs::create_dir_all(&backup_dir).expect("create backup"); + std::fs::write(backup_dir.join("sentinel"), b"previous").expect("write sentinel"); + + slot.recover_interrupted_install(temp.path()); + + assert_eq!( + std::fs::read(temp.path().join(TTS_MODEL_DIR_NAME).join("sentinel")) + .expect("restored sentinel"), + b"previous" + ); + assert!(!backup_dir.exists()); +} + +#[test] +fn interrupted_install_replaces_incomplete_destination_with_backup() { + let temp = tempfile::tempdir().expect("tempdir"); + let slot = tts_model_slot(); + let model_dir = temp.path().join(TTS_MODEL_DIR_NAME); + let backup_dir = temp.path().join("pocket-tts.old"); + std::fs::create_dir_all(&model_dir).expect("create incomplete destination"); + std::fs::write(model_dir.join("incomplete"), b"april").expect("write incomplete file"); + std::fs::create_dir_all(&backup_dir).expect("create backup"); + std::fs::write(backup_dir.join("sentinel"), b"previous").expect("write sentinel"); + + slot.recover_interrupted_install(temp.path()); + + assert_eq!( + std::fs::read(model_dir.join("sentinel")).expect("restored sentinel"), + b"previous" + ); + assert!(!model_dir.join("incomplete").exists()); + assert!(!backup_dir.exists()); +} + +#[test] +fn ready_destination_removes_stale_backup() { + let temp = tempfile::tempdir().expect("tempdir"); + let slot = tts_model_slot(); + let model_dir = create_ready_model_dir(temp.path()); + let backup_dir = temp.path().join("pocket-tts.old"); + std::fs::create_dir_all(&backup_dir).expect("create backup"); + std::fs::write(backup_dir.join("sentinel"), b"previous").expect("write sentinel"); + + slot.recover_interrupted_install(temp.path()); + + assert!(slot.is_ready(temp.path())); + assert!(model_dir.exists()); + assert!(!backup_dir.exists()); +} diff --git a/desktop/src-tauri/src/huddle/pocket.rs b/desktop/src-tauri/src/huddle/pocket.rs index ee1faf928a..2154a25c22 100644 --- a/desktop/src-tauri/src/huddle/pocket.rs +++ b/desktop/src-tauri/src/huddle/pocket.rs @@ -1,184 +1,52 @@ -//! Pocket TTS engine wrapper around sherpa-onnx's `OfflineTts`. +//! April 2026 Pocket TTS engine for Buzz Desktop. //! -//! 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". +//! The `english_2026-04` bundle uses SentencePiece tokenization, a learned +//! voice BOS embedding, recurrent FlowLM state, and stateful Mimi decoding. +//! Buzz selects the upstream three-graph INT8 variant while retaining the +//! full-precision Mimi encoder and text conditioner specified by that variant. //! //! ## Attribution //! -//! - **Model**: Kyutai *Pocket TTS* — Charles, Roebel, et al., 2026. -//! arXiv:2509.06926. Original repository: . -//! Licensed CC-BY-4.0. -//! - **Mimi neural codec**: Kyutai, bundled in the same release. CC-BY-4.0. -//! - **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 -//! (f, conversation)" preset from the Kyutai TTS demo -//! (), which maps to `vctk/p333_023_enhanced.wav` -//! in . CC-BY-4.0, base recording -//! from the VCTK corpus, enhanced by ai-coustics. -//! -//! Buzz ships these files unmodified; see the on-disk `MODEL_LICENSE.txt` -//! sidecar written by `huddle::models` during install for the canonical -//! CC-BY-4.0 §3(a)(1) attribution block. -//! -//! ## Engine-module contract (see `huddle::tts`) +//! - Pocket TTS and Mimi: Kyutai, CC-BY-4.0. +//! - ONNX export: KevinAHM/pocket-tts-onnx, CC-BY-4.0. +//! - Reference voice: Kyutai's Mary preset (VCTK p333), CC-BY-4.0. //! -//! `pocket.rs` exposes a fixed surface used by `tts.rs`. Mirroring this -//! contract is what lets the TTS pipeline stay engine-agnostic: -//! -//! - `SAMPLE_RATE: u32` — engine output sample rate in Hz. -//! - `DEFAULT_VOICE: &str` — default voice name (without extension). -//! - `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)` -//! → `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. +//! `huddle::models` writes the complete attribution beside the cached bytes. -use std::collections::HashMap; use std::path::{Path, PathBuf}; +use std::sync::Mutex; + +use sherpa_onnx::Wave; -use sherpa_onnx::{GenerationConfig, OfflineTts, OfflineTtsConfig, Wave}; +#[path = "pocket_april.rs"] +mod pocket_april; +#[path = "pocket_models.rs"] +mod pocket_models; -// ── Engine-module contract: public consts ───────────────────────────────────── +use pocket_april::{prepare_april_prompt, AprilPocketTts}; +pub(crate) use pocket_models::{ + april_model_info, PocketModelArtifact, APRIL_BUNDLE_ID, APRIL_MODEL_ID, APRIL_MODEL_REVISION, +}; -/// Pocket TTS emits 24 kHz mono PCM. Matches the previous Kokoro output rate, -/// so the rodio sink and inter-sentence silence buffer in `tts.rs` remain valid. +/// Pocket TTS emits 24 kHz mono PCM. pub const SAMPLE_RATE: u32 = 24_000; -/// Name (without extension) of the bundled reference voice. The model directory -/// is expected to contain `.` after install. +/// Bundled reference voice name without its extension. pub const DEFAULT_VOICE: &str = "reference_sample"; -/// Voice files for Pocket TTS are reference audio (WAV). Distinct from the -/// Kokoro `.bin` style vectors — the model conditions on raw waveform samples, -/// not a precomputed embedding, so the extension change is honest. +/// Pocket voice files are reference WAVs. pub const VOICE_FILE_EXT: &str = "wav"; -// ── Tuning ──────────────────────────────────────────────────────────────────── - -/// Single-threaded ONNX execution for predictable CPU contention with the STT -/// pipeline. Matches `STT_NUM_THREADS` in `stt.rs`; raise only if a benchmark -/// argues for it. -const TTS_NUM_THREADS: i32 = 1; - -/// LRU cache size for cloned voice embeddings inside the sherpa-onnx engine. -/// We bind to one voice per pipeline today, but the upstream example uses 16 -/// and the cost is negligible — keep room for future multi-voice support. -const VOICE_EMBEDDING_CACHE_CAPACITY: i32 = 16; - -/// Pocket TTS is a consistency-based LM. Generation quality saturates at one -/// denoising step — the upstream `GenerationConfig` default of 5 multiplies -/// synthesis time by ~5× with no audible benefit on this model. -const SYNTH_NUM_STEPS: i32 = 1; - -/// Leave the generated audio's silences untouched (1.0 is the identity). -/// -/// sherpa-onnx's `ScaleSilence` (`offline-tts.cc`) is *not* pre/post padding -/// control: it finds every interior silence run ≥ 0.2 s (|s| ≤ 0.01) and -/// multiplies its length by this factor. The previous value of 0.0 — set -/// under the mistaken belief it disabled lead-in/lead-out padding — deleted -/// every natural pause inside an utterance: clause breaks, breaths, the gap -/// after a comma. Words slammed together and endings cut abruptly. The -/// reference Pocket TTS pipeline does not post-process silence at all; -/// 1.0 restores parity. -const SYNTH_SILENCE_SCALE: f32 = 1.0; - -/// sherpa-onnx upstream default for `max_frames` (LM steps), in -/// `offline-tts-pocket-impl.h:Generate`. 500 steps ≈ 40 s of audio at the -/// Mimi 12.5 Hz frame rate. Referenced only by the regression test below; -/// production code path never raises (or even reads) this value — we just -/// leave sherpa-onnx's own default in place by not setting the override. -#[cfg(test)] -const SHERPA_ONNX_MAX_FRAMES_DEFAULT: i32 = 500; - -/// Tight `max_frames` we ask for on short, padded prompts to bound the -/// original "monster breathing" runaway. 100 LM steps ≈ 8 s of audio — -/// roomy for any one-to-four-word utterance the user is likely to elicit -/// while still well short of the 40 s upstream default. Chosen with slack so -/// we never *truncate* a legitimate short reply. -const SHORT_PROMPT_MAX_FRAMES: i32 = 100; - -/// Word-count threshold (inclusive) below which we pad the prompt with -/// leading spaces and cap `max_frames` tighter than the upstream default. -/// Matches upstream `pocket_tts.models.tts_model.prepare_text_prompt`. Above -/// this threshold we leave sherpa-onnx's own defaults in place — overriding -/// them caused the "first 'yep' is just static" regression seen on -/// 2026-05-18, where dropping `frames_after_eos` below the upstream default -/// of 3 clipped the leading audio of multi-clause sentences. -const SHORT_PROMPT_WORD_THRESHOLD: usize = 4; - -/// Number of leading spaces prepended to short prompts. The upstream Python -/// uses exactly 8 — keep parity rather than tuning blindly. -/// -/// This is upstream's *only* mitigation for the FlowLM cold-start smear on -/// short utterances (kyutai-labs/pocket-tts #91, #70): the autoregressive -/// generation has a 2–3 step "settle" period where the first phoneme can be -/// smeared. A previous revision added a sacrificial `". . "` prefix plus an -/// amplitude-threshold trim to strip the rendered prefix from the output — -/// but the trim's absolute threshold (0.02 against raw peaks of ~0.076) sat -/// in soft-onset territory and could eat real word starts, and its tuning -/// was calibrated against `silence_scale = 0.0` audio. Deleted in favour of -/// upstream parity: accept the occasional smeared first syllable rather -/// than risk trimming real speech. -const SHORT_PROMPT_PAD_SPACES: usize = 8; - -/// sherpa-onnx's documented `frames_after_eos` default. We deliberately do -/// *not* override this knob — the previous attempt to bump it for short -/// inputs and lower it for long inputs lowered it below the upstream default -/// of 3, which clipped the leading audio of multi-clause sentences (the -/// "first 'yep' is static" regression). The constant exists only for the -/// regression test below. Source: `offline-tts-pocket-impl.h:Generate`. -#[cfg(test)] -const SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT: i32 = 3; - -// ── ONNX file names (five Pocket TTS sessions plus two JSON tables) ─────────── +const TTS_NUM_THREADS: usize = 1; -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.onnx"; -const FILE_TEXT_COND: &str = "text_conditioner.onnx"; -const FILE_VOCAB: &str = "vocab.json"; -const FILE_TOKEN_SCORES: &str = "token_scores.json"; - -// ── Voice style ─────────────────────────────────────────────────────────────── - -/// Loaded reference voice — normalised f32 PCM samples plus their sample rate. -/// -/// Pocket TTS takes a reference waveform per generation call (not a -/// precomputed style embedding), so we keep the samples in memory and clone -/// the small `Vec` into each `GenerationConfig` rather than re-reading the -/// WAV from disk on every sentence. +/// Loaded reference voice samples and their original sample rate. #[derive(Debug, Clone)] pub struct VoiceStyle { samples: Vec, sample_rate: i32, } -/// Load a reference voice WAV from disk. -/// -/// Accepts any sample rate sherpa-onnx's `Wave::read` can decode — Pocket TTS -/// resamples internally using `reference_sample_rate`. The bundled -/// `reference_sample.wav` ("Mary" — VCTK p333, enhanced) is 32 kHz mono. +/// Load a Pocket reference voice WAV from disk. pub fn load_voice_style(path: &Path) -> Result { let path_str = path .to_str() @@ -195,199 +63,46 @@ pub fn load_voice_style(path: &Path) -> Result { }) } -// ── Engine ──────────────────────────────────────────────────────────────────── - -/// Pocket TTS engine handle. Cheap to construct (one `OfflineTts::create` -/// call). Owned by the TTS worker thread for the lifetime of a huddle session. -/// -/// `OfflineTts` does not implement `Debug`, so we don't derive it here — the -/// pipeline only needs to move the engine into the worker thread and call -/// `synth_chunk` on it, never to print it. +/// Resident April INT8 Pocket TTS engine. pub struct PocketTts { - inner: OfflineTts, + inner: Mutex, } -/// Build the Pocket TTS engine from the model directory installed by -/// `huddle::models`. Returns `Err` if any expected ONNX or JSON file is -/// missing — readiness is normally enforced by `is_tts_ready` upstream, but -/// the check is repeated here so a manually-modified model dir produces a -/// clear error string instead of an opaque sherpa-onnx `None`. +/// Load Buzz Desktop's pinned April INT8 model. pub fn load_text_to_speech(model_dir: &str) -> Result { let dir = PathBuf::from(model_dir); - for name in [ - FILE_LM_MAIN, - FILE_LM_FLOW, - FILE_ENCODER, - FILE_DECODER, - FILE_TEXT_COND, - FILE_VOCAB, - FILE_TOKEN_SCORES, - ] { - let p = dir.join(name); - if !p.is_file() { - return Err(format!("missing Pocket TTS file: {}", p.display())); - } - } - - let to_str = |name: &str| -> String { dir.join(name).to_string_lossy().into_owned() }; - - // Build the config by mutating defaults — mirrors `stt.rs` and stays - // resilient if sherpa-onnx adds unrelated model-family fields. - let mut cfg = OfflineTtsConfig::default(); - cfg.model.pocket.lm_main = Some(to_str(FILE_LM_MAIN)); - cfg.model.pocket.lm_flow = Some(to_str(FILE_LM_FLOW)); - cfg.model.pocket.encoder = Some(to_str(FILE_ENCODER)); - cfg.model.pocket.decoder = Some(to_str(FILE_DECODER)); - cfg.model.pocket.text_conditioner = Some(to_str(FILE_TEXT_COND)); - cfg.model.pocket.vocab_json = Some(to_str(FILE_VOCAB)); - cfg.model.pocket.token_scores_json = Some(to_str(FILE_TOKEN_SCORES)); - cfg.model.pocket.voice_embedding_cache_capacity = VOICE_EMBEDDING_CACHE_CAPACITY; - cfg.model.num_threads = TTS_NUM_THREADS; - // Explicit — defaults are not part of the API contract, and noisy debug - // logging in release builds would be expensive on every synthesized chunk. - cfg.model.debug = false; - - let inner = OfflineTts::create(&cfg) - .ok_or_else(|| "OfflineTts::create returned None for Pocket TTS".to_string())?; - Ok(PocketTts { inner }) -} - -// ── Prompt preparation ──────────────────────────────────────────────────────── - -/// Result of [`prepare_pocket_prompt`]: a synthesizer-ready prompt plus the -/// per-call generation overrides derived from the original text. -/// -/// `None` for either override means "leave sherpa-onnx's documented default -/// in place". The pipeline only sets `max_frames` (and only for short -/// padded inputs) so it can bound the original "monster breathing" runaway -/// without disturbing the rest of the LM sampling envelope. -#[derive(Debug, Clone, PartialEq)] -pub(crate) struct PreparedPrompt { - /// Text to hand to `OfflineTts::generate_with_config`. Capitalized, - /// punctuation-terminated, and (for short inputs) left-padded with - /// spaces — upstream's mitigation for the FlowLM cold-start smear. - pub text: String, - /// Value to pass via `GenerationConfig.extra["max_frames"]`, or `None` to - /// keep the upstream default of 500 LM steps. We only override on short - /// padded prompts where we have a tight expectation on output length. - pub max_frames: Option, -} - -/// Mirror of the *text-preparation* half of upstream -/// `pocket_tts.models.tts_model.prepare_text_prompt`. Sherpa-onnx's C++ -/// Pocket TTS impl does not run these preparation steps, so short / -/// unpunctuated / lowercase inputs can trigger up to 40 s of runaway -/// generation when the EOS logit never crosses its threshold. We replicate -/// the upstream Python recipe here: -/// -/// 1. Collapse interior whitespace (already done by `preprocess_for_tts`, but -/// cheap to re-check after sentence splitting). -/// 2. Capitalize the first letter. -/// 3. Append `.` if the text doesn't end in punctuation. -/// 4. If fewer than five words, prepend `SHORT_PROMPT_PAD_SPACES` spaces -/// (upstream's cold-start mitigation — see the constant's docstring) and -/// return a tight [`SHORT_PROMPT_MAX_FRAMES`] cap so the LM can't run -/// away if EOS still doesn't fire. -/// -/// We do **not** override `frames_after_eos` — sherpa-onnx's default of 3 -/// is what we want. An earlier version set it to 1 on long inputs, which -/// clipped the leading audio of multi-clause sentences ("first 'yep' is -/// just static" regression). Tests `prepare_prompt_never_lowers_frames_…` -/// lock this in. -/// -/// Returns `None` only if the input is empty after trimming — caller should -/// skip synthesis in that case. -pub(crate) fn prepare_pocket_prompt(input: &str) -> Option { - let trimmed = input.trim(); - if trimmed.is_empty() { - return None; - } - - // Collapse stray double-spaces / embedded newlines that may slip past - // `preprocess_for_tts` when sentences are spliced back together. - let mut cleaned = String::with_capacity(trimmed.len()); - let mut last_was_space = false; - for ch in trimmed.chars() { - let is_ws = ch.is_whitespace(); - if is_ws { - if !last_was_space { - cleaned.push(' '); - } - last_was_space = true; - } else { - cleaned.push(ch); - last_was_space = false; + for artifact in april_model_info().artifacts { + let path = dir.join(artifact.filename); + if !path.is_file() { + return Err(format!( + "incomplete Pocket TTS {} INT8 bundle: missing {}", + APRIL_BUNDLE_ID, + path.display() + )); } } - - // Capitalize first character. Uses `to_uppercase` (multi-codepoint safe). - let first = cleaned.chars().next().expect("cleaned non-empty above"); - if first.is_lowercase() { - let upper: String = first.to_uppercase().collect(); - let mut iter = cleaned.chars(); - iter.next(); - cleaned = upper + iter.as_str(); - } - - // Ensure terminal punctuation. Anything not in `.!?;:,` gets a period. - // The upstream Python only checks `isalnum` → period, but for our agent - // text we already may end in `!` `?` `.` etc. — treat any of those as OK. - let last = cleaned - .chars() - .next_back() - .expect("cleaned non-empty above"); - if !matches!(last, '.' | '!' | '?' | ';' | ':' | ',') { - cleaned.push('.'); - } - - // Word count of the *cleaned but not padded* text — padding is whitespace - // only and would just lie to the threshold check below. - let word_count = cleaned.split_whitespace().count(); - - let (final_text, max_frames) = if word_count <= SHORT_PROMPT_WORD_THRESHOLD { - let mut padded = String::with_capacity(cleaned.len() + SHORT_PROMPT_PAD_SPACES); - for _ in 0..SHORT_PROMPT_PAD_SPACES { - padded.push(' '); - } - padded.push_str(&cleaned); - (padded, Some(SHORT_PROMPT_MAX_FRAMES)) - } else { - // For everything ≥5 words, fall back to upstream defaults. Overriding - // these is what caused the "first 'yep' is static" regression — the - // upstream LM has been tuned for `frames_after_eos = 3` and - // `max_frames = 500`, and there's no clear win in second-guessing. - (cleaned, None) - }; - - Some(PreparedPrompt { - text: final_text, - max_frames, - }) -} - -/// Build the `GenerationConfig.extra` HashMap from a [`PreparedPrompt`]. -/// -/// Centralised so the regression test below can assert that we **never** -/// emit a `frames_after_eos` override — the previous attempt to override -/// that knob (setting it to 1 for ≥5-word inputs) clipped the leading -/// audio of multi-clause sentences (the "first 'yep' is static" bug on -/// 2026-05-18). The upstream sherpa-onnx default of 3 is what we want, and -/// the right way to keep it is to not set it at all. -fn build_generation_extra(prepared: &PreparedPrompt) -> Option> { - prepared.max_frames.map(|mf| { - let mut h: HashMap = HashMap::with_capacity(1); - h.insert("max_frames".to_string(), serde_json::Value::from(mf)); - h + Ok(PocketTts { + inner: Mutex::new(AprilPocketTts::load(&dir, TTS_NUM_THREADS)?), }) } impl PocketTts { - /// Synthesise `text` with the given reference voice. + /// 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, 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. /// - /// `_lang` and `_steps` are accepted for API compatibility with the - /// previous Kokoro engine. Pocket TTS infers language from the input text - /// directly and is a one-step consistency model. Returns an empty buffer - /// for whitespace-only input. + /// 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, @@ -395,57 +110,21 @@ impl PocketTts { style: &VoiceStyle, _steps: usize, ) -> Result, String> { - // Mirror upstream pocket-tts prompt prep — without this short or - // unpunctuated inputs can cause the LM's EOS logit to never trip, - // producing up to 40 s of "monster breathing" garbage on the first - // utterance. See `prepare_pocket_prompt` for the full recipe. - let prepared = match prepare_pocket_prompt(text) { - Some(p) => p, - None => return Ok(Vec::new()), + let Some(prepared) = prepare_april_prompt(text) else { + return Ok(Vec::new()); }; - - // Per-call generation hints sherpa-onnx forwards to - // `offline-tts-pocket-impl.h`. We only override `max_frames`, and - // only for short padded prompts where we have a tight expectation - // on output length — that bounds the original runaway without - // disturbing the rest of the LM sampling envelope. See - // `prepare_pocket_prompt` docs for the regression history. - let extra = build_generation_extra(&prepared); - - let cfg = GenerationConfig { - 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() - }; - - // No progress callback — synthesis is fast enough that returning the - // whole buffer at once keeps the lookahead pipelining in `tts.rs` - // simple. `None:: bool>` pins the callback type for the - // `generate_with_config` generic parameter. - let audio = self + let mut engine = self .inner - .generate_with_config(&prepared.text, &cfg, None:: bool>) - .ok_or_else(|| { - format!( - "Pocket TTS synthesis failed for text ({} chars)", - prepared.text.len() - ) - })?; - - let sample_rate = audio.sample_rate(); - if sample_rate != SAMPLE_RATE as i32 { - eprintln!( - "buzz-desktop: Pocket TTS returned unexpected sample rate {sample_rate}Hz \ - (expected {SAMPLE_RATE}Hz); playback speed may be wrong" - ); + .lock() + .map_err(|_| "Pocket TTS engine lock poisoned".to_string())?; + let chunks = engine.split_prompt(&prepared)?; + let mut samples = Vec::new(); + for chunk in chunks { + let prepared = prepare_april_prompt(&chunk) + .ok_or_else(|| "Pocket TTS prompt chunk became empty".to_string())?; + samples.extend(engine.synth_chunk(&prepared, style)?); } - - Ok(audio.samples().to_vec()) + Ok(samples) } } @@ -453,202 +132,35 @@ impl PocketTts { mod tests { use super::*; - // ── prepare_pocket_prompt ──────────────────────────────────────────────── - - #[test] - fn prepare_prompt_returns_none_for_empty_input() { - assert!(prepare_pocket_prompt("").is_none()); - assert!(prepare_pocket_prompt(" ").is_none()); - assert!(prepare_pocket_prompt("\n\t ").is_none()); - } - - /// Helper: the exact leading sequence prepended to every short prompt — - /// 8 spaces of padding (upstream's cold-start mitigation). - /// Centralising this keeps the assertions readable. - fn short_prefix() -> String { - " ".repeat(SHORT_PROMPT_PAD_SPACES) - } - - #[test] - fn prepare_prompt_pads_and_capitalizes_one_word() { - // The "yep" case Tyler hit in production — bare lowercase one-word - // utterance with no punctuation. Must be padded with the short-prompt - // space pad, capitalized, terminated, with a tight `max_frames` cap - // to bound runaway gen. - let out = prepare_pocket_prompt("yep").expect("non-empty"); - assert_eq!(out.text, format!("{}Yep.", short_prefix())); - assert_eq!(out.max_frames, Some(SHORT_PROMPT_MAX_FRAMES)); - const { - assert!( - SHORT_PROMPT_MAX_FRAMES < SHERPA_ONNX_MAX_FRAMES_DEFAULT, - "short cap must be tighter than the upstream default" - ); - } - } - - #[test] - fn prepare_prompt_preserves_existing_punctuation() { - let out = prepare_pocket_prompt("yes!").expect("non-empty"); - assert_eq!(out.text, format!("{}Yes!", short_prefix())); // exclamation kept - let out = prepare_pocket_prompt("really?").expect("non-empty"); - assert_eq!(out.text, format!("{}Really?", short_prefix())); - } - #[test] - fn prepare_prompt_threshold_is_inclusive_at_four_words() { - // 4 words = short (padded + tight max_frames); 5 words = long - // (no padding, no overrides — upstream defaults stand). - let four = prepare_pocket_prompt("one two three four").expect("non-empty"); - assert_eq!( - four.text, - format!("{}One two three four.", short_prefix()), - "four-word input should get exactly the space pad" - ); - assert_eq!(four.max_frames, Some(SHORT_PROMPT_MAX_FRAMES)); - - let five = prepare_pocket_prompt("one two three four five").expect("non-empty"); - assert!( - !five.text.starts_with(' '), - "five-word input should NOT be padded" - ); - assert_eq!( - five.max_frames, None, - "long inputs must leave sherpa-onnx's max_frames default in place" - ); - } - - #[test] - fn prepare_prompt_does_not_pad_long_text() { - let long = "This is a longer sentence that the model should handle just fine."; - let out = prepare_pocket_prompt(long).expect("non-empty"); - assert!(!out.text.starts_with(' ')); - assert_eq!(out.max_frames, None); - assert!(out.text.ends_with('.')); + fn desktop_model_is_april_int8_only() { + let info = april_model_info(); + assert_eq!(info.max_token_per_chunk, 50); + assert_eq!(info.sample_rate, SAMPLE_RATE); + assert!(info + .artifacts + .iter() + .any(|artifact| artifact.filename == "flow_lm_main_int8.onnx")); + assert!(!info + .artifacts + .iter() + .any(|artifact| artifact.filename == "flow_lm_main.onnx")); } #[test] - fn prepare_prompt_collapses_whitespace() { - let out = prepare_pocket_prompt("Hello world\n\nfriend").expect("non-empty"); - // 3 words → short → padded. Interior whitespace collapsed. - assert_eq!(out.text, format!("{}Hello world friend.", short_prefix())); - } - - #[test] - fn prepare_prompt_does_not_double_capitalize_already_uppercase() { - let out = prepare_pocket_prompt("HELLO there").expect("non-empty"); - assert_eq!(out.text, format!("{}HELLO there.", short_prefix())); - } - - #[test] - fn prepare_prompt_handles_non_ascii_first_letter() { - // Cyrillic lowercase 'д' → uppercase 'Д'. Must not panic / produce - // mojibake. - let out = prepare_pocket_prompt("дa").expect("non-empty"); - assert!(out.text.contains("Дa.")); - } - - /// REGRESSION GUARD: short prompts must receive *only* whitespace - /// padding — no sacrificial text. A previous revision prepended a - /// `". . "` cold-start absorber and trimmed the rendered audio back out - /// with an amplitude threshold that could eat soft word onsets. If - /// non-whitespace ever reappears in the pad, the synth output will - /// contain audio for text the user never wrote. - #[test] - fn prepare_prompt_pad_is_whitespace_only() { - let out = prepare_pocket_prompt("I'm happy.").expect("non-empty"); - let pad_len = out.text.len() - "I'm happy.".len(); - assert!( - out.text[..pad_len].chars().all(|c| c == ' '), - "short-prompt pad must be spaces only, got {:?}", - &out.text[..pad_len] - ); - assert_eq!(out.text, format!("{}I'm happy.", short_prefix())); - } - - // ── build_generation_extra ─────────────────────────────────────────────── - // - // These tests pin down a behaviour we've now regressed twice on: - // 1) Not padding/punctuating short inputs → 40 s of "monster breathing" - // (pre-773a2a1). - // 2) Setting `frames_after_eos = 1` on long inputs → clipped leading - // audio of multi-clause sentences, e.g. "Yep, I can hear you. …" - // came out as a static burst (the 773a2a1 regression Tyler hit on - // 2026-05-18 ~14:30 UTC). - // - // The contract we enforce going forward: we **only** override - // `max_frames`, and only for ≤4-word inputs. Every other knob is left - // at sherpa-onnx's documented default (notably `frames_after_eos = 3`). - - #[test] - fn build_extra_short_prompt_sets_only_max_frames() { - let prepared = prepare_pocket_prompt("yep").expect("non-empty"); - let extra = build_generation_extra(&prepared).expect("short prompts get extra"); - // Exactly one key — `max_frames` — and nothing else. - assert_eq!(extra.len(), 1, "extra has unexpected keys: {extra:?}"); - assert_eq!( - extra.get("max_frames"), - Some(&serde_json::Value::from(SHORT_PROMPT_MAX_FRAMES)) - ); - assert!( - !extra.contains_key("frames_after_eos"), - "frames_after_eos must never be set — upstream default of {SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT} is what we want" - ); - } - - #[test] - fn build_extra_long_prompt_is_none() { - // ≥5 words: no extras at all. This is the key fix for the "first - // 'yep' in 'Yep, I can hear you. …' is static" regression — we - // were previously forcing `frames_after_eos = 1` on this path. - let prepared = prepare_pocket_prompt("Yep, I can hear you.").expect("non-empty"); - assert_eq!( - build_generation_extra(&prepared), - None, - "long prompts must not override any LM knob" - ); - } - - #[test] - fn build_extra_never_lowers_frames_after_eos_for_any_word_count() { - // Sweep a range of prompt lengths and assert the `extra` map (when - // present) never carries a `frames_after_eos` override that's lower - // than the upstream sherpa-onnx default. Implemented as a structural - // check — we just never set the key — but worth a property test in - // case someone reintroduces the override in the future. - let prompts: &[&str] = &[ - "hi", - "hi there", - "yes please", - "one two three four", - "one two three four five", - "a slightly longer reply, hopefully fine", - "This is a multi-clause sentence. It has two parts.", - "really really really really really long prompt with lots of words just to be sure", - ]; - for &p in prompts { - let prepared = prepare_pocket_prompt(p).expect("non-empty"); - if let Some(extra) = build_generation_extra(&prepared) { - if let Some(v) = extra.get("frames_after_eos") { - let n = v.as_i64().expect("frames_after_eos should be int"); - assert!( - n >= SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT as i64, - "prompt {p:?} set frames_after_eos={n}, below upstream default of {SHERPA_ONNX_FRAMES_AFTER_EOS_DEFAULT}" - ); - } - } - } - } - - #[test] - fn short_prompt_max_frames_is_below_upstream_default() { - // Sanity: the override only ever *lowers* the cap, never raises it. - const { - assert!(SHORT_PROMPT_MAX_FRAMES < SHERPA_ONNX_MAX_FRAMES_DEFAULT); - } - // …and is still large enough for a one-to-four-word reply. At Mimi's - // 12.5 Hz frame rate, 100 frames = 8 s, which is roomy. - const { - assert!(SHORT_PROMPT_MAX_FRAMES >= 50, "would risk truncation"); - } + #[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"] + fn production_api_emits_non_silent_april_int8_pcm() { + 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"); + let samples = engine + .synth_chunk("Bright birds begin beside the bay.", "en", &style, 1) + .expect("synthesize through the production API"); + + assert!(!samples.is_empty()); + assert!(samples.iter().all(|sample| sample.is_finite())); + assert!(samples.iter().any(|sample| sample.abs() > 1.0e-6)); } } diff --git a/desktop/src-tauri/src/huddle/pocket_april.rs b/desktop/src-tauri/src/huddle/pocket_april.rs new file mode 100644 index 0000000000..43826df5c9 --- /dev/null +++ b/desktop/src-tauri/src/huddle/pocket_april.rs @@ -0,0 +1,940 @@ +//! Native ONNX loader for Pocket TTS `english_2026-04`. +//! +//! The bundle uses SentencePiece, prepends a learned BOS voice embedding, and +//! describes recurrent state tensors in `bundle.json`. This module supplies +//! that frontend and state loop while reusing the ONNX Runtime linked by the +//! Desktop speech stack. + +use std::borrow::Cow; +use std::f32::consts::TAU; +use std::fs; +use std::path::{Path, PathBuf}; + +use ort::session::{Session, SessionInputValue}; +use ort::value::{DynValue, Tensor}; +use rand::{Rng, RngExt}; +use sentencepiece_model::SentencePieceModel; +use serde::Deserialize; +use sherpa_onnx::LinearResampler; +use tokenizers::models::unigram::Unigram; +use tokenizers::pre_tokenizers::metaspace::{Metaspace, PrependScheme}; +use tokenizers::Tokenizer; + +use super::VoiceStyle; + +const FILE_BUNDLE: &str = "bundle.json"; +const FILE_MIMI_ENCODER: &str = "mimi_encoder.onnx"; +const FILE_TEXT_CONDITIONER: &str = "text_conditioner.onnx"; +const FILE_FLOW_MAIN_INT8: &str = "flow_lm_main_int8.onnx"; +const FILE_FLOW_INT8: &str = "flow_lm_flow_int8.onnx"; +const FILE_MIMI_DECODER_INT8: &str = "mimi_decoder_int8.onnx"; + +const MODEL_LANGUAGE: &str = "english_2026-04"; +const DEFAULT_TEMPERATURE: f32 = 0.7; +const EOS_LOGIT_THRESHOLD: f32 = -4.0; +const DECODER_CHUNK_FRAMES: usize = 12; +const TOKENS_PER_SECOND_ESTIMATE: f32 = 3.0; +const GENERATION_SECONDS_PADDING: f32 = 2.0; + +#[derive(Debug, Deserialize)] +struct Bundle { + schema_version: u32, + language: String, + sample_rate: usize, + frame_rate: f32, + samples_per_frame: usize, + latent_dim: usize, + conditioning_dim: usize, + insert_bos_before_voice: bool, + pad_with_spaces_for_short_inputs: bool, + remove_semicolons: bool, + model_recommended_frames_after_eos: Option, + max_token_per_chunk: usize, + tokenizer_file: String, + bos_before_voice_file: String, + flow_lm_state_manifest: Vec, + mimi_state_manifest: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +struct StateSpec { + input_name: String, + output_name: String, + dtype: StateDtype, + shape: Vec, + fill: StateFill, +} + +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(rename_all = "lowercase")] +enum StateDtype { + #[serde(rename = "float32")] + Float32, + #[serde(rename = "int64")] + Int64, + Bool, +} + +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(rename_all = "lowercase")] +enum StateFill { + Empty, + Nan, + Ones, + Zeros, +} + +struct StateValue { + spec: StateSpec, + value: DynValue, +} + +struct CachedVoice { + samples_ptr: usize, + samples_len: usize, + sample_rate: i32, + embeddings: Vec, +} + +pub(crate) struct AprilPocketTts { + bundle: Bundle, + tokenizer: Tokenizer, + bos_embedding: Vec, + mimi_encoder: Session, + text_conditioner: Session, + flow_main: Session, + flow: Session, + mimi_decoder: Session, + cached_voice: Option, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct AprilPreparedPrompt { + pub(crate) text: String, + pub(crate) frames_after_eos: usize, +} + +pub(crate) fn prepare_april_prompt(input: &str) -> Option { + let trimmed = input.trim(); + if trimmed.is_empty() { + return None; + } + + let mut cleaned = String::with_capacity(trimmed.len()); + let mut last_was_space = false; + for ch in trimmed.chars() { + if ch.is_whitespace() { + if !last_was_space { + cleaned.push(' '); + } + last_was_space = true; + } else { + cleaned.push(ch); + last_was_space = false; + } + } + + let first = cleaned.chars().next().expect("cleaned non-empty above"); + if first.is_lowercase() { + let upper: String = first.to_uppercase().collect(); + let mut iter = cleaned.chars(); + iter.next(); + cleaned = upper + iter.as_str(); + } + + let last = cleaned + .chars() + .next_back() + .expect("cleaned non-empty above"); + if last.is_alphanumeric() { + cleaned.push('.'); + } + + let word_count = cleaned.split_whitespace().count(); + Some(AprilPreparedPrompt { + text: cleaned, + // Mirror the bundle's upstream heuristic: three generated frames plus + // two trailing frames for short prompts, one plus two otherwise. + frames_after_eos: if word_count <= 4 { 5 } else { 3 }, + }) +} + +impl AprilPocketTts { + pub(crate) fn load(dir: &Path, num_threads: usize) -> Result { + if num_threads == 0 { + return Err("Pocket TTS num_threads must be at least 1".to_string()); + } + let bundle_path = dir.join(FILE_BUNDLE); + let bundle: Bundle = serde_json::from_slice( + &fs::read(&bundle_path) + .map_err(|err| format!("read {}: {err}", bundle_path.display()))?, + ) + .map_err(|err| format!("parse {}: {err}", bundle_path.display()))?; + + if bundle.schema_version != 2 { + return Err(format!( + "unsupported Pocket TTS bundle schema {} in {}", + bundle.schema_version, + bundle_path.display() + )); + } + if bundle.language != MODEL_LANGUAGE { + return Err(format!( + "expected Pocket TTS language {MODEL_LANGUAGE}, got {}", + bundle.language + )); + } + if bundle.sample_rate != 24_000 + || bundle.frame_rate != 12.5 + || bundle.samples_per_frame != 1_920 + || bundle.latent_dim != 32 + || bundle.conditioning_dim != 1024 + { + return Err(format!( + "unexpected Pocket TTS dimensions: sample_rate={}, frame_rate={}, samples_per_frame={}, latent_dim={}, conditioning_dim={}", + bundle.sample_rate, + bundle.frame_rate, + bundle.samples_per_frame, + bundle.latent_dim, + bundle.conditioning_dim + )); + } + if !bundle.insert_bos_before_voice { + return Err("April Pocket TTS bundle must insert BOS before voice".to_string()); + } + if bundle.pad_with_spaces_for_short_inputs + || bundle.remove_semicolons + || bundle.model_recommended_frames_after_eos.is_some() + || bundle.max_token_per_chunk != 50 + { + return Err("unsupported April Pocket TTS prompt-policy metadata".to_string()); + } + + let tokenizer_path = dir.join(&bundle.tokenizer_file); + let tokenizer = load_tokenizer(&tokenizer_path)?; + let bos_path = dir.join(&bundle.bos_before_voice_file); + let bos_embedding = read_npy_f32(&bos_path)?; + if bos_embedding.len() != bundle.conditioning_dim { + return Err(format!( + "{} has {} values; expected {}", + bos_path.display(), + bos_embedding.len(), + bundle.conditioning_dim + )); + } + + let flow_main = FILE_FLOW_MAIN_INT8; + let flow = FILE_FLOW_INT8; + let mimi_decoder = FILE_MIMI_DECODER_INT8; + + Ok(Self { + // The INT8 layout quantizes only the three generation graphs; + // voice encoding and text conditioning remain full precision. + mimi_encoder: load_session(dir.join(FILE_MIMI_ENCODER), num_threads)?, + text_conditioner: load_session(dir.join(FILE_TEXT_CONDITIONER), num_threads)?, + flow_main: load_session(dir.join(flow_main), num_threads)?, + flow: load_session(dir.join(flow), num_threads)?, + mimi_decoder: load_session(dir.join(mimi_decoder), num_threads)?, + bundle, + tokenizer, + bos_embedding, + cached_voice: None, + }) + } + + pub(crate) fn split_prompt( + &self, + prepared: &AprilPreparedPrompt, + ) -> Result, String> { + if self.token_count(&prepared.text)? <= self.bundle.max_token_per_chunk { + return Ok(vec![prepared.text.clone()]); + } + + let mut chunks = Vec::new(); + let mut current = String::new(); + for word in prepared.text.split_whitespace() { + let candidate = if current.is_empty() { + word.to_string() + } else { + format!("{current} {word}") + }; + if self.prepared_token_count(&candidate)? <= self.bundle.max_token_per_chunk { + current = candidate; + continue; + } + if !current.is_empty() { + chunks.push(std::mem::take(&mut current)); + } + + if self.prepared_token_count(word)? <= self.bundle.max_token_per_chunk { + current = word.to_string(); + continue; + } + + let mut fragment = String::new(); + for ch in word.chars() { + let candidate = format!("{fragment}{ch}"); + if !fragment.is_empty() + && self.prepared_token_count(&candidate)? > self.bundle.max_token_per_chunk + { + chunks.push(std::mem::take(&mut fragment)); + } + fragment.push(ch); + } + current = fragment; + } + if !current.is_empty() { + chunks.push(current); + } + + chunks + .into_iter() + .map(|text| { + let chunk = prepare_april_prompt(&text) + .ok_or_else(|| "Pocket TTS prompt chunk became empty".to_string())?; + let token_count = self.token_count(&chunk.text)?; + if token_count > self.bundle.max_token_per_chunk { + return Err(format!( + "Pocket TTS prompt chunk has {token_count} tokens; maximum is {}", + self.bundle.max_token_per_chunk + )); + } + Ok(chunk.text) + }) + .collect() + } + + pub(crate) fn synth_chunk( + &mut self, + prepared: &AprilPreparedPrompt, + style: &VoiceStyle, + ) -> Result, String> { + let voice_embeddings = self.voice_embeddings(style)?; + let mut flow_state = self.condition_voice(&voice_embeddings)?; + let token_ids = self + .tokenizer + .encode(prepared.text.as_str(), false) + .map_err(|err| format!("tokenize Pocket TTS prompt: {err}"))? + .get_ids() + .iter() + .copied() + .map(i64::from) + .collect::>(); + if token_ids.is_empty() { + return Ok(Vec::new()); + } + if token_ids.len() > self.bundle.max_token_per_chunk { + return Err(format!( + "Pocket TTS prompt has {} tokens; split_text_into_chunks maximum is {}", + token_ids.len(), + self.bundle.max_token_per_chunk + )); + } + + let token_count = token_ids.len(); + let text_embeddings = self.text_embeddings(token_ids)?; + self.run_flow_main_prefix(&text_embeddings, &mut flow_state)?; + let 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) + } + + fn prepared_token_count(&self, text: &str) -> Result { + let prepared = prepare_april_prompt(text) + .ok_or_else(|| "Pocket TTS prompt chunk became empty".to_string())?; + self.token_count(&prepared.text) + } + + fn token_count(&self, text: &str) -> Result { + Ok(self + .tokenizer + .encode(text, false) + .map_err(|err| format!("tokenize Pocket TTS prompt: {err}"))? + .get_ids() + .len()) + } + + fn voice_embeddings(&mut self, style: &VoiceStyle) -> Result, String> { + let key = ( + style.samples.as_ptr() as usize, + style.samples.len(), + style.sample_rate, + ); + if let Some(cached) = &self.cached_voice { + if (cached.samples_ptr, cached.samples_len, cached.sample_rate) == key { + return Ok(cached.embeddings.clone()); + } + } + + let samples = if style.sample_rate == self.bundle.sample_rate as i32 { + style.samples.clone() + } else { + LinearResampler::create(style.sample_rate, self.bundle.sample_rate as i32) + .ok_or_else(|| { + format!( + "create Pocket TTS resampler {}Hz -> {}Hz", + style.sample_rate, self.bundle.sample_rate + ) + })? + .resample(&style.samples, true) + }; + let audio = Tensor::from_array(( + vec![1_i64, 1, samples.len() as i64], + samples.into_boxed_slice(), + )) + .map_err(ort_error("create voice audio tensor"))?; + let outputs = self + .mimi_encoder + .run(ort::inputs!["audio" => audio]) + .map_err(ort_error("run Mimi encoder"))?; + let (_, encoded) = outputs[0] + .try_extract_tensor::() + .map_err(ort_error("extract Mimi encoder output"))?; + if !encoded.len().is_multiple_of(self.bundle.conditioning_dim) { + return Err(format!( + "Mimi encoder returned {} values, not divisible by {}", + encoded.len(), + self.bundle.conditioning_dim + )); + } + let mut embeddings = + Vec::with_capacity(self.bos_embedding.len().saturating_add(encoded.len())); + embeddings.extend_from_slice(&self.bos_embedding); + embeddings.extend_from_slice(encoded); + self.cached_voice = Some(CachedVoice { + samples_ptr: key.0, + samples_len: key.1, + sample_rate: key.2, + embeddings: embeddings.clone(), + }); + Ok(embeddings) + } + + fn condition_voice(&mut self, embeddings: &[f32]) -> Result, String> { + let frames = embeddings.len() / self.bundle.conditioning_dim; + let sequence = Tensor::::new( + &ort::memory::Allocator::default(), + [1_i64, 0, self.bundle.latent_dim as i64], + ) + .map_err(ort_error("create empty voice sequence"))?; + let text_embeddings = Tensor::from_array(( + vec![1_i64, frames as i64, self.bundle.conditioning_dim as i64], + embeddings.to_vec().into_boxed_slice(), + )) + .map_err(ort_error("create voice embedding tensor"))?; + let mut state = initialize_state(&self.bundle.flow_lm_state_manifest)?; + let mut inputs = vec![ + (Cow::Borrowed("sequence"), SessionInputValue::from(sequence)), + ( + Cow::Borrowed("text_embeddings"), + SessionInputValue::from(text_embeddings), + ), + ]; + append_state_inputs(&mut inputs, &state); + let mut outputs = self + .flow_main + .run(inputs) + .map_err(ort_error("condition Pocket TTS voice"))?; + replace_state_from_outputs(&mut state, &mut outputs)?; + Ok(state) + } + + fn text_embeddings(&mut self, token_ids: Vec) -> Result, String> { + let tokens = Tensor::from_array(( + vec![1_i64, token_ids.len() as i64], + token_ids.into_boxed_slice(), + )) + .map_err(ort_error("create token tensor"))?; + let outputs = self + .text_conditioner + .run(ort::inputs!["token_ids" => tokens]) + .map_err(ort_error("run text conditioner"))?; + let (_, embeddings) = outputs[0] + .try_extract_tensor::() + .map_err(ort_error("extract text embeddings"))?; + Ok(embeddings.to_vec()) + } + + fn run_flow_main_prefix( + &mut self, + text_embeddings: &[f32], + state: &mut [StateValue], + ) -> Result<(), String> { + if !text_embeddings + .len() + .is_multiple_of(self.bundle.conditioning_dim) + { + return Err(format!( + "text conditioner returned {} values, not divisible by {}", + text_embeddings.len(), + self.bundle.conditioning_dim + )); + } + let frames = text_embeddings.len() / self.bundle.conditioning_dim; + let sequence = Tensor::::new( + &ort::memory::Allocator::default(), + [1_i64, 0, self.bundle.latent_dim as i64], + ) + .map_err(ort_error("create empty text sequence"))?; + let text_embeddings = Tensor::from_array(( + vec![1_i64, frames as i64, self.bundle.conditioning_dim as i64], + text_embeddings.to_vec().into_boxed_slice(), + )) + .map_err(ort_error("create text embedding tensor"))?; + let mut inputs = vec![ + (Cow::Borrowed("sequence"), SessionInputValue::from(sequence)), + ( + Cow::Borrowed("text_embeddings"), + SessionInputValue::from(text_embeddings), + ), + ]; + append_state_inputs(&mut inputs, state); + let mut outputs = self + .flow_main + .run(inputs) + .map_err(ort_error("prime Pocket TTS text state"))?; + replace_state_from_outputs(state, &mut outputs) + } + + fn generate_latents( + &mut self, + max_frames: usize, + frames_after_eos: usize, + state: &mut [StateValue], + ) -> Result, String> { + let mut current = vec![f32::NAN; self.bundle.latent_dim]; + let mut latents = Vec::with_capacity(max_frames * self.bundle.latent_dim); + let mut eos_step = None; + let mut rng = rand::rng(); + + for step in 0..max_frames { + let sequence = Tensor::from_array(( + vec![1_i64, 1, self.bundle.latent_dim as i64], + current.clone().into_boxed_slice(), + )) + .map_err(ort_error("create latent input"))?; + let text_embeddings = Tensor::::new( + &ort::memory::Allocator::default(), + [1_i64, 0, self.bundle.conditioning_dim as i64], + ) + .map_err(ort_error("create empty text input"))?; + let mut inputs = vec![ + (Cow::Borrowed("sequence"), SessionInputValue::from(sequence)), + ( + Cow::Borrowed("text_embeddings"), + SessionInputValue::from(text_embeddings), + ), + ]; + append_state_inputs(&mut inputs, state); + let mut outputs = self + .flow_main + .run(inputs) + .map_err(ort_error("run Pocket TTS Flow LM"))?; + let conditioning = outputs[0] + .try_extract_tensor::() + .map_err(ort_error("extract Flow LM conditioning"))? + .1 + .to_vec(); + let eos_logit = outputs[1] + .try_extract_tensor::() + .map_err(ort_error("extract Flow LM EOS logit"))? + .1 + .first() + .copied() + .ok_or_else(|| "Flow LM returned empty EOS logit".to_string())?; + replace_state_from_outputs(state, &mut outputs)?; + + if eos_logit > EOS_LOGIT_THRESHOLD && eos_step.is_none() { + eos_step = Some(step); + } + if eos_step.is_some_and(|eos| step >= eos + frames_after_eos) { + break; + } + + let mut noise = + normal_noise(&mut rng, self.bundle.latent_dim, DEFAULT_TEMPERATURE.sqrt()); + let conditioning = Tensor::from_array(( + vec![1_i64, self.bundle.conditioning_dim as i64], + conditioning.into_boxed_slice(), + )) + .map_err(ort_error("create flow conditioning"))?; + let s = Tensor::from_array((vec![1_i64, 1], vec![0.0_f32].into_boxed_slice())) + .map_err(ort_error("create flow start tensor"))?; + let t = Tensor::from_array((vec![1_i64, 1], vec![1.0_f32].into_boxed_slice())) + .map_err(ort_error("create flow end tensor"))?; + let x = Tensor::from_array(( + vec![1_i64, self.bundle.latent_dim as i64], + noise.clone().into_boxed_slice(), + )) + .map_err(ort_error("create flow noise tensor"))?; + let outputs = self + .flow + .run(ort::inputs![ + "c" => conditioning, + "s" => s, + "t" => t, + "x" => x, + ]) + .map_err(ort_error("run Pocket TTS flow"))?; + let flow = outputs[0] + .try_extract_tensor::() + .map_err(ort_error("extract Pocket TTS flow"))? + .1; + if flow.len() != noise.len() { + return Err(format!( + "flow returned {} values; expected {}", + flow.len(), + noise.len() + )); + } + for (sample, delta) in noise.iter_mut().zip(flow) { + *sample += *delta; + } + current.clone_from(&noise); + latents.extend_from_slice(&noise); + } + Ok(latents) + } + + fn decode_latents(&mut self, latents: &[f32]) -> Result, String> { + if latents.is_empty() { + return Ok(Vec::new()); + } + if !latents.len().is_multiple_of(self.bundle.latent_dim) { + return Err(format!( + "latent buffer has {} values, not divisible by {}", + latents.len(), + self.bundle.latent_dim + )); + } + let frame_count = latents.len() / self.bundle.latent_dim; + let mut state = initialize_state(&self.bundle.mimi_state_manifest)?; + let mut audio = Vec::new(); + + for start in (0..frame_count).step_by(DECODER_CHUNK_FRAMES) { + let end = (start + DECODER_CHUNK_FRAMES).min(frame_count); + let values = + latents[start * self.bundle.latent_dim..end * self.bundle.latent_dim].to_vec(); + let latent = Tensor::from_array(( + vec![1_i64, (end - start) as i64, self.bundle.latent_dim as i64], + values.into_boxed_slice(), + )) + .map_err(ort_error("create Mimi latent tensor"))?; + let mut inputs = vec![(Cow::Borrowed("latent"), SessionInputValue::from(latent))]; + append_state_inputs(&mut inputs, &state); + let mut outputs = self + .mimi_decoder + .run(inputs) + .map_err(ort_error("run Mimi decoder"))?; + let samples = outputs[0] + .try_extract_tensor::() + .map_err(ort_error("extract Mimi audio"))? + .1; + audio.extend_from_slice(samples); + replace_state_from_outputs(&mut state, &mut outputs)?; + } + Ok(audio) + } +} + +fn load_session(path: PathBuf, num_threads: usize) -> Result { + if !path.is_file() { + return Err(format!("missing Pocket TTS file: {}", path.display())); + } + Session::builder() + .map_err(ort_error("create ONNX session builder"))? + .with_intra_threads(num_threads) + .map_err(|err| format!("configure ONNX intra-op threads: {err}"))? + .with_inter_threads(1) + .map_err(|err| format!("configure ONNX inter-op threads: {err}"))? + .commit_from_file(&path) + .map_err(|err| format!("load {}: {err}", path.display())) +} + +fn load_tokenizer(path: &Path) -> Result { + let sentencepiece = SentencePieceModel::from_file(path) + .map_err(|err| format!("load {}: {err}", path.display()))?; + let trainer = sentencepiece + .trainer() + .ok_or_else(|| format!("{} has no SentencePiece trainer metadata", path.display()))?; + let normalizer = sentencepiece.normalizer().ok_or_else(|| { + format!( + "{} has no SentencePiece normalizer metadata", + path.display() + ) + })?; + if normalizer.name() != "identity" { + return Err(format!( + "{} uses unsupported SentencePiece normalizer {:?}", + path.display(), + normalizer.name() + )); + } + + let vocab = sentencepiece + .pieces() + .iter() + .map(|piece| (piece.piece().to_owned(), f64::from(piece.score()))) + .collect(); + let mut tokenizer = Tokenizer::new( + Unigram::from( + vocab, + Some(trainer.unk_id() as usize), + trainer.byte_fallback(), + ) + .map_err(|err| format!("construct tokenizer from {}: {err}", path.display()))?, + ); + // SentencePiece's identity normalizer still escapes spaces as U+2581 and + // prepends one marker to the input before unigram segmentation. + tokenizer.with_pre_tokenizer(Some(Metaspace::new('▁', PrependScheme::Always, false))); + Ok(tokenizer) +} + +fn initialize_state(specs: &[StateSpec]) -> Result, String> { + specs + .iter() + .cloned() + .map(|spec| { + let len = shape_len(&spec.shape)?; + let value = match spec.dtype { + StateDtype::Float32 => { + let fill = match spec.fill { + StateFill::Nan => f32::NAN, + StateFill::Empty | StateFill::Zeros => 0.0, + StateFill::Ones => 1.0, + }; + if len == 0 { + Tensor::::new(&ort::memory::Allocator::default(), spec.shape.clone()) + .map_err(ort_error("create empty float state tensor"))? + .into_dyn() + } else { + Tensor::from_array((spec.shape.clone(), vec![fill; len].into_boxed_slice())) + .map_err(ort_error("create float state tensor"))? + .into_dyn() + } + } + StateDtype::Int64 => { + let fill = i64::from(matches!(spec.fill, StateFill::Ones)); + if len == 0 { + Tensor::::new(&ort::memory::Allocator::default(), spec.shape.clone()) + .map_err(ort_error("create empty integer state tensor"))? + .into_dyn() + } else { + Tensor::from_array((spec.shape.clone(), vec![fill; len].into_boxed_slice())) + .map_err(ort_error("create integer state tensor"))? + .into_dyn() + } + } + StateDtype::Bool => { + let fill = matches!(spec.fill, StateFill::Ones); + if len == 0 { + Tensor::::new(&ort::memory::Allocator::default(), spec.shape.clone()) + .map_err(ort_error("create empty bool state tensor"))? + .into_dyn() + } else { + Tensor::from_array((spec.shape.clone(), vec![fill; len].into_boxed_slice())) + .map_err(ort_error("create bool state tensor"))? + .into_dyn() + } + } + }; + Ok(StateValue { spec, value }) + }) + .collect() +} + +fn append_state_inputs<'a>( + inputs: &mut Vec<(Cow<'a, str>, SessionInputValue<'a>)>, + state: &'a [StateValue], +) { + for value in state { + inputs.push(( + Cow::Borrowed(value.spec.input_name.as_str()), + SessionInputValue::from(&value.value), + )); + } +} + +fn replace_state_from_outputs( + state: &mut [StateValue], + outputs: &mut ort::session::SessionOutputs<'_>, +) -> Result<(), String> { + for value in state { + value.value = outputs + .remove(&value.spec.output_name) + .ok_or_else(|| format!("missing state output {}", value.spec.output_name))?; + } + Ok(()) +} + +fn shape_len(shape: &[i64]) -> Result { + shape.iter().try_fold(1_usize, |len, &dim| { + let dim = usize::try_from(dim).map_err(|_| format!("negative state dimension {dim}"))?; + len.checked_mul(dim) + .ok_or_else(|| format!("state shape overflows usize: {shape:?}")) + }) +} + +fn estimate_max_frames(token_count: usize, frame_rate: f32) -> usize { + ((token_count as f32 / TOKENS_PER_SECOND_ESTIMATE + GENERATION_SECONDS_PADDING) * frame_rate) + .ceil() as usize +} + +fn normal_noise(rng: &mut impl Rng, len: usize, std_dev: f32) -> Vec { + let mut out = Vec::with_capacity(len); + while out.len() < len { + let u1 = rng.random::().max(f32::MIN_POSITIVE); + let u2 = rng.random::(); + let radius = (-2.0_f32 * u1.ln()).sqrt() * std_dev; + out.push(radius * (TAU * u2).cos()); + if out.len() < len { + out.push(radius * (TAU * u2).sin()); + } + } + out +} + +fn read_npy_f32(path: &Path) -> Result, String> { + let bytes = fs::read(path).map_err(|err| format!("read {}: {err}", path.display()))?; + if bytes.len() < 10 || &bytes[..6] != b"\x93NUMPY" { + return Err(format!("{} is not a NumPy array", path.display())); + } + let major = bytes[6]; + let header_len_bytes = match major { + 1 => 2, + 2 | 3 => 4, + _ => { + return Err(format!( + "unsupported NumPy version {major} in {}", + path.display() + )) + } + }; + let header_start = 8 + header_len_bytes; + if bytes.len() < header_start { + return Err(format!("truncated NumPy header in {}", path.display())); + } + let header_len = if header_len_bytes == 2 { + u16::from_le_bytes([bytes[8], bytes[9]]) as usize + } else { + u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize + }; + let data_start = header_start + .checked_add(header_len) + .ok_or_else(|| format!("NumPy header overflow in {}", path.display()))?; + if data_start > bytes.len() { + return Err(format!("truncated NumPy data in {}", path.display())); + } + let header = std::str::from_utf8(&bytes[header_start..data_start]) + .map_err(|err| format!("invalid NumPy header in {}: {err}", path.display()))?; + if !(header.contains("'descr': ' impl FnOnce(ort::Error) -> String { + move |err| format!("{context}: {err}") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn shape_len_supports_empty_state_dimensions() { + assert_eq!(shape_len(&[1, 128, 0]).expect("shape"), 0); + assert_eq!(shape_len(&[2, 1, 8, 1000, 64]).expect("shape"), 1_024_000); + } + + #[test] + fn normal_noise_has_requested_length() { + let mut rng = rand::rng(); + assert_eq!(normal_noise(&mut rng, 1, 1.0).len(), 1); + assert_eq!(normal_noise(&mut rng, 32, 1.0).len(), 32); + } + + #[test] + fn generation_frame_estimate_scales_with_token_count() { + assert_eq!(estimate_max_frames(3, 12.5), 38); + assert_eq!(estimate_max_frames(300, 12.5), 1_275); + } + + #[test] + #[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"] + fn tokenizer_matches_sentencepiece_reference_including_unknown_words() { + let dir = std::env::var("BUZZ_POCKET_TEST_MODEL_DIR") + .expect("set BUZZ_POCKET_TEST_MODEL_DIR to the verified April bundle"); + let tokenizer = + load_tokenizer(&Path::new(&dir).join("tokenizer.model")).expect("load April tokenizer"); + let cases: &[(&str, &[u32])] = &[ + ("Yep.", &[2462, 263]), + ("Hello there.", &[2994, 310, 263]), + ( + "quizzaciously xyzzy.", + &[ + 260, 1157, 1818, 362, 1814, 323, 260, 568, 327, 1818, 327, 263, + ], + ), + ("I'm listening.", &[268, 264, 283, 260, 604, 273, 263]), + ]; + for (text, expected) in cases { + let encoding = tokenizer.encode(*text, false).expect("tokenize"); + assert_eq!(encoding.get_ids(), *expected, "{text}"); + } + } + + #[test] + #[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"] + fn loader_splits_oversized_prompts_at_bundle_token_limit() { + let dir = std::env::var("BUZZ_POCKET_TEST_MODEL_DIR") + .expect("set BUZZ_POCKET_TEST_MODEL_DIR to the verified April bundle"); + let engine = AprilPocketTts::load(Path::new(&dir), 1).expect("load April bundle"); + let text = "This deliberately long sentence repeats ordinary English words so the exact SentencePiece token limit is exercised without relying on punctuation, and it keeps adding more material until the prompt must be divided into multiple independently safe generation chunks before the recurrent state cache can be exhausted."; + let prepared = prepare_april_prompt(text).expect("prepare prompt"); + let chunks = engine.split_prompt(&prepared).expect("split prompt"); + + assert!(chunks.len() > 1); + assert!(chunks.iter().all(|chunk| { + engine.token_count(chunk).expect("tokenize chunk") <= engine.bundle.max_token_per_chunk + })); + } + + #[test] + #[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"] + fn gary_provost_long_sentence_respects_bundle_token_limit() { + let dir = std::env::var("BUZZ_POCKET_TEST_MODEL_DIR") + .expect("set BUZZ_POCKET_TEST_MODEL_DIR to the verified April bundle"); + let engine = AprilPocketTts::load(Path::new(&dir), 1).expect("load April bundle"); + 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 prepared = prepare_april_prompt(text).expect("prepare prompt"); + let chunks = engine.split_prompt(&prepared).expect("split long sentence"); + let token_counts: Vec<_> = chunks + .iter() + .map(|chunk| engine.token_count(chunk).expect("count tokens")) + .collect(); + + assert_eq!( + chunks, + [ + "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.", + ] + ); + assert_eq!(token_counts, [48, 44]); + } +} diff --git a/desktop/src-tauri/src/huddle/pocket_models.rs b/desktop/src-tauri/src/huddle/pocket_models.rs new file mode 100644 index 0000000000..de34c77a70 --- /dev/null +++ b/desktop/src-tauri/src/huddle/pocket_models.rs @@ -0,0 +1,130 @@ +//! Immutable capabilities for Buzz Desktop's April Pocket TTS bundle. + +/// Pinned upstream export repository. +pub const APRIL_MODEL_ID: &str = "KevinAHM/pocket-tts-onnx"; + +/// Pinned revision containing the `english_2026-04` bundle. +pub const APRIL_MODEL_REVISION: &str = "58a6d00cf13d239b6748cb0769f35c580a8f606c"; + +/// Language bundle selected from the pinned export. +pub const APRIL_BUNDLE_ID: &str = "english_2026-04"; + +/// Maximum input size declared by the April bundle. +pub const APRIL_MAX_TOKEN_PER_CHUNK: usize = 50; + +/// One immutable artifact required by the April INT8 runtime. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PocketModelArtifact { + pub filename: &'static str, + pub sha256: &'static str, + pub size_bytes: u64, + pub quantized: bool, +} + +/// Capabilities of Buzz Desktop's sole Pocket model. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PocketModelInfo { + pub bundle_id: &'static str, + pub source_model_id: &'static str, + pub revision: &'static str, + pub sample_rate: u32, + pub max_token_per_chunk: usize, + pub artifacts: &'static [PocketModelArtifact], + pub quantized_components: &'static [&'static str], +} + +const INT8_ARTIFACTS: [PocketModelArtifact; 8] = [ + PocketModelArtifact { + filename: "bundle.json", + sha256: "bab643150f437f37df080a710520ff39ed9ebd9a339f8ebdc739f7eddfc28b3f", + size_bytes: 24_381, + quantized: false, + }, + PocketModelArtifact { + filename: "bos_before_voice.npy", + sha256: "f46edf4f7007b7ba4ea58831f49d003e59e167b4641c44bb3addfe9231a780b1", + size_bytes: 4_224, + quantized: false, + }, + PocketModelArtifact { + filename: "tokenizer.model", + sha256: "d461765ae179566678c93091c5fa6f2984c31bbe990bf1aa62d92c64d91bc3f6", + size_bytes: 59_339, + quantized: false, + }, + PocketModelArtifact { + filename: "flow_lm_main_int8.onnx", + sha256: "f9bd8106b79a0192c1c43399ab938fb24900a95c1c599870d75a884e99000116", + size_bytes: 76_341_079, + quantized: true, + }, + PocketModelArtifact { + filename: "flow_lm_flow_int8.onnx", + sha256: "3dd781ee5abee9e195320bf0106bebd6372a852b3b36352524ee78b40554635d", + size_bytes: 9_962_530, + quantized: true, + }, + PocketModelArtifact { + filename: "mimi_decoder_int8.onnx", + sha256: "3630450a3297a101792a6ac66619ebc70ab916b265e6220c2afaef8b1673f925", + size_bytes: 22_684_077, + quantized: true, + }, + PocketModelArtifact { + filename: "mimi_encoder.onnx", + sha256: "853e2ca623b8782d94c3745ec6133bfdff7ce33d9b11128bd29ea03f28d76e3d", + size_bytes: 39_768_446, + quantized: false, + }, + PocketModelArtifact { + filename: "text_conditioner.onnx", + sha256: "4ecee995fb69f85c7a7493d11f7b5ee15d9950facc7ab3f5c9c49ef1e03847bb", + size_bytes: 16_388_344, + quantized: false, + }, +]; + +const INT8_COMPONENTS: [&str; 3] = ["flow_lm_main", "flow_lm_flow", "mimi_decoder"]; + +/// Return immutable metadata for Buzz Desktop's April INT8 model. +pub const fn april_model_info() -> PocketModelInfo { + PocketModelInfo { + bundle_id: APRIL_BUNDLE_ID, + source_model_id: APRIL_MODEL_ID, + revision: APRIL_MODEL_REVISION, + sample_rate: 24_000, + max_token_per_chunk: APRIL_MAX_TOKEN_PER_CHUNK, + artifacts: &INT8_ARTIFACTS, + quantized_components: &INT8_COMPONENTS, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn metadata_matches_pinned_int8_layout() { + let info = april_model_info(); + assert_eq!(info.artifacts.len(), 8); + assert_eq!( + info.quantized_components, + ["flow_lm_main", "flow_lm_flow", "mimi_decoder"] + ); + assert_eq!( + info.artifacts + .iter() + .map(|artifact| artifact.size_bytes) + .sum::(), + 165_232_420 + ); + assert!(info + .artifacts + .iter() + .any(|artifact| { artifact.filename == "mimi_encoder.onnx" && !artifact.quantized })); + assert!(!info + .artifacts + .iter() + .any(|artifact| artifact.filename == "mimi_encoder_int8.onnx")); + } +} diff --git a/desktop/src-tauri/src/huddle/tts.rs b/desktop/src-tauri/src/huddle/tts.rs index 63a435cd8e..a8e6837cc8 100644 --- a/desktop/src-tauri/src/huddle/tts.rs +++ b/desktop/src-tauri/src/huddle/tts.rs @@ -101,13 +101,9 @@ const SENTENCE_LEAD_IN_SAMPLES: usize = (SAMPLE_RATE as f64 * 0.020) as usize; /// 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. +/// This character budget performs only coarse sentence packing. The April +/// engine applies its SentencePiece tokenizer afterward and refines every +/// result at the bundle's exact 50-token boundary. const MAX_CHUNK_CHARS: usize = 200; /// Silence inserted between sentences by the TTS pipeline (seconds). @@ -502,7 +498,22 @@ fn tts_worker( .into_iter() .filter(|s| !s.trim().is_empty()) .collect(); - let chunks = group_sentences_into_chunks(&sentences, MAX_CHUNK_CHARS); + 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; + } for chunk in &chunks { if handle_cancel_or_shutdown( @@ -734,9 +745,8 @@ fn build_sentence_append_buffer( /// 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. +/// than `max_chars` becomes its own chunk here, then the Pocket engine splits +/// it at the April bundle's exact token limit before synthesis. /// /// Sentences within a chunk are joined with a single space; sentence-ending /// punctuation is preserved by `split_sentences`, so the model sees natural diff --git a/desktop/src-tauri/src/huddle/tts_tests.rs b/desktop/src-tauri/src/huddle/tts_tests.rs index 7887f8bbdb..585850c69f 100644 --- a/desktop/src-tauri/src/huddle/tts_tests.rs +++ b/desktop/src-tauri/src/huddle/tts_tests.rs @@ -949,9 +949,8 @@ fn chunk_grouping_packs_up_to_budget_then_spills() { 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. +/// A single sentence longer than the coarse budget is passed through here; +/// the loaded April engine subsequently enforces its exact 50-token limit. #[test] fn chunk_grouping_oversized_sentence_passes_through() { let long = "word ".repeat(60).trim_end().to_string() + "."; From 21c01799deb6f38e3997f716f589ebadf02cc3ef Mon Sep 17 00:00:00 2001 From: John Tennant Date: Tue, 28 Jul 2026 16:20:00 -0400 Subject: [PATCH 2/6] fix(desktop): preserve Pocket token split continuity Signed-off-by: John Tennant --- desktop/src-tauri/examples/pocket_bench.rs | 116 -------------- .../src-tauri/examples/pocket_clip_probe.rs | 122 -------------- .../src-tauri/examples/pocket_onset_probe.rs | 149 ------------------ desktop/src-tauri/src/huddle/pocket.rs | 12 -- desktop/src-tauri/src/huddle/tts.rs | 53 ++----- 5 files changed, 14 insertions(+), 438 deletions(-) delete mode 100644 desktop/src-tauri/examples/pocket_bench.rs delete mode 100644 desktop/src-tauri/examples/pocket_clip_probe.rs delete mode 100644 desktop/src-tauri/examples/pocket_onset_probe.rs diff --git a/desktop/src-tauri/examples/pocket_bench.rs b/desktop/src-tauri/examples/pocket_bench.rs deleted file mode 100644 index b4f5635a95..0000000000 --- a/desktop/src-tauri/examples/pocket_bench.rs +++ /dev/null @@ -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:: 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:: 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" - } - ); -} diff --git a/desktop/src-tauri/examples/pocket_clip_probe.rs b/desktop/src-tauri/examples/pocket_clip_probe.rs deleted file mode 100644 index ad8657599f..0000000000 --- a/desktop/src-tauri/examples/pocket_clip_probe.rs +++ /dev/null @@ -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:: 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:: 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::() / 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" - } - ); -} diff --git a/desktop/src-tauri/examples/pocket_onset_probe.rs b/desktop/src-tauri/examples/pocket_onset_probe.rs deleted file mode 100644 index 05b4d0193c..0000000000 --- a/desktop/src-tauri/examples/pocket_onset_probe.rs +++ /dev/null @@ -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:: 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:: 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 - ); - } -} diff --git a/desktop/src-tauri/src/huddle/pocket.rs b/desktop/src-tauri/src/huddle/pocket.rs index 2154a25c22..02db11d42c 100644 --- a/desktop/src-tauri/src/huddle/pocket.rs +++ b/desktop/src-tauri/src/huddle/pocket.rs @@ -87,18 +87,6 @@ pub fn load_text_to_speech(model_dir: &str) -> Result { } 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, 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 diff --git a/desktop/src-tauri/src/huddle/tts.rs b/desktop/src-tauri/src/huddle/tts.rs index a8e6837cc8..2122723f7c 100644 --- a/desktop/src-tauri/src/huddle/tts.rs +++ b/desktop/src-tauri/src/huddle/tts.rs @@ -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 = 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) -> Vec { samples.into_iter().map(|s| s.clamp(-1.0, 1.0)).collect() } @@ -678,14 +657,10 @@ fn clamp_to_full_scale(samples: Vec) -> Vec { /// /// # 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 From 132dafda915ebe32b84347f672e8b9ff2d2fadfd Mon Sep 17 00:00:00 2001 From: John Tennant Date: Tue, 28 Jul 2026 16:27:42 -0400 Subject: [PATCH 3/6] fix(desktop): stream Pocket model units without seams Signed-off-by: John Tennant --- desktop/src-tauri/src/huddle/pocket.rs | 12 ++ desktop/src-tauri/src/huddle/tts.rs | 146 ++++++++++++++-------- desktop/src-tauri/src/huddle/tts_tests.rs | 33 ++++- 3 files changed, 132 insertions(+), 59 deletions(-) diff --git a/desktop/src-tauri/src/huddle/pocket.rs b/desktop/src-tauri/src/huddle/pocket.rs index 02db11d42c..2154a25c22 100644 --- a/desktop/src-tauri/src/huddle/pocket.rs +++ b/desktop/src-tauri/src/huddle/pocket.rs @@ -87,6 +87,18 @@ pub fn load_text_to_speech(model_dir: &str) -> Result { } 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, 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 diff --git a/desktop/src-tauri/src/huddle/tts.rs b/desktop/src-tauri/src/huddle/tts.rs index 2122723f7c..2a31b89893 100644 --- a/desktop/src-tauri/src/huddle/tts.rs +++ b/desktop/src-tauri/src/huddle/tts.rs @@ -488,18 +488,17 @@ fn tts_worker( // 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). 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. + // greedily up to MAX_CHUNK_CHARS. Playback of each model unit overlaps + // synthesis of the next one. The Pocket engine applies its exact + // 50-token split; keeping those units within one playback chunk avoids + // adding fades and pauses at token-only boundaries. 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 chunk in &chunks { + 'playback_chunks: for chunk in &chunks { if handle_cancel_or_shutdown( &cancel, &shutdown, @@ -516,51 +515,76 @@ fn tts_worker( continue; } - 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 - // onset (see `apply_fade_out` docstring + the - // 2026-05-18 "first little sound is missing" regression). - apply_fade_out(&mut audio); - - // Build one contiguous buffer per synthesized sentence: - // lead-in cushion + audio + trailing gap. Keeping this as - // a single rodio source preserves the original queue/drain - // semantics (one append per sentence) while still giving - // every chunk a quiet device warm-up window. - let buf = - build_sentence_append_buffer(&mut first_append, audio, silence_buf_len); - - // Check-and-append under `player_ops`, serialized with - // the monitor: a barge-in may have arrived during - // synthesis (the blocking window the monitor thread - // exists for). Don't append the now-stale sentence — the - // human interrupted; speaking it anyway would talk over - // them. Holding the lock for the check + append means the - // monitor can never clear between our check passing and - // the buffer landing. The flag is deliberately NOT - // consumed here: the loop-top handle_cancel_or_shutdown - // does the full consume (drain queue, reset lead-in) on - // the next iteration. - let _ops = lock_player_ops(&player_ops); - if cancel.load(Ordering::Acquire) { - // Nothing appended; the loop-top consume re-arms - // `first_append` (the flag is still set — the worker - // is its only consumer). - break; - } - player.append(SamplesBuffer::new(channels, rate, buf)); - // NOTE: tts_active is set AFTER player.append(), not - // before. Setting it before synthesis would cause STT to - // discard user speech during the synthesis window as - // "echo" even though no audio is actually playing yet. - // See crossfire review C3. - tts_active.store(true, Ordering::Release); + let model_chunks = match engine.split_text_into_chunks(text) { + Ok(model_chunks) => model_chunks, + Err(error) => { + eprintln!("buzz-desktop: TTS chunking failed: {error}"); + break; } - Ok(_) => {} - Err(e) => { - eprintln!("buzz-desktop: TTS synth failed: {e}"); + }; + let model_chunk_count = model_chunks.len(); + for (model_chunk_index, model_chunk) in model_chunks.iter().enumerate() { + if handle_cancel_or_shutdown( + &cancel, + &shutdown, + &tts_active, + &text_rx, + Some((&player, &player_ops)), + ) { + first_append = true; + break 'playback_chunks; + } + + let ends_playback_chunk = model_chunk_index + 1 == model_chunk_count; + match engine.synth_chunk(model_chunk, "en", &style, SYNTH_STEPS) { + Ok(samples) if !samples.is_empty() => { + let mut audio = clamp_to_full_scale(samples); + if ends_playback_chunk { + // Fade only at the playback-chunk boundary. Applying + // it at the model's internal token boundary would + // create an audible dip between contiguous units. + apply_fade_out(&mut audio); + } + + let buf = build_sentence_append_buffer( + &mut first_append, + audio, + silence_buf_len, + model_chunk_index == 0, + ends_playback_chunk, + ); + + // Check-and-append under `player_ops`, serialized with + // the monitor: a barge-in may have arrived during + // synthesis (the blocking window the monitor thread + // exists for). Don't append the now-stale sentence — the + // human interrupted; speaking it anyway would talk over + // them. Holding the lock for the check + append means the + // monitor can never clear between our check passing and + // the buffer landing. The flag is deliberately NOT + // consumed here: the loop-top handle_cancel_or_shutdown + // does the full consume (drain queue, reset lead-in) on + // the next iteration. + let _ops = lock_player_ops(&player_ops); + if cancel.load(Ordering::Acquire) { + // Nothing appended; the loop-top consume re-arms + // `first_append` (the flag is still set — the worker + // is its only consumer). + break; + } + player.append(SamplesBuffer::new(channels, rate, buf)); + // NOTE: tts_active is set AFTER player.append(), not + // before. Setting it before synthesis would cause STT to + // discard user speech during the synthesis window as + // "echo" even though no audio is actually playing yet. + // See crossfire review C3. + tts_active.store(true, Ordering::Release); + } + Ok(_) => {} + Err(e) => { + eprintln!("buzz-desktop: TTS synth failed: {e}"); + break 'playback_chunks; + } } } } @@ -696,18 +720,34 @@ fn apply_fade_out(samples: &mut [f32]) { /// The worker uses it in the idle branch of the main loop to distinguish /// "never queued anything since last drain" from "drained after speaking", /// which controls when `tts_active` is released and the lead-in re-armed. +/// Add silence only at the outer playback boundary. +/// +/// A playback chunk may contain several model-sized synthesis units. The first +/// unit receives the onset cushion, the last receives the remaining gap, and +/// intermediate units stay sample-contiguous. fn build_sentence_append_buffer( first_append: &mut bool, audio: Vec, silence_buf_len: usize, + starts_playback_chunk: bool, + ends_playback_chunk: bool, ) -> Vec { if *first_append { *first_append = false; } - let trailing_silence_len = silence_buf_len.saturating_sub(SENTENCE_LEAD_IN_SAMPLES); - let mut buf = Vec::with_capacity(SENTENCE_LEAD_IN_SAMPLES + audio.len() + trailing_silence_len); - buf.extend(std::iter::repeat_n(0.0_f32, SENTENCE_LEAD_IN_SAMPLES)); + 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 buf = Vec::with_capacity(lead_in_len + audio.len() + trailing_silence_len); + buf.extend(std::iter::repeat_n(0.0_f32, lead_in_len)); buf.extend(audio); buf.extend(std::iter::repeat_n(0.0_f32, trailing_silence_len)); buf diff --git a/desktop/src-tauri/src/huddle/tts_tests.rs b/desktop/src-tauri/src/huddle/tts_tests.rs index 585850c69f..79825e2260 100644 --- a/desktop/src-tauri/src/huddle/tts_tests.rs +++ b/desktop/src-tauri/src/huddle/tts_tests.rs @@ -812,6 +812,8 @@ fn lead_in_pad_is_present_for_every_sentence_chunk() { &mut first, vec![0.5_f32; SENTENCE_AUDIO_LEN], SILENCE_BUF_LEN, + true, + true, ); assert_eq!(buf.len(), SENTENCE_AUDIO_LEN + SILENCE_BUF_LEN); @@ -840,11 +842,11 @@ fn lead_in_pad_is_present_for_every_sentence_chunk() { #[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); + 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); + 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); } @@ -853,7 +855,7 @@ fn build_sentence_append_buffer_flips_first_append() { #[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); + 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); } @@ -863,8 +865,10 @@ fn first_sentence_leading_silence_is_exactly_lead_in() { 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); - let second_buf = build_sentence_append_buffer(&mut first, vec![0.5; 100], silence_buf_len); + 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]; @@ -877,7 +881,7 @@ fn sentence_gap_budget_is_preserved() { #[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); + 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)); @@ -888,6 +892,23 @@ fn sentence_append_buffer_is_one_contiguous_source() { ); } +/// 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); +} + // ── clamp_to_full_scale tests ───────────────────────────────────────────── /// In-range speech audio passes through bit-exact — no gain is applied. From b9054668c78c1611206f384c6aae642b9758a07f Mon Sep 17 00:00:00 2001 From: John Tennant Date: Tue, 28 Jul 2026 16:32:08 -0400 Subject: [PATCH 4/6] fix(desktop): continue TTS after unit errors Signed-off-by: John Tennant --- desktop/src-tauri/src/huddle/tts.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desktop/src-tauri/src/huddle/tts.rs b/desktop/src-tauri/src/huddle/tts.rs index 2a31b89893..7f7ae06a7c 100644 --- a/desktop/src-tauri/src/huddle/tts.rs +++ b/desktop/src-tauri/src/huddle/tts.rs @@ -583,7 +583,7 @@ fn tts_worker( Ok(_) => {} Err(e) => { eprintln!("buzz-desktop: TTS synth failed: {e}"); - break 'playback_chunks; + break; } } } From f8a795ac5421603b4105068494c8c843f2c30492 Mon Sep 17 00:00:00 2001 From: John Tennant Date: Tue, 28 Jul 2026 16:40:04 -0400 Subject: [PATCH 5/6] test(desktop): split Pocket token boundary tests Signed-off-by: John Tennant --- desktop/src-tauri/src/huddle/tts_tests.rs | 30 ++----------------- .../src/huddle/tts_tests/token_split.rs | 24 +++++++++++++++ 2 files changed, 27 insertions(+), 27 deletions(-) create mode 100644 desktop/src-tauri/src/huddle/tts_tests/token_split.rs diff --git a/desktop/src-tauri/src/huddle/tts_tests.rs b/desktop/src-tauri/src/huddle/tts_tests.rs index 79825e2260..1908b096b1 100644 --- a/desktop/src-tauri/src/huddle/tts_tests.rs +++ b/desktop/src-tauri/src/huddle/tts_tests.rs @@ -9,6 +9,9 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc; use std::sync::{Arc, Mutex}; +#[path = "tts_tests/token_split.rs"] +mod token_split; + // ── Remote interrupt tracker ────────────────────────────────────────────── // // Models the per-peer frame counting logic in the recv task of @@ -785,16 +788,6 @@ fn apply_fade_out_single_sample() { assert_eq!(samples[0], 1.0); } -/// Sanity-check the per-sentence cushion length: 20 ms at 24 kHz must -/// land at exactly 480 samples. This is a const computation, so the -/// real value of this test is documenting *why* 20 ms was chosen — it -/// covers a typical CoreAudio buffer turnover (256–1024 samples) -/// without being audible as user-facing latency. -#[test] -fn sentence_lead_in_is_sane() { - assert_eq!(SENTENCE_LEAD_IN_SAMPLES, 480, "20 ms × 24 kHz"); -} - // ── build_sentence_append_buffer tests ─────────────────────────────────── /// REGRESSION: every chunk needs an onset cushion; synthesized chunks @@ -892,23 +885,6 @@ fn sentence_append_buffer_is_one_contiguous_source() { ); } -/// 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); -} - // ── clamp_to_full_scale tests ───────────────────────────────────────────── /// In-range speech audio passes through bit-exact — no gain is applied. diff --git a/desktop/src-tauri/src/huddle/tts_tests/token_split.rs b/desktop/src-tauri/src/huddle/tts_tests/token_split.rs new file mode 100644 index 0000000000..b9249c9afc --- /dev/null +++ b/desktop/src-tauri/src/huddle/tts_tests/token_split.rs @@ -0,0 +1,24 @@ +use super::*; + +/// The onset cushion covers 20 ms at the production sample rate. +#[test] +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); +} From 0e7abb406fc49fafca36e063c8137b744b3aa897 Mon Sep 17 00:00:00 2001 From: John Tennant Date: Tue, 28 Jul 2026 17:29:27 -0400 Subject: [PATCH 6/6] fix(desktop): cushion Pocket playback underruns Signed-off-by: John Tennant --- desktop/src-tauri/src/huddle/tts.rs | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/desktop/src-tauri/src/huddle/tts.rs b/desktop/src-tauri/src/huddle/tts.rs index 7f7ae06a7c..f084cca6d5 100644 --- a/desktop/src-tauri/src/huddle/tts.rs +++ b/desktop/src-tauri/src/huddle/tts.rs @@ -550,7 +550,7 @@ fn tts_worker( &mut first_append, audio, silence_buf_len, - model_chunk_index == 0, + model_chunk_index == 0 || player.empty(), ends_playback_chunk, ); @@ -699,13 +699,12 @@ fn apply_fade_out(samples: &mut [f32]) { } } -/// Build the single buffer appended to the rodio `Player` for one synthesised -/// sentence. +/// Build one buffer appended to the rodio `Player` for a synthesis unit. /// -/// Every sentence chunk gets a short lead-in pad immediately before its audio. -/// This matters for chunks that start with soft first phonemes (`I'm`, `I've`): -/// the synthesized buffer can begin with speech within the first millisecond, -/// so the playback layer must provide the device/mixer cushion. +/// Every playback boundary gets a short lead-in pad immediately before its +/// audio. This matters for chunks that start with soft first phonemes (`I'm`, +/// `I've`): the synthesized buffer can begin with speech within the first +/// millisecond, so the playback layer must provide the device/mixer cushion. /// To keep the audible gap unchanged, the trailing silence after this chunk is /// shortened by the same amount (`silence_buf_len - SENTENCE_LEAD_IN_SAMPLES`): /// sentence N contributes 80 ms of post-speech silence and sentence N+1 @@ -716,15 +715,15 @@ fn apply_fade_out(samples: &mut [f32]) { /// tracked source per synthesized sentence, avoiding source-boundary/drain /// regressions from enqueueing the lead-in, audio, and tail as separate sounds. /// +/// A playback chunk may contain several model-sized synthesis units. Only the +/// first unit receives the onset cushion and only the last receives the +/// remaining gap. If playback underruns while the next unit is synthesized, +/// that unit becomes a new playback boundary and receives a fresh cushion. +/// /// `first_append` is flipped on the first call after the player goes idle. /// The worker uses it in the idle branch of the main loop to distinguish /// "never queued anything since last drain" from "drained after speaking", /// which controls when `tts_active` is released and the lead-in re-armed. -/// Add silence only at the outer playback boundary. -/// -/// A playback chunk may contain several model-sized synthesis units. The first -/// unit receives the onset cushion, the last receives the remaining gap, and -/// intermediate units stay sample-contiguous. fn build_sentence_append_buffer( first_append: &mut bool, audio: Vec,