feat: import Chinese-localized Buzz source snapshot
Docker image / Build (linux/amd64) (push) Has been cancelled
Docker image / Build (linux/arm64) (push) Has been cancelled
Docker image / Merge release multi-arch manifest (push) Has been cancelled
Docker image / Merge debug multi-arch manifest (push) Has been cancelled
Docker image / Build public push gateway (linux/amd64) (push) Has been cancelled
Docker image / Build public push gateway (linux/arm64) (push) Has been cancelled
Docker image / Publish public push gateway image (push) Has been cancelled
Sprig image / Build (linux/amd64) (push) Has been cancelled
Sprig image / Build (linux/arm64) (push) Has been cancelled
Sprig image / Merge multi-arch manifest (push) Has been cancelled
Harbor Buzz Orchestra / Python tests and lint (push) Has been cancelled
CI / Detect Changed Paths (push) Has been cancelled
CI / Rust Lint (push) Has been cancelled
CI / Unit Tests (push) Has been cancelled
CI / Desktop Core (push) Has been cancelled
CI / Desktop Smoke E2E (1) (push) Has been cancelled
CI / Desktop Smoke E2E (2) (push) Has been cancelled
CI / Desktop Smoke E2E (3) (push) Has been cancelled
CI / Desktop Smoke E2E (4) (push) Has been cancelled
CI / Desktop (push) Has been cancelled
CI / Desktop E2E Relay (push) Has been cancelled
CI / Desktop E2E Integration (1/2) (push) Has been cancelled
CI / Desktop E2E Integration (2/2) (push) Has been cancelled
CI / Desktop E2E Integration (push) Has been cancelled
CI / Backend Integration (relay e2e) (push) Has been cancelled
CI / Relay E2E (push) Has been cancelled
CI / Web (push) Has been cancelled
CI / Mobile (push) Has been cancelled
CI / Security (push) Has been cancelled
CI / Dead Token Reference Guard (push) Has been cancelled
CI / Server Cross-Compile (aarch64-unknown-linux-musl) (push) Has been cancelled
CI / Server Cross-Compile (x86_64-unknown-linux-musl) (push) Has been cancelled
CI / Windows Rust (x86_64-pc-windows-msvc) (push) Has been cancelled
CI / Desktop Build (macOS) (push) Has been cancelled
helm chart / lint + unittest + render matrix (push) Has been cancelled
helm chart / install on kind (gated) (push) Has been cancelled
helm chart / publish chart to GHCR (push) Has been cancelled
Mesh Lifecycle / Relay-Driven Mesh Lifecycle Smoke (push) Has been cancelled
Sprig / Build (aarch64-unknown-linux-musl) (push) Has been cancelled
Sprig / Build (x86_64-unknown-linux-musl) (push) Has been cancelled
Sprig / Publish rolling release (push) Has been cancelled
Sprig / Publish tagged release (push) Has been cancelled
Docker image / Build (linux/amd64) (push) Has been cancelled
Docker image / Build (linux/arm64) (push) Has been cancelled
Docker image / Merge release multi-arch manifest (push) Has been cancelled
Docker image / Merge debug multi-arch manifest (push) Has been cancelled
Docker image / Build public push gateway (linux/amd64) (push) Has been cancelled
Docker image / Build public push gateway (linux/arm64) (push) Has been cancelled
Docker image / Publish public push gateway image (push) Has been cancelled
Sprig image / Build (linux/amd64) (push) Has been cancelled
Sprig image / Build (linux/arm64) (push) Has been cancelled
Sprig image / Merge multi-arch manifest (push) Has been cancelled
Harbor Buzz Orchestra / Python tests and lint (push) Has been cancelled
CI / Detect Changed Paths (push) Has been cancelled
CI / Rust Lint (push) Has been cancelled
CI / Unit Tests (push) Has been cancelled
CI / Desktop Core (push) Has been cancelled
CI / Desktop Smoke E2E (1) (push) Has been cancelled
CI / Desktop Smoke E2E (2) (push) Has been cancelled
CI / Desktop Smoke E2E (3) (push) Has been cancelled
CI / Desktop Smoke E2E (4) (push) Has been cancelled
CI / Desktop (push) Has been cancelled
CI / Desktop E2E Relay (push) Has been cancelled
CI / Desktop E2E Integration (1/2) (push) Has been cancelled
CI / Desktop E2E Integration (2/2) (push) Has been cancelled
CI / Desktop E2E Integration (push) Has been cancelled
CI / Backend Integration (relay e2e) (push) Has been cancelled
CI / Relay E2E (push) Has been cancelled
CI / Web (push) Has been cancelled
CI / Mobile (push) Has been cancelled
CI / Security (push) Has been cancelled
CI / Dead Token Reference Guard (push) Has been cancelled
CI / Server Cross-Compile (aarch64-unknown-linux-musl) (push) Has been cancelled
CI / Server Cross-Compile (x86_64-unknown-linux-musl) (push) Has been cancelled
CI / Windows Rust (x86_64-pc-windows-msvc) (push) Has been cancelled
CI / Desktop Build (macOS) (push) Has been cancelled
helm chart / lint + unittest + render matrix (push) Has been cancelled
helm chart / install on kind (gated) (push) Has been cancelled
helm chart / publish chart to GHCR (push) Has been cancelled
Mesh Lifecycle / Relay-Driven Mesh Lifecycle Smoke (push) Has been cancelled
Sprig / Build (aarch64-unknown-linux-musl) (push) Has been cancelled
Sprig / Build (x86_64-unknown-linux-musl) (push) Has been cancelled
Sprig / Publish rolling release (push) Has been cancelled
Sprig / Publish tagged release (push) Has been cancelled
Signed-off-by: cls_宁波本机 <908705107@qq.com>
This commit is contained in:
@@ -0,0 +1,25 @@
|
||||
[package]
|
||||
name = "buzz-voice"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
rust-version.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Reusable local voice primitives for Buzz"
|
||||
|
||||
[dependencies]
|
||||
atomic-write-file = "0.3"
|
||||
hex = { workspace = true }
|
||||
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"
|
||||
sentencepiece-model = "0.1"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
sha2 = { workspace = true }
|
||||
sherpa-onnx = "1.12"
|
||||
symphonia = { version = "0.5", default-features = false, features = ["aac", "aiff", "alac", "flac", "isomp4", "mp3", "ogg", "pcm", "vorbis", "wav"] }
|
||||
tokenizers = { version = "0.22", default-features = false, features = ["fancy-regex"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
@@ -0,0 +1,730 @@
|
||||
//! Device-local Pocket reference voice validation, canonicalization, and storage.
|
||||
|
||||
use std::{
|
||||
fs,
|
||||
io::Write,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
||||
use atomic_write_file::AtomicWriteFile;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use symphonia::core::{
|
||||
audio::SampleBuffer, codecs::DecoderOptions, errors::Error as SymphoniaError,
|
||||
formats::FormatOptions, io::MediaSourceStream, meta::MetadataOptions, probe::Hint,
|
||||
};
|
||||
|
||||
const MAX_SOURCE_BYTES: u64 = 25 * 1024 * 1024;
|
||||
const MIN_SAMPLE_RATE: u32 = 8_000;
|
||||
const MAX_SAMPLE_RATE: u32 = 96_000;
|
||||
const MIN_DURATION_SECONDS: f64 = 2.0;
|
||||
const MAX_DURATION_SECONDS: f64 = 30.0;
|
||||
pub const CANONICAL_SAMPLE_RATE: u32 = 32_000;
|
||||
const REGISTRY_VERSION: u32 = 1;
|
||||
const REGISTRY_FILE: &str = "registry.json";
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ImportedVoice {
|
||||
pub key: String,
|
||||
pub display_name: String,
|
||||
pub content_hash: String,
|
||||
pub file_name: String,
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ImportedVoiceRegistry {
|
||||
version: u32,
|
||||
voices: Vec<ImportedVoice>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PocketVoiceLibrary {
|
||||
root: PathBuf,
|
||||
}
|
||||
|
||||
impl PocketVoiceLibrary {
|
||||
pub fn new(root: impl Into<PathBuf>) -> Self {
|
||||
Self { root: root.into() }
|
||||
}
|
||||
|
||||
pub fn root(&self) -> &Path {
|
||||
&self.root
|
||||
}
|
||||
|
||||
fn registry_path(&self) -> PathBuf {
|
||||
self.root.join(REGISTRY_FILE)
|
||||
}
|
||||
|
||||
pub fn load(&self) -> Result<Vec<ImportedVoice>, String> {
|
||||
let path = self.registry_path();
|
||||
if !path.exists() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let bytes =
|
||||
fs::read(&path).map_err(|error| format!("could not read imported voices: {error}"))?;
|
||||
let registry: ImportedVoiceRegistry = serde_json::from_slice(&bytes)
|
||||
.map_err(|error| format!("imported voice registry is invalid: {error}"))?;
|
||||
if registry.version > REGISTRY_VERSION {
|
||||
return Err(format!(
|
||||
"imported voice registry version {} is newer than this Buzz build supports",
|
||||
registry.version
|
||||
));
|
||||
}
|
||||
Ok(registry
|
||||
.voices
|
||||
.into_iter()
|
||||
.filter(valid_identity)
|
||||
.filter(|voice| self.resolve_file(voice).is_ok())
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn save(&self, voices: &[ImportedVoice]) -> Result<(), String> {
|
||||
ensure_storage_dir(&self.root)?;
|
||||
let payload = serde_json::to_vec_pretty(&ImportedVoiceRegistry {
|
||||
version: REGISTRY_VERSION,
|
||||
voices: voices.to_vec(),
|
||||
})
|
||||
.map_err(|error| format!("could not encode imported voice registry: {error}"))?;
|
||||
atomic_write_restricted(&self.registry_path(), &payload)
|
||||
.map_err(|error| format!("could not save imported voice registry: {error}"))
|
||||
}
|
||||
|
||||
pub fn resolve_file(&self, voice: &ImportedVoice) -> Result<PathBuf, String> {
|
||||
if !valid_identity(voice) {
|
||||
return Err("Imported voice registry contains an invalid file identity".to_string());
|
||||
}
|
||||
let path = self.root.join(&voice.file_name);
|
||||
if !is_regular_file_without_symlink(&path) {
|
||||
return Err(format!("Imported voice {} is missing", voice.display_name));
|
||||
}
|
||||
let bytes =
|
||||
fs::read(&path).map_err(|error| format!("could not verify imported voice: {error}"))?;
|
||||
if hex::encode(Sha256::digest(bytes)) != voice.content_hash {
|
||||
return Err(format!(
|
||||
"Imported voice {} does not match its content identity",
|
||||
voice.display_name
|
||||
));
|
||||
}
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
pub fn find(&self, key: &str) -> Result<Option<ImportedVoice>, String> {
|
||||
Ok(self.load()?.into_iter().find(|voice| voice.key == key))
|
||||
}
|
||||
|
||||
pub fn import_path(&self, source: &Path) -> Result<ImportedVoice, String> {
|
||||
let metadata = fs::metadata(source)
|
||||
.map_err(|error| format!("could not inspect selected audio: {error}"))?;
|
||||
if metadata.len() > MAX_SOURCE_BYTES {
|
||||
return Err("Voice audio must be 25 MB or smaller".to_string());
|
||||
}
|
||||
let extension = source
|
||||
.extension()
|
||||
.and_then(|extension| extension.to_str())
|
||||
.map(str::to_ascii_lowercase)
|
||||
.ok_or_else(|| "Voice audio must have a supported file extension".to_string())?;
|
||||
let samples = if extension == "wav" {
|
||||
let source_bytes = fs::read(source)
|
||||
.map_err(|error| format!("could not read selected audio: {error}"))?;
|
||||
decode_wav(&source_bytes)?
|
||||
} else {
|
||||
decode_media(source, &extension)?
|
||||
};
|
||||
let canonical_samples = resample_linear(&samples.samples, samples.sample_rate);
|
||||
let canonical = encode_pcm16_wav(&canonical_samples, CANONICAL_SAMPLE_RATE);
|
||||
let hash = hex::encode(Sha256::digest(&canonical));
|
||||
let key = format!("pocket:imported:{hash}");
|
||||
let file_name = format!("{hash}.wav");
|
||||
let display_name = source
|
||||
.file_stem()
|
||||
.and_then(|name| name.to_str())
|
||||
.map(str::trim)
|
||||
.filter(|name| !name.is_empty())
|
||||
.unwrap_or("Imported voice")
|
||||
.chars()
|
||||
.take(80)
|
||||
.collect::<String>();
|
||||
|
||||
ensure_storage_dir(&self.root)?;
|
||||
let file_path = self.root.join(&file_name);
|
||||
let file_created = !file_path.exists();
|
||||
if file_created {
|
||||
atomic_write_restricted(&file_path, &canonical)
|
||||
.map_err(|error| format!("could not save imported voice audio: {error}"))?;
|
||||
} else {
|
||||
if !is_regular_file_without_symlink(&file_path) {
|
||||
return Err("Imported voice storage contains an unsafe file entry".to_string());
|
||||
}
|
||||
let existing = fs::read(&file_path)
|
||||
.map_err(|error| format!("could not verify imported voice audio: {error}"))?;
|
||||
if hex::encode(Sha256::digest(&existing)) != hash {
|
||||
return Err("Imported voice storage contains mismatched audio data".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let mut imported = ImportedVoice {
|
||||
key,
|
||||
display_name,
|
||||
content_hash: hash,
|
||||
file_name,
|
||||
};
|
||||
let mut voices = self.load()?;
|
||||
if let Some(existing) = voices
|
||||
.iter()
|
||||
.find(|voice| voice.content_hash == imported.content_hash)
|
||||
{
|
||||
imported = existing.clone();
|
||||
} else {
|
||||
voices.push(imported.clone());
|
||||
}
|
||||
if let Err(error) = self.save(&voices) {
|
||||
if file_created {
|
||||
let _ = fs::remove_file(&file_path);
|
||||
}
|
||||
return Err(error);
|
||||
}
|
||||
Ok(imported)
|
||||
}
|
||||
|
||||
pub fn delete(&self, key: &str) -> Result<(), String> {
|
||||
let mut voices = self.load()?;
|
||||
let index = voices
|
||||
.iter()
|
||||
.position(|voice| voice.key == key)
|
||||
.ok_or_else(|| format!("Unknown imported voice: {key}"))?;
|
||||
let previous_voices = voices.clone();
|
||||
let removed = voices.remove(index);
|
||||
self.save(&voices)?;
|
||||
let path = self.root.join(removed.file_name);
|
||||
match fs::remove_file(path) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
|
||||
Err(error) => {
|
||||
self.save(&previous_voices).map_err(|rollback_error| {
|
||||
format!(
|
||||
"Imported voice audio could not be deleted ({error}), and its registry \
|
||||
entry could not be restored ({rollback_error})"
|
||||
)
|
||||
})?;
|
||||
Err(format!(
|
||||
"Imported voice audio could not be deleted: {error}"
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct PcmStats {
|
||||
pub sample_count: usize,
|
||||
pub sample_rate: u32,
|
||||
pub duration_seconds: f64,
|
||||
pub peak: f32,
|
||||
pub rms: f32,
|
||||
pub non_silent_samples: usize,
|
||||
}
|
||||
|
||||
impl PcmStats {
|
||||
pub fn analyze(samples: &[f32], sample_rate: u32) -> Self {
|
||||
let peak = samples
|
||||
.iter()
|
||||
.filter(|sample| sample.is_finite())
|
||||
.fold(0.0_f32, |peak, sample| peak.max(sample.abs()));
|
||||
let square_sum = samples
|
||||
.iter()
|
||||
.filter(|sample| sample.is_finite())
|
||||
.map(|sample| sample * sample)
|
||||
.sum::<f32>();
|
||||
let rms = if samples.is_empty() {
|
||||
0.0
|
||||
} else {
|
||||
(square_sum / samples.len() as f32).sqrt()
|
||||
};
|
||||
Self {
|
||||
sample_count: samples.len(),
|
||||
sample_rate,
|
||||
duration_seconds: if sample_rate == 0 {
|
||||
0.0
|
||||
} else {
|
||||
samples.len() as f64 / f64::from(sample_rate)
|
||||
},
|
||||
peak,
|
||||
rms,
|
||||
non_silent_samples: samples
|
||||
.iter()
|
||||
.filter(|sample| sample.is_finite() && sample.abs() >= 0.001)
|
||||
.count(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_non_silent(self) -> bool {
|
||||
self.peak >= 0.001 && self.rms >= 0.0001 && self.non_silent_samples > 0
|
||||
}
|
||||
}
|
||||
|
||||
pub fn write_pcm16_wav(path: &Path, samples: &[f32], sample_rate: u32) -> Result<(), String> {
|
||||
let bytes = encode_pcm16_wav(samples, sample_rate);
|
||||
fs::write(path, bytes).map_err(|error| format!("could not write PCM evidence: {error}"))
|
||||
}
|
||||
|
||||
fn ensure_storage_dir(path: &Path) -> Result<(), String> {
|
||||
fs::create_dir_all(path)
|
||||
.map_err(|error| format!("could not create local voice storage: {error}"))?;
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
fs::set_permissions(path, fs::Permissions::from_mode(0o700))
|
||||
.map_err(|error| format!("could not restrict local voice storage: {error}"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn atomic_write_restricted(path: &Path, payload: &[u8]) -> Result<(), String> {
|
||||
let resolved = fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf());
|
||||
let mut file = AtomicWriteFile::open(&resolved)
|
||||
.map_err(|error| format!("open {} for atomic write: {error}", resolved.display()))?;
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
file.set_permissions(fs::Permissions::from_mode(0o600))
|
||||
.map_err(|error| format!("set {} permissions: {error}", resolved.display()))?;
|
||||
}
|
||||
file.write_all(payload)
|
||||
.map_err(|error| format!("write {}: {error}", resolved.display()))?;
|
||||
file.commit()
|
||||
.map_err(|error| format!("commit {}: {error}", resolved.display()))
|
||||
}
|
||||
|
||||
fn valid_hash(hash: &str) -> bool {
|
||||
hash.len() == 64 && hash.bytes().all(|byte| byte.is_ascii_hexdigit())
|
||||
}
|
||||
|
||||
fn valid_identity(voice: &ImportedVoice) -> bool {
|
||||
valid_hash(&voice.content_hash)
|
||||
&& voice.key == format!("pocket:imported:{}", voice.content_hash)
|
||||
&& voice.file_name == format!("{}.wav", voice.content_hash)
|
||||
}
|
||||
|
||||
fn is_regular_file_without_symlink(path: &Path) -> bool {
|
||||
fs::symlink_metadata(path)
|
||||
.is_ok_and(|metadata| metadata.file_type().is_file() && !metadata.file_type().is_symlink())
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct DecodedAudio {
|
||||
sample_rate: u32,
|
||||
samples: Vec<f32>,
|
||||
}
|
||||
|
||||
fn decode_wav(bytes: &[u8]) -> Result<DecodedAudio, String> {
|
||||
if bytes.len() < 12 || &bytes[..4] != b"RIFF" || &bytes[8..12] != b"WAVE" {
|
||||
return Err("Selected file is not a valid RIFF/WAVE file".to_string());
|
||||
}
|
||||
let mut offset = 12usize;
|
||||
let mut format = None;
|
||||
let mut data = None;
|
||||
while offset.checked_add(8).is_some_and(|end| end <= bytes.len()) {
|
||||
let id = &bytes[offset..offset + 4];
|
||||
let size =
|
||||
u32::from_le_bytes(bytes[offset + 4..offset + 8].try_into().unwrap_or([0; 4])) as usize;
|
||||
let start = offset + 8;
|
||||
let end = start.checked_add(size).ok_or("WAV chunk size overflow")?;
|
||||
if end > bytes.len() {
|
||||
return Err("Selected WAV contains a truncated chunk".to_string());
|
||||
}
|
||||
if id == b"fmt " {
|
||||
format = Some(&bytes[start..end]);
|
||||
} else if id == b"data" {
|
||||
data = Some(&bytes[start..end]);
|
||||
}
|
||||
offset = end + (size & 1);
|
||||
}
|
||||
let format = format.ok_or("Selected WAV has no format chunk")?;
|
||||
let data = data.ok_or("Selected WAV has no audio data")?;
|
||||
if format.len() < 16 {
|
||||
return Err("Selected WAV has an invalid format chunk".to_string());
|
||||
}
|
||||
let encoding = u16::from_le_bytes(format[0..2].try_into().unwrap_or([0; 2]));
|
||||
let encoding = if encoding == 0xfffe && format.len() >= 40 {
|
||||
u16::from_le_bytes(format[24..26].try_into().unwrap_or([0; 2]))
|
||||
} else {
|
||||
encoding
|
||||
};
|
||||
let channels = u16::from_le_bytes(format[2..4].try_into().unwrap_or([0; 2]));
|
||||
let sample_rate = u32::from_le_bytes(format[4..8].try_into().unwrap_or([0; 4]));
|
||||
let block_align = u16::from_le_bytes(format[12..14].try_into().unwrap_or([0; 2])) as usize;
|
||||
let bits = u16::from_le_bytes(format[14..16].try_into().unwrap_or([0; 2]));
|
||||
if channels == 0 || channels > 8 {
|
||||
return Err("Voice WAV must contain between 1 and 8 channels".to_string());
|
||||
}
|
||||
if !(MIN_SAMPLE_RATE..=MAX_SAMPLE_RATE).contains(&sample_rate) {
|
||||
return Err("Voice WAV sample rate must be between 8 and 96 kHz".to_string());
|
||||
}
|
||||
let bytes_per_sample = usize::from(bits.div_ceil(8));
|
||||
if block_align != bytes_per_sample * usize::from(channels)
|
||||
|| block_align == 0
|
||||
|| data.len() % block_align != 0
|
||||
{
|
||||
return Err("Voice WAV has invalid sample alignment".to_string());
|
||||
}
|
||||
if !matches!((encoding, bits), (1, 8 | 16 | 24 | 32) | (3, 32)) {
|
||||
return Err("Voice WAV must contain PCM or 32-bit float audio".to_string());
|
||||
}
|
||||
let frames = data.len() / block_align;
|
||||
let duration = frames as f64 / f64::from(sample_rate);
|
||||
if !(MIN_DURATION_SECONDS..=MAX_DURATION_SECONDS).contains(&duration) {
|
||||
return Err("Voice WAV must be between 2 and 30 seconds long".to_string());
|
||||
}
|
||||
|
||||
let mut samples = Vec::with_capacity(frames);
|
||||
for frame in data.chunks_exact(block_align) {
|
||||
let mut mono = 0.0_f32;
|
||||
for chunk in frame.chunks_exact(bytes_per_sample) {
|
||||
let sample = match (encoding, bits) {
|
||||
(1, 8) => (f32::from(chunk[0]) - 128.0) / 128.0,
|
||||
(1, 16) => f32::from(i16::from_le_bytes([chunk[0], chunk[1]])) / 32768.0,
|
||||
(1, 24) => {
|
||||
let raw = i32::from_le_bytes([
|
||||
chunk[0],
|
||||
chunk[1],
|
||||
chunk[2],
|
||||
if chunk[2] & 0x80 == 0 { 0 } else { 0xff },
|
||||
]);
|
||||
raw as f32 / 8_388_608.0
|
||||
}
|
||||
(1, 32) => {
|
||||
i32::from_le_bytes(chunk.try_into().map_err(|_| "invalid PCM sample")?) as f32
|
||||
/ 2_147_483_648.0
|
||||
}
|
||||
(3, 32) => f32::from_le_bytes(
|
||||
chunk
|
||||
.try_into()
|
||||
.map_err(|_| "invalid floating-point sample")?,
|
||||
),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
if !sample.is_finite() {
|
||||
return Err("Voice WAV contains non-finite samples".to_string());
|
||||
}
|
||||
mono += sample;
|
||||
}
|
||||
samples.push((mono / f32::from(channels)).clamp(-1.0, 1.0));
|
||||
}
|
||||
let stats = PcmStats::analyze(&samples, sample_rate);
|
||||
if !stats.is_non_silent() {
|
||||
return Err("Voice WAV is silent or too quiet to clone".to_string());
|
||||
}
|
||||
Ok(DecodedAudio {
|
||||
sample_rate,
|
||||
samples,
|
||||
})
|
||||
}
|
||||
|
||||
fn decode_media(source: &Path, extension: &str) -> Result<DecodedAudio, String> {
|
||||
let supported = ["m4a", "mp3", "flac", "ogg", "oga", "aif", "aiff"];
|
||||
if !supported.contains(&extension) {
|
||||
return Err(format!(
|
||||
"Unsupported voice audio format .{extension}. Choose WAV, M4A, MP3, FLAC, OGG, or AIFF"
|
||||
));
|
||||
}
|
||||
|
||||
let file = fs::File::open(source)
|
||||
.map_err(|error| format!("could not read selected audio: {error}"))?;
|
||||
let media = MediaSourceStream::new(Box::new(file), Default::default());
|
||||
let mut hint = Hint::new();
|
||||
hint.with_extension(extension);
|
||||
let probed = symphonia::default::get_probe()
|
||||
.format(
|
||||
&hint,
|
||||
media,
|
||||
&FormatOptions::default(),
|
||||
&MetadataOptions::default(),
|
||||
)
|
||||
.map_err(|error| format!("could not recognize selected audio: {error}"))?;
|
||||
let mut format = probed.format;
|
||||
let track = format
|
||||
.default_track()
|
||||
.ok_or_else(|| "Selected audio has no decodable track".to_string())?;
|
||||
let track_id = track.id;
|
||||
let mut decoder = symphonia::default::get_codecs()
|
||||
.make(&track.codec_params, &DecoderOptions::default())
|
||||
.map_err(|error| format!("could not initialize audio decoder: {error}"))?;
|
||||
let mut sample_rate = None;
|
||||
let mut samples = Vec::new();
|
||||
|
||||
loop {
|
||||
let packet = match format.next_packet() {
|
||||
Ok(packet) => packet,
|
||||
Err(SymphoniaError::ResetRequired) => {
|
||||
return Err("Selected audio changes format mid-stream".to_string());
|
||||
}
|
||||
Err(SymphoniaError::IoError(error))
|
||||
if error.kind() == std::io::ErrorKind::UnexpectedEof =>
|
||||
{
|
||||
break;
|
||||
}
|
||||
Err(error) => return Err(format!("could not read selected audio: {error}")),
|
||||
};
|
||||
if packet.track_id() != track_id {
|
||||
continue;
|
||||
}
|
||||
let decoded = match decoder.decode(&packet) {
|
||||
Ok(decoded) => decoded,
|
||||
Err(SymphoniaError::DecodeError(_)) => continue,
|
||||
Err(error) => return Err(format!("could not decode selected audio: {error}")),
|
||||
};
|
||||
let spec = *decoded.spec();
|
||||
if !(MIN_SAMPLE_RATE..=MAX_SAMPLE_RATE).contains(&spec.rate) {
|
||||
return Err("Voice audio sample rate must be between 8 and 96 kHz".to_string());
|
||||
}
|
||||
if sample_rate.is_some_and(|rate| rate != spec.rate) {
|
||||
return Err("Selected audio changes sample rate mid-stream".to_string());
|
||||
}
|
||||
sample_rate = Some(spec.rate);
|
||||
let channels = spec.channels.count();
|
||||
if channels == 0 || channels > 8 {
|
||||
return Err("Voice audio must contain between 1 and 8 channels".to_string());
|
||||
}
|
||||
let mut buffer = SampleBuffer::<f32>::new(decoded.capacity() as u64, spec);
|
||||
buffer.copy_interleaved_ref(decoded);
|
||||
for frame in buffer.samples().chunks_exact(channels) {
|
||||
let mono = frame.iter().copied().sum::<f32>() / channels as f32;
|
||||
if !mono.is_finite() {
|
||||
return Err("Voice audio contains non-finite samples".to_string());
|
||||
}
|
||||
samples.push(mono.clamp(-1.0, 1.0));
|
||||
}
|
||||
if samples.len() as f64 > MAX_DURATION_SECONDS * f64::from(spec.rate) {
|
||||
return Err("Voice audio must be between 2 and 30 seconds long".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let sample_rate =
|
||||
sample_rate.ok_or_else(|| "Selected audio contains no samples".to_string())?;
|
||||
validate_decoded_audio(&samples, sample_rate)?;
|
||||
Ok(DecodedAudio {
|
||||
sample_rate,
|
||||
samples,
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_decoded_audio(samples: &[f32], sample_rate: u32) -> Result<(), String> {
|
||||
let stats = PcmStats::analyze(samples, sample_rate);
|
||||
if !(MIN_DURATION_SECONDS..=MAX_DURATION_SECONDS).contains(&stats.duration_seconds) {
|
||||
return Err("Voice audio must be between 2 and 30 seconds long".to_string());
|
||||
}
|
||||
if !stats.is_non_silent() {
|
||||
return Err("Voice audio is silent or too quiet to clone".to_string());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn resample_linear(samples: &[f32], source_rate: u32) -> Vec<f32> {
|
||||
if source_rate == CANONICAL_SAMPLE_RATE {
|
||||
return samples.to_vec();
|
||||
}
|
||||
let output_len = ((samples.len() as u64 * u64::from(CANONICAL_SAMPLE_RATE)
|
||||
+ u64::from(source_rate) / 2)
|
||||
/ u64::from(source_rate)) as usize;
|
||||
(0..output_len)
|
||||
.map(|index| {
|
||||
let source = index as f64 * f64::from(source_rate) / f64::from(CANONICAL_SAMPLE_RATE);
|
||||
let left = source.floor() as usize;
|
||||
let fraction = (source - left as f64) as f32;
|
||||
let a = samples[left.min(samples.len() - 1)];
|
||||
let b = samples[(left + 1).min(samples.len() - 1)];
|
||||
a + (b - a) * fraction
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn encode_pcm16_wav(samples: &[f32], sample_rate: u32) -> Vec<u8> {
|
||||
let data_len = (samples.len() * 2) as u32;
|
||||
let mut bytes = Vec::with_capacity(44 + data_len as usize);
|
||||
bytes.extend_from_slice(b"RIFF");
|
||||
bytes.extend_from_slice(&(36 + data_len).to_le_bytes());
|
||||
bytes.extend_from_slice(b"WAVEfmt ");
|
||||
bytes.extend_from_slice(&16_u32.to_le_bytes());
|
||||
bytes.extend_from_slice(&1_u16.to_le_bytes());
|
||||
bytes.extend_from_slice(&1_u16.to_le_bytes());
|
||||
bytes.extend_from_slice(&sample_rate.to_le_bytes());
|
||||
bytes.extend_from_slice(&(sample_rate * 2).to_le_bytes());
|
||||
bytes.extend_from_slice(&2_u16.to_le_bytes());
|
||||
bytes.extend_from_slice(&16_u16.to_le_bytes());
|
||||
bytes.extend_from_slice(b"data");
|
||||
bytes.extend_from_slice(&data_len.to_le_bytes());
|
||||
for sample in samples {
|
||||
let value = (sample.clamp(-1.0, 1.0) * f32::from(i16::MAX)).round() as i16;
|
||||
bytes.extend_from_slice(&value.to_le_bytes());
|
||||
}
|
||||
bytes
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn fixture(sample_rate: u32, seconds: usize, amplitude: f32) -> Vec<u8> {
|
||||
let samples = (0..sample_rate as usize * seconds)
|
||||
.map(|index| {
|
||||
amplitude
|
||||
* (std::f32::consts::TAU * 220.0 * index as f32 / sample_rate as f32).sin()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
encode_pcm16_wav(&samples, sample_rate)
|
||||
}
|
||||
|
||||
fn stereo_fixture(sample_rate: u32, seconds: usize, amplitude: f32) -> Vec<u8> {
|
||||
let mono = fixture(sample_rate, seconds, amplitude);
|
||||
let mono_data = &mono[44..];
|
||||
let mut stereo_data = Vec::with_capacity(mono_data.len() * 2);
|
||||
for sample in mono_data.chunks_exact(2) {
|
||||
stereo_data.extend_from_slice(sample);
|
||||
stereo_data.extend_from_slice(sample);
|
||||
}
|
||||
let mut stereo = mono[..44].to_vec();
|
||||
stereo[4..8].copy_from_slice(&(36 + stereo_data.len() as u32).to_le_bytes());
|
||||
stereo[22..24].copy_from_slice(&2_u16.to_le_bytes());
|
||||
stereo[28..32].copy_from_slice(&(sample_rate * 4).to_le_bytes());
|
||||
stereo[32..34].copy_from_slice(&4_u16.to_le_bytes());
|
||||
stereo[40..44].copy_from_slice(&(stereo_data.len() as u32).to_le_bytes());
|
||||
stereo.extend_from_slice(&stereo_data);
|
||||
stereo
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn imports_persists_reloads_and_deletes_canonical_voice() {
|
||||
let temp = tempfile::tempdir().expect("temp voice workspace");
|
||||
let source = temp.path().join("My voice.wav");
|
||||
fs::write(&source, fixture(44_100, 2, 0.5)).expect("write source");
|
||||
let library = PocketVoiceLibrary::new(temp.path().join("library"));
|
||||
|
||||
let imported = library.import_path(&source).expect("import voice");
|
||||
assert!(imported.key.starts_with("pocket:imported:"));
|
||||
assert_eq!(imported.display_name, "My voice");
|
||||
|
||||
let relaunched = PocketVoiceLibrary::new(library.root());
|
||||
assert_eq!(
|
||||
relaunched.load().expect("reload registry"),
|
||||
vec![imported.clone()]
|
||||
);
|
||||
let stored = relaunched
|
||||
.resolve_file(&imported)
|
||||
.expect("resolve stored voice");
|
||||
let decoded = decode_wav(&fs::read(&stored).expect("read stored voice"))
|
||||
.expect("decode canonical voice");
|
||||
assert_eq!(decoded.sample_rate, CANONICAL_SAMPLE_RATE);
|
||||
assert_eq!(decoded.samples.len(), CANONICAL_SAMPLE_RATE as usize * 2);
|
||||
|
||||
assert_eq!(
|
||||
relaunched.import_path(&source).expect("idempotent import"),
|
||||
imported
|
||||
);
|
||||
assert_eq!(relaunched.load().expect("deduplicated registry").len(), 1);
|
||||
|
||||
relaunched.delete(&imported.key).expect("delete voice");
|
||||
assert!(relaunched.load().expect("empty registry").is_empty());
|
||||
assert!(!stored.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn common_stereo_audio_is_downmixed_to_canonical_mono() {
|
||||
let temp = tempfile::tempdir().expect("temp voice workspace");
|
||||
let source = temp.path().join("stereo.wav");
|
||||
fs::write(&source, stereo_fixture(44_100, 2, 0.5)).expect("write stereo");
|
||||
let library = PocketVoiceLibrary::new(temp.path().join("library"));
|
||||
|
||||
let imported = library.import_path(&source).expect("import stereo");
|
||||
let stored = library
|
||||
.resolve_file(&imported)
|
||||
.expect("resolve stored voice");
|
||||
let decoded = decode_wav(&fs::read(stored).expect("read stored voice"))
|
||||
.expect("decode canonical voice");
|
||||
assert_eq!(decoded.sample_rate, CANONICAL_SAMPLE_RATE);
|
||||
assert_eq!(decoded.samples.len(), CANONICAL_SAMPLE_RATE as usize * 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires BUZZ_VOICE_IMPORT_TEST_DIR with common-format fixtures"]
|
||||
fn imports_common_audio_format_fixtures() {
|
||||
let fixtures =
|
||||
PathBuf::from(std::env::var("BUZZ_VOICE_IMPORT_TEST_DIR").expect("fixture directory"));
|
||||
let temp = tempfile::tempdir().expect("temp voice workspace");
|
||||
let library = PocketVoiceLibrary::new(temp.path().join("library"));
|
||||
|
||||
for file_name in [
|
||||
"voice.wav",
|
||||
"voice.m4a",
|
||||
"voice.mp3",
|
||||
"voice.flac",
|
||||
"voice.ogg",
|
||||
"voice.aiff",
|
||||
] {
|
||||
let imported = library
|
||||
.import_path(&fixtures.join(file_name))
|
||||
.unwrap_or_else(|error| panic!("import {file_name}: {error}"));
|
||||
let stored = library
|
||||
.resolve_file(&imported)
|
||||
.unwrap_or_else(|error| panic!("resolve {file_name}: {error}"));
|
||||
let decoded = decode_wav(&fs::read(stored).expect("read canonical voice"))
|
||||
.expect("decode canonical voice");
|
||||
assert_eq!(decoded.sample_rate, CANONICAL_SAMPLE_RATE);
|
||||
assert!(decoded.samples.len() >= CANONICAL_SAMPLE_RATE as usize * 2);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_unsupported_and_silent_files_do_not_mutate_registry() {
|
||||
let temp = tempfile::tempdir().expect("temp voice workspace");
|
||||
let library = PocketVoiceLibrary::new(temp.path().join("library"));
|
||||
|
||||
let garbage = temp.path().join("garbage.wav");
|
||||
fs::write(&garbage, b"not a wave").expect("write garbage");
|
||||
assert!(library
|
||||
.import_path(&garbage)
|
||||
.expect_err("garbage rejected")
|
||||
.contains("RIFF/WAVE"));
|
||||
|
||||
let silent = temp.path().join("silent.wav");
|
||||
fs::write(&silent, fixture(32_000, 2, 0.0)).expect("write silence");
|
||||
assert!(library
|
||||
.import_path(&silent)
|
||||
.expect_err("silence rejected")
|
||||
.contains("silent"));
|
||||
|
||||
let unsupported_container = temp.path().join("voice.txt");
|
||||
fs::write(&unsupported_container, b"not audio").expect("write unsupported container");
|
||||
assert!(library
|
||||
.import_path(&unsupported_container)
|
||||
.expect_err("container rejected")
|
||||
.contains("Unsupported voice audio format"));
|
||||
|
||||
let mut unsupported = fixture(32_000, 2, 0.5);
|
||||
unsupported[20..22].copy_from_slice(&6_u16.to_le_bytes());
|
||||
let unsupported_path = temp.path().join("unsupported.wav");
|
||||
fs::write(&unsupported_path, unsupported).expect("write unsupported");
|
||||
assert!(library
|
||||
.import_path(&unsupported_path)
|
||||
.expect_err("unsupported rejected")
|
||||
.contains("PCM or 32-bit float"));
|
||||
|
||||
assert!(library.load().expect("unchanged registry").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pcm_analysis_distinguishes_signal_from_silence() {
|
||||
let signal = (0..24_000)
|
||||
.map(|index| (std::f32::consts::TAU * 440.0 * index as f32 / 24_000.0).sin() * 0.5)
|
||||
.collect::<Vec<_>>();
|
||||
let signal_stats = PcmStats::analyze(&signal, 24_000);
|
||||
assert!(signal_stats.is_non_silent());
|
||||
assert_eq!(signal_stats.duration_seconds, 1.0);
|
||||
assert!(signal_stats.peak > 0.49);
|
||||
assert!(signal_stats.rms > 0.3);
|
||||
|
||||
let silence = vec![0.0; 24_000];
|
||||
assert!(!PcmStats::analyze(&silence, 24_000).is_non_silent());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
//! Reusable local voice primitives for Buzz.
|
||||
|
||||
pub mod imported;
|
||||
pub mod pocket;
|
||||
|
||||
pub use pocket::{
|
||||
april_model_info, load_text_to_speech, load_voice_style, PocketModelInfo, PocketTts,
|
||||
VoiceStyle, DEFAULT_VOICE, SAMPLE_RATE, VOICE_FILE_EXT,
|
||||
};
|
||||
|
||||
/// One immutable artifact required by the April Pocket bundle.
|
||||
///
|
||||
/// `filename` is the bundle-relative file name, `sha256` pins its contents,
|
||||
/// `size_bytes` supports download progress and validation, and `quantized`
|
||||
/// identifies the INT8 components.
|
||||
pub type PocketModelArtifact = pocket::PocketModelArtifact;
|
||||
|
||||
/// Language bundle selected from the pinned export.
|
||||
pub const APRIL_BUNDLE_ID: &str = pocket::APRIL_BUNDLE_ID;
|
||||
/// Pinned upstream export repository.
|
||||
pub const APRIL_MODEL_ID: &str = pocket::APRIL_MODEL_ID;
|
||||
/// Pinned revision containing the April bundle.
|
||||
pub const APRIL_MODEL_REVISION: &str = pocket::APRIL_MODEL_REVISION;
|
||||
@@ -0,0 +1,167 @@
|
||||
//! April 2026 Pocket TTS engine for Buzz Desktop.
|
||||
//!
|
||||
//! 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
|
||||
//!
|
||||
//! - 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.
|
||||
//!
|
||||
//! `huddle::models` writes the complete attribution beside the cached bytes.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Mutex;
|
||||
|
||||
use sherpa_onnx::Wave;
|
||||
|
||||
#[path = "pocket_april.rs"]
|
||||
mod pocket_april;
|
||||
#[path = "pocket_models.rs"]
|
||||
mod pocket_models;
|
||||
|
||||
use pocket_april::{prepare_april_prompt, AprilPocketTts};
|
||||
pub use pocket_models::{
|
||||
april_model_info, PocketModelArtifact, PocketModelInfo, APRIL_BUNDLE_ID, APRIL_MODEL_ID,
|
||||
APRIL_MODEL_REVISION,
|
||||
};
|
||||
|
||||
/// Pocket TTS emits 24 kHz mono PCM.
|
||||
pub const SAMPLE_RATE: u32 = 24_000;
|
||||
|
||||
/// Bundled reference voice name without its extension.
|
||||
pub const DEFAULT_VOICE: &str = "reference_sample";
|
||||
|
||||
/// Pocket voice files are reference WAVs.
|
||||
pub const VOICE_FILE_EXT: &str = "wav";
|
||||
|
||||
const TTS_NUM_THREADS: usize = 1;
|
||||
|
||||
/// Loaded reference voice samples and their original sample rate.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct VoiceStyle {
|
||||
samples: Vec<f32>,
|
||||
sample_rate: i32,
|
||||
}
|
||||
|
||||
/// Load a Pocket reference voice WAV from disk.
|
||||
pub fn load_voice_style(path: &Path) -> Result<VoiceStyle, String> {
|
||||
let path_str = path
|
||||
.to_str()
|
||||
.ok_or_else(|| format!("voice path is not valid UTF-8: {}", path.display()))?;
|
||||
let wave = Wave::read(path_str)
|
||||
.ok_or_else(|| format!("could not read voice WAV at {}", path.display()))?;
|
||||
let samples = wave.samples().to_vec();
|
||||
if samples.is_empty() {
|
||||
return Err(format!("voice WAV is empty: {}", path.display()));
|
||||
}
|
||||
Ok(VoiceStyle {
|
||||
samples,
|
||||
sample_rate: wave.sample_rate(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Resident April INT8 Pocket TTS engine.
|
||||
pub struct PocketTts {
|
||||
inner: Mutex<AprilPocketTts>,
|
||||
}
|
||||
|
||||
/// 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 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()
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(PocketTts {
|
||||
inner: Mutex::new(AprilPocketTts::load(&dir, TTS_NUM_THREADS)?),
|
||||
})
|
||||
}
|
||||
|
||||
impl PocketTts {
|
||||
/// Split text into synthesis units that satisfy the bundle's exact
|
||||
/// 50-token input limit.
|
||||
pub fn split_text_into_chunks(&self, text: &str) -> Result<Vec<String>, String> {
|
||||
let Some(prepared) = prepare_april_prompt(text) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
self.inner
|
||||
.lock()
|
||||
.map_err(|_| "Pocket TTS engine lock poisoned".to_string())?
|
||||
.split_prompt(&prepared)
|
||||
}
|
||||
|
||||
/// Synthesize text with the supplied reference voice.
|
||||
///
|
||||
/// Pocket detects language from text and this model uses one synthesis
|
||||
/// step, so `_lang` and `_steps` intentionally do not affect output.
|
||||
pub fn synth_chunk(
|
||||
&self,
|
||||
text: &str,
|
||||
_lang: &str,
|
||||
style: &VoiceStyle,
|
||||
_steps: usize,
|
||||
) -> Result<Vec<f32>, String> {
|
||||
let Some(prepared) = prepare_april_prompt(text) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let mut engine = self
|
||||
.inner
|
||||
.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(samples)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
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]
|
||||
#[ignore = "requires BUZZ_POCKET_TEST_MODEL_DIR"]
|
||||
fn production_api_emits_non_silent_april_int8_pcm() {
|
||||
let dir = std::env::var("BUZZ_POCKET_TEST_MODEL_DIR")
|
||||
.expect("set BUZZ_POCKET_TEST_MODEL_DIR to an April INT8 model directory");
|
||||
let engine = load_text_to_speech(&dir).expect("load April INT8 engine");
|
||||
let style = load_voice_style(&Path::new(&dir).join("reference_sample.wav"))
|
||||
.expect("load reference voice");
|
||||
let samples = engine
|
||||
.synth_chunk("Bright birds begin beside the bay.", "en", &style, 1)
|
||||
.expect("synthesize through the production API");
|
||||
|
||||
assert!(!samples.is_empty());
|
||||
assert!(samples.iter().all(|sample| sample.is_finite()));
|
||||
assert!(samples.iter().any(|sample| sample.abs() > 1.0e-6));
|
||||
}
|
||||
}
|
||||
@@ -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]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
//! 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 {
|
||||
/// Language bundle selected from the pinned export.
|
||||
pub bundle_id: &'static str,
|
||||
/// Upstream model repository.
|
||||
pub source_model_id: &'static str,
|
||||
/// Pinned upstream model revision.
|
||||
pub revision: &'static str,
|
||||
/// PCM output sample rate.
|
||||
pub sample_rate: u32,
|
||||
/// Maximum input size declared by the bundle.
|
||||
pub max_token_per_chunk: usize,
|
||||
/// Immutable files required by the runtime.
|
||||
pub artifacts: &'static [PocketModelArtifact],
|
||||
/// Components quantized in the selected bundle.
|
||||
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"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
use std::{
|
||||
fs,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
||||
use buzz_voice::{
|
||||
imported::{write_pcm16_wav, PcmStats, PocketVoiceLibrary},
|
||||
pocket::{load_text_to_speech, load_voice_style, DEFAULT_VOICE, SAMPLE_RATE, VOICE_FILE_EXT},
|
||||
};
|
||||
|
||||
const PREVIEW_TEXT: &str = "This is an objective Pocket voice preview.";
|
||||
|
||||
fn required_path(name: &str) -> PathBuf {
|
||||
std::env::var_os(name)
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| panic!("{name} must point to the required local test path"))
|
||||
}
|
||||
|
||||
fn checked_in_voice() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
.join("../../desktop/src-tauri/resources/pocket-voices/eve.wav")
|
||||
}
|
||||
|
||||
fn evidence_dir() -> PathBuf {
|
||||
std::env::var_os("BUZZ_VOICE_EVIDENCE_DIR")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("../../target/buzz-voice-evidence")
|
||||
})
|
||||
}
|
||||
|
||||
fn synthesize(model_dir: &Path, voice_path: &Path, text: &str) -> (Vec<f32>, PcmStats) {
|
||||
let engine = load_text_to_speech(
|
||||
model_dir
|
||||
.to_str()
|
||||
.expect("Pocket model path must be valid UTF-8"),
|
||||
)
|
||||
.expect("load Pocket model");
|
||||
let style = load_voice_style(voice_path).expect("load selected voice");
|
||||
let samples = engine
|
||||
.synth_chunk(text, "en", &style, 1)
|
||||
.expect("synthesize preview");
|
||||
let stats = PcmStats::analyze(&samples, SAMPLE_RATE);
|
||||
assert!(
|
||||
stats.is_non_silent(),
|
||||
"generated PCM must be non-silent: {stats:?}"
|
||||
);
|
||||
assert!(
|
||||
stats.duration_seconds > 0.2,
|
||||
"generated PCM is unexpectedly short: {stats:?}"
|
||||
);
|
||||
(samples, stats)
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires BUZZ_POCKET_MODEL_DIR and runs the installed Pocket ONNX model"]
|
||||
fn objective_import_synthesis_delete_and_mary_fallback() {
|
||||
let model_dir = required_path("BUZZ_POCKET_MODEL_DIR");
|
||||
let temp = tempfile::tempdir().expect("temporary voice workspace");
|
||||
let source = temp.path().join("Imported Eve.wav");
|
||||
fs::copy(checked_in_voice(), &source).expect("copy checked-in voice fixture");
|
||||
|
||||
let library_root = temp.path().join("library");
|
||||
let library = PocketVoiceLibrary::new(&library_root);
|
||||
let imported = library.import_path(&source).expect("import valid WAV");
|
||||
assert_eq!(
|
||||
library.find(&imported.key).expect("read selection"),
|
||||
Some(imported.clone())
|
||||
);
|
||||
|
||||
drop(library);
|
||||
let relaunched = PocketVoiceLibrary::new(&library_root);
|
||||
let selected = relaunched
|
||||
.find(&imported.key)
|
||||
.expect("reload persisted selection")
|
||||
.expect("selected imported voice survived relaunch");
|
||||
let imported_path = relaunched
|
||||
.resolve_file(&selected)
|
||||
.expect("resolve persisted imported voice");
|
||||
let (imported_pcm, imported_stats) = synthesize(&model_dir, &imported_path, PREVIEW_TEXT);
|
||||
|
||||
let evidence = evidence_dir();
|
||||
fs::create_dir_all(&evidence).expect("create evidence directory");
|
||||
let imported_wav = evidence.join("imported-preview.wav");
|
||||
write_pcm16_wav(&imported_wav, &imported_pcm, SAMPLE_RATE)
|
||||
.expect("write imported preview evidence");
|
||||
|
||||
relaunched
|
||||
.delete(&imported.key)
|
||||
.expect("delete imported voice");
|
||||
assert_eq!(
|
||||
relaunched.find(&imported.key).expect("reload after delete"),
|
||||
None
|
||||
);
|
||||
|
||||
let mary_path = model_dir.join(format!("{DEFAULT_VOICE}.{VOICE_FILE_EXT}"));
|
||||
assert_eq!(
|
||||
mary_path.file_name().and_then(|name| name.to_str()),
|
||||
Some("reference_sample.wav"),
|
||||
"fallback must remain the deterministic Mary reference"
|
||||
);
|
||||
let (mary_pcm, mary_stats) = synthesize(&model_dir, &mary_path, PREVIEW_TEXT);
|
||||
let mary_wav = evidence.join("mary-fallback-preview.wav");
|
||||
write_pcm16_wav(&mary_wav, &mary_pcm, SAMPLE_RATE)
|
||||
.expect("write Mary fallback preview evidence");
|
||||
|
||||
println!(
|
||||
"{}",
|
||||
serde_json::json!({
|
||||
"importedKey": imported.key,
|
||||
"persistence": "reloaded",
|
||||
"afterDelete": "pocket:mary",
|
||||
"importedPreview": {
|
||||
"path": imported_wav,
|
||||
"samples": imported_stats.sample_count,
|
||||
"sampleRate": imported_stats.sample_rate,
|
||||
"durationSeconds": imported_stats.duration_seconds,
|
||||
"peak": imported_stats.peak,
|
||||
"rms": imported_stats.rms,
|
||||
"nonSilentSamples": imported_stats.non_silent_samples,
|
||||
},
|
||||
"maryFallbackPreview": {
|
||||
"path": mary_wav,
|
||||
"samples": mary_stats.sample_count,
|
||||
"sampleRate": mary_stats.sample_rate,
|
||||
"durationSeconds": mary_stats.duration_seconds,
|
||||
"peak": mary_stats.peak,
|
||||
"rms": mary_stats.rms,
|
||||
"nonSilentSamples": mary_stats.non_silent_samples,
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user