mirror of
https://github.com/block/buzz.git
synced 2026-11-01 15:01:12 +00:00
Upgrade Desktop Pocket TTS to April INT8
Signed-off-by: John Tennant <jtennant@block.xyz>
This commit is contained in:
John Tennant
472
desktop/src-tauri/Cargo.lock
generated
472
desktop/src-tauri/Cargo.lock
generated
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 -- \
|
||||
//! <int8-model-dir> <fp32-model-dir> <output-dir> [--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<KeyItem>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct KeyItem {
|
||||
id: String,
|
||||
kind: String,
|
||||
text: String,
|
||||
clips: Vec<KeyClip>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct KeyClip {
|
||||
file: String,
|
||||
precision: Precision,
|
||||
chunking: Chunking,
|
||||
cold_start: bool,
|
||||
idle_minutes: Option<u64>,
|
||||
synthesis_ms: u128,
|
||||
audio_seconds: f32,
|
||||
}
|
||||
|
||||
struct Voice {
|
||||
samples: Vec<f32>,
|
||||
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::<u64>()
|
||||
.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<String> = 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<String> = 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<String>, label: &str) -> Result<PathBuf, String> {
|
||||
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<Engine, String> {
|
||||
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<Vec<f32>, 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::<fn(&[f32], f32) -> bool>)
|
||||
.ok_or_else(|| format!("synthesis failed for {chunk:?}"))?;
|
||||
let mut samples: Vec<f32> = 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<f32> {
|
||||
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<f32>, 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<u8>)> = (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())
|
||||
}
|
||||
@@ -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<u64>,
|
||||
status: Arc<Mutex<ModelStatus>>,
|
||||
just_ready: Arc<AtomicBool>, // 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<u64>) -> 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<PathBuf> {
|
||||
@@ -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<F, Fut>(
|
||||
&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<u64> {
|
||||
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<Self> {
|
||||
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}"))?;
|
||||
let actual = sha256_file(&dest).await?;
|
||||
if actual != expected {
|
||||
if bytes != artifact.size_bytes {
|
||||
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} size check failed: expected {} bytes, got {bytes}",
|
||||
artifact.size_bytes
|
||||
));
|
||||
}
|
||||
let actual = sha256_file(&dest).await?;
|
||||
if actual != artifact.sha256 {
|
||||
let _ = tokio::fs::remove_dir_all(&temp_dir).await;
|
||||
return Err(format!(
|
||||
"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;
|
||||
|
||||
148
desktop/src-tauri/src/huddle/models_tests.rs
Normal file
148
desktop/src-tauri/src/huddle/models_tests.rs
Normal file
@@ -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::<Vec<_>>();
|
||||
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());
|
||||
}
|
||||
@@ -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: <https://huggingface.co/kyutai/pocket-tts>.
|
||||
//! Licensed CC-BY-4.0.
|
||||
//! - **Mimi neural codec**: Kyutai, bundled in the same release. CC-BY-4.0.
|
||||
//! - **ONNX export**: KevinAHM —
|
||||
//! <https://huggingface.co/KevinAHM/pocket-tts-onnx>. CC-BY-4.0.
|
||||
//! - **sherpa-onnx repackage**: csukuangfj / k2-fsa —
|
||||
//! <https://huggingface.co/csukuangfj2/sherpa-onnx-pocket-tts-2026-01-26>.
|
||||
//! 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
|
||||
//! (<https://kyutai.org/tts>), which maps to `vctk/p333_023_enhanced.wav`
|
||||
//! in <https://huggingface.co/kyutai/tts-voices>. CC-BY-4.0, base recording
|
||||
//! from the VCTK corpus, enhanced by ai-coustics.
|
||||
//! - 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.
|
||||
//!
|
||||
//! 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.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<Engine, String>`
|
||||
//! - `load_voice_style(path)` → `Result<VoiceStyle, String>`
|
||||
//! - `Engine::synth_chunk(&self, text, lang, &VoiceStyle, steps)`
|
||||
//! → `Result<Vec<f32>, 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::{GenerationConfig, OfflineTts, OfflineTtsConfig, Wave};
|
||||
use sherpa_onnx::Wave;
|
||||
|
||||
// ── Engine-module contract: public consts ─────────────────────────────────────
|
||||
#[path = "pocket_april.rs"]
|
||||
mod pocket_april;
|
||||
#[path = "pocket_models.rs"]
|
||||
mod pocket_models;
|
||||
|
||||
/// 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.
|
||||
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.
|
||||
pub const SAMPLE_RATE: u32 = 24_000;
|
||||
|
||||
/// Name (without extension) of the bundled reference voice. The model directory
|
||||
/// is expected to contain `<DEFAULT_VOICE>.<VOICE_FILE_EXT>` 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 ────────────────────────────────────────────────────────────────────
|
||||
const TTS_NUM_THREADS: usize = 1;
|
||||
|
||||
/// 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 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<f32>,
|
||||
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<VoiceStyle, String> {
|
||||
let path_str = path
|
||||
.to_str()
|
||||
@@ -195,199 +63,46 @@ pub fn load_voice_style(path: &Path) -> Result<VoiceStyle, String> {
|
||||
})
|
||||
}
|
||||
|
||||
// ── 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<AprilPocketTts>,
|
||||
}
|
||||
|
||||
/// 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<PocketTts, String> {
|
||||
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()));
|
||||
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()
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
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<i32>,
|
||||
}
|
||||
|
||||
/// 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<PreparedPrompt> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
// 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<HashMap<String, serde_json::Value>> {
|
||||
prepared.max_frames.map(|mf| {
|
||||
let mut h: HashMap<String, serde_json::Value> = 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<Vec<String>, String> {
|
||||
let Some(prepared) = prepare_april_prompt(text) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
self.inner
|
||||
.lock()
|
||||
.map_err(|_| "Pocket TTS engine lock poisoned".to_string())?
|
||||
.split_prompt(&prepared)
|
||||
}
|
||||
|
||||
/// Synthesize text with the supplied reference voice.
|
||||
///
|
||||
/// `_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<Vec<f32>, 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::<fn(...) -> 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::<fn(&[f32], f32) -> 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)
|
||||
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_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"
|
||||
);
|
||||
}
|
||||
}
|
||||
#[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");
|
||||
|
||||
#[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('.'));
|
||||
}
|
||||
|
||||
#[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");
|
||||
}
|
||||
assert!(!samples.is_empty());
|
||||
assert!(samples.iter().all(|sample| sample.is_finite()));
|
||||
assert!(samples.iter().any(|sample| sample.abs() > 1.0e-6));
|
||||
}
|
||||
}
|
||||
|
||||
940
desktop/src-tauri/src/huddle/pocket_april.rs
Normal file
940
desktop/src-tauri/src/huddle/pocket_april.rs
Normal file
@@ -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<usize>,
|
||||
max_token_per_chunk: usize,
|
||||
tokenizer_file: String,
|
||||
bos_before_voice_file: String,
|
||||
flow_lm_state_manifest: Vec<StateSpec>,
|
||||
mimi_state_manifest: Vec<StateSpec>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
struct StateSpec {
|
||||
input_name: String,
|
||||
output_name: String,
|
||||
dtype: StateDtype,
|
||||
shape: Vec<i64>,
|
||||
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<f32>,
|
||||
}
|
||||
|
||||
pub(crate) struct AprilPocketTts {
|
||||
bundle: Bundle,
|
||||
tokenizer: Tokenizer,
|
||||
bos_embedding: Vec<f32>,
|
||||
mimi_encoder: Session,
|
||||
text_conditioner: Session,
|
||||
flow_main: Session,
|
||||
flow: Session,
|
||||
mimi_decoder: Session,
|
||||
cached_voice: Option<CachedVoice>,
|
||||
}
|
||||
|
||||
#[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<AprilPreparedPrompt> {
|
||||
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<Self, String> {
|
||||
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<Vec<String>, 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<Vec<f32>, 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::<Vec<_>>();
|
||||
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<usize, String> {
|
||||
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<usize, String> {
|
||||
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<Vec<f32>, 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::<f32>()
|
||||
.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<Vec<StateValue>, String> {
|
||||
let frames = embeddings.len() / self.bundle.conditioning_dim;
|
||||
let sequence = Tensor::<f32>::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<i64>) -> Result<Vec<f32>, 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::<f32>()
|
||||
.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::<f32>::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<Vec<f32>, 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::<f32>::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::<f32>()
|
||||
.map_err(ort_error("extract Flow LM conditioning"))?
|
||||
.1
|
||||
.to_vec();
|
||||
let eos_logit = outputs[1]
|
||||
.try_extract_tensor::<f32>()
|
||||
.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::<f32>()
|
||||
.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<Vec<f32>, 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::<f32>()
|
||||
.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<Session, String> {
|
||||
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<Tokenizer, String> {
|
||||
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<Vec<StateValue>, 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::<f32>::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::<i64>::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::<bool>::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<usize, String> {
|
||||
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<f32> {
|
||||
let mut out = Vec::with_capacity(len);
|
||||
while out.len() < len {
|
||||
let u1 = rng.random::<f32>().max(f32::MIN_POSITIVE);
|
||||
let u2 = rng.random::<f32>();
|
||||
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<Vec<f32>, 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': '<f4'") || header.contains("\"descr\": \"<f4\""))
|
||||
|| header.contains("fortran_order': True")
|
||||
|| header.contains("fortran_order\": true")
|
||||
{
|
||||
return Err(format!(
|
||||
"{} must be a little-endian, C-order float32 NumPy array",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
let data = &bytes[data_start..];
|
||||
if data.len() % 4 != 0 {
|
||||
return Err(format!("misaligned float32 data in {}", path.display()));
|
||||
}
|
||||
Ok(data
|
||||
.chunks_exact(4)
|
||||
.map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn ort_error(context: &'static str) -> 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]);
|
||||
}
|
||||
}
|
||||
130
desktop/src-tauri/src/huddle/pocket_models.rs
Normal file
130
desktop/src-tauri/src/huddle/pocket_models.rs
Normal file
@@ -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::<u64>(),
|
||||
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"));
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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() + ".";
|
||||
|
||||
Reference in New Issue
Block a user