1use super::adapter::{
9 TrustMode, ADAPTER_FAKE_SINE_V1, ADAPTER_KITTEN_ONNX_V1, ADAPTER_KOKORO_ONNX_V0,
10};
11use super::catalogue::{
12 ensure_voice_pack, lookup_model, onnx_path, resolve_voice_for_model, validate_speaking_rate,
13 voices_path,
14};
15use super::chunk::{prepare_tts_chunks_with, PhonemeCodec, TtsChunk, CHUNK_PAUSE_MS};
16use super::conformance::synthesize_fake_sine_ms;
17use super::npz::load_voices_npz;
18use super::pack::{load_pack_dir, stage_verified_for_load};
19use super::pcm_post::{
20 duration_ms_from_pcm, trim_trailing_silence, validate_raw_pcm, TailTrimPolicy, PEAK_LIMIT,
21};
22use super::provider::{BackendKind, SynthesisOptions, SynthesisProvider, SynthesisResult};
23use super::validate::{
24 normalize_tts_language, prepare_text, resolve_sample_rate, DEFAULT_MAX_CHARS,
25};
26use super::wav::peak_guard_f32_to_i16;
27use crate::error::{ProviderError, Result, UserError};
28use crate::runtime::{
29 LoadKey, ModelRegistry, OpContext, RegistryConfig, RegistryPin, ResidencyWeight,
30 ResourceGovernor, Singleflight,
31};
32use async_trait::async_trait;
33use once_cell::sync::Lazy;
34use ort::session::Session;
35use ort::value::Tensor;
36use std::collections::HashMap;
37use std::path::{Path, PathBuf};
38use std::sync::{Arc, Mutex};
39use std::time::Duration;
40
41pub struct TtsSessionPool {
45 flight: Singleflight<LoadedPack>,
46 registry: ModelRegistry<LoadedPack>,
47}
48
49impl Default for TtsSessionPool {
50 fn default() -> Self {
51 Self::new()
52 }
53}
54
55impl TtsSessionPool {
56 pub fn new() -> Self {
57 Self {
58 flight: Singleflight::default(),
59 registry: ModelRegistry::new(RegistryConfig::default()),
60 }
61 }
62
63 pub fn resident_len(&self) -> usize {
64 self.registry.len()
65 }
66
67 fn weight_for(onnx: &Path) -> ResidencyWeight {
68 let disk = std::fs::metadata(onnx).map(|m| m.len()).unwrap_or(0);
69 let bytes = disk.saturating_mul(2).max(32 * 1024 * 1024);
71 ResidencyWeight { bytes }
72 }
73
74 fn get_or_load_pin<F>(
77 &self,
78 key: LoadKey,
79 weight: ResidencyWeight,
80 gov: &Arc<ResourceGovernor>,
81 loader: F,
82 ) -> Result<RegistryPin<LoadedPack>>
83 where
84 F: FnOnce() -> Result<LoadedPack>,
85 {
86 if weight.bytes > self.registry.config().max_resident_bytes {
87 return Err(ProviderError::Overload {
88 reason: format!(
89 "TTS pack weight {} exceeds residency budget {}",
90 weight.bytes,
91 self.registry.config().max_resident_bytes
92 ),
93 }
94 .into());
95 }
96 loop {
97 if let Some(pin) = self.registry.get_and_pin(&key) {
98 return Ok(pin);
99 }
100 match self.flight.begin_or_wait_guard(key.clone()) {
101 Ok(None) => continue,
102 Err(message) => {
103 return Err(ProviderError::ModelLoad {
104 model: key.id.clone(),
105 reason: message,
106 }
107 .into());
108 }
109 Ok(Some(leader)) => {
110 let gov = Arc::clone(gov);
111 let load_result = (|| -> Result<Arc<LoadedPack>> {
112 let _permit = gov.acquire(crate::runtime::PermitKind::ModelLoad, None)?;
113 Ok(Arc::new(loader()?))
114 })();
115 match load_result {
116 Ok(pack) => {
117 match self.registry.insert_and_pin(
118 key.clone(),
119 Arc::clone(&pack),
120 weight,
121 ) {
122 Ok(pin) => {
123 leader.success();
124 return Ok(pin);
125 }
126 Err(e) => {
127 leader.fail(e.to_string());
128 drop(pack);
129 return Err(e);
130 }
131 }
132 }
133 Err(e) => {
134 leader.fail(e.to_string());
135 return Err(e);
136 }
137 }
138 }
139 }
140 }
141 }
142
143 pub fn clear(&self) {
147 let _ = self.registry.clear_idle();
148 }
151}
152
153static PROCESS_TTS_POOL: Lazy<Arc<TtsSessionPool>> = Lazy::new(|| Arc::new(TtsSessionPool::new()));
154
155pub fn process_global_tts_pool() -> Arc<TtsSessionPool> {
157 Arc::clone(&PROCESS_TTS_POOL)
158}
159
160pub struct LocalTtsProvider {
171 cache_dir: PathBuf,
172 show_progress: bool,
173 local_only: bool,
174 max_chars: usize,
175 pool: Arc<TtsSessionPool>,
176 governor: Arc<ResourceGovernor>,
177}
178
179struct LoadedPack {
180 session: Mutex<Session>,
181 voices: HashMap<String, super::npz::VoiceMatrix>,
182 sample_rate_hz: u32,
183 max_phoneme_tokens: usize,
185 speed_priors: HashMap<String, f32>,
187}
188
189impl LocalTtsProvider {
190 pub fn new(cache_dir: PathBuf) -> Self {
191 Self::with_runtime(
192 cache_dir,
193 process_global_tts_pool(),
194 ResourceGovernor::process_global(),
195 )
196 }
197
198 pub fn with_runtime(
200 cache_dir: PathBuf,
201 pool: Arc<TtsSessionPool>,
202 governor: Arc<ResourceGovernor>,
203 ) -> Self {
204 Self {
205 cache_dir,
206 show_progress: false,
207 local_only: false,
208 max_chars: DEFAULT_MAX_CHARS,
209 pool,
210 governor,
211 }
212 }
213
214 pub fn with_progress(mut self, v: bool) -> Self {
215 self.show_progress = v;
216 self
217 }
218
219 pub fn with_local_only(mut self, v: bool) -> Self {
220 self.local_only = v;
221 self
222 }
223
224 pub fn with_max_chars(mut self, n: usize) -> Self {
225 self.max_chars = n.max(1);
226 self
227 }
228
229 pub fn pool(&self) -> &Arc<TtsSessionPool> {
230 &self.pool
231 }
232
233 pub fn clear_sessions(&self) {
241 self.pool.clear();
242 }
243
244 async fn ensure_loaded_pin(
245 &self,
246 model: &str,
247 local_only: bool,
248 ) -> Result<RegistryPin<LoadedPack>> {
249 let key = LoadKey::tts(model, self.cache_dir.join(model).display().to_string());
250 if let Some(pin) = self.pool.registry.get_and_pin(&key) {
251 return Ok(pin);
252 }
253
254 let info = lookup_model(model)?;
255 let _pack_dir = ensure_voice_pack(
256 &self.cache_dir,
257 model,
258 self.show_progress,
259 local_only || self.local_only,
260 )
261 .await?;
262
263 let onnx = onnx_path(&self.cache_dir, info);
264 let voices_file = voices_path(&self.cache_dir, info);
265 let speed_priors = load_speed_priors(&self.cache_dir, info);
266 let sample_rate = info.sample_rate_hz;
267 let catalogue_max = info.max_phoneme_tokens;
268 let key_for_load = key.clone();
269 let model_id = model.to_string();
270 let weight = TtsSessionPool::weight_for(&onnx);
271 let pool = Arc::clone(&self.pool);
272 let gov = Arc::clone(&self.governor);
273 let cache_root = self.cache_dir.clone();
274 let onnx_sha = Some(info.onnx.sha256.to_string());
275 let voices_sha = Some(info.voices.sha256.to_string());
276 let onnx_leaf = onnx
277 .file_name()
278 .and_then(|s| s.to_str())
279 .unwrap_or("model.onnx")
280 .to_string();
281 let voices_leaf = voices_file
282 .file_name()
283 .and_then(|s| s.to_str())
284 .unwrap_or("voices.npz")
285 .to_string();
286
287 let pin = tokio::task::spawn_blocking(move || {
288 pool.get_or_load_pin(key_for_load, weight, &gov, || {
289 let onnx_load =
290 stage_verified_for_load(&onnx, onnx_sha.as_deref(), &cache_root, &onnx_leaf)?;
291 let voices_load = stage_verified_for_load(
292 &voices_file,
293 voices_sha.as_deref(),
294 &cache_root,
295 &voices_leaf,
296 )?;
297 load_pack(
298 &onnx_load,
299 &voices_load,
300 sample_rate,
301 speed_priors,
302 catalogue_max,
303 &model_id,
304 )
305 })
306 })
307 .await
308 .map_err(|e| crate::error::TranscriptionError::internal(format!("TTS load join: {e}")))??;
309
310 Ok(pin)
311 }
312
313 async fn ensure_loaded_from_pack_pin(
316 &self,
317 pack_dir: &Path,
318 allow_unverified: bool,
319 ) -> Result<(RegistryPin<LoadedPack>, super::adapter::ModelPackManifest)> {
320 let (root, manifest) = load_pack_dir(pack_dir, allow_unverified)?;
321 if manifest.adapter_id != ADAPTER_KITTEN_ONNX_V1
322 && manifest.adapter_id != ADAPTER_KOKORO_ONNX_V0
323 {
324 return Err(UserError::InvalidConfig {
325 reason: format!(
326 "local pack override synthesis currently supports adapters \
327 '{ADAPTER_KITTEN_ONNX_V1}' and '{ADAPTER_KOKORO_ONNX_V0}' \
328 (got '{}'); use `aurum tts inspect` / conformance for others",
329 manifest.adapter_id
330 ),
331 }
332 .into());
333 }
334 let onnx_name = manifest
335 .artifact("onnx")
336 .map(|a| a.filename.clone())
337 .ok_or_else(|| UserError::InvalidConfig {
338 reason: "pack missing onnx artifact".into(),
339 })?;
340 let voices_name = manifest
341 .artifact("voices")
342 .map(|a| a.filename.clone())
343 .ok_or_else(|| UserError::InvalidConfig {
344 reason: "pack missing voices artifact".into(),
345 })?;
346 let config_name = manifest
347 .artifact("config")
348 .map(|a| a.filename.clone())
349 .unwrap_or_else(|| "config.json".into());
350 let onnx = root.join(&onnx_name);
351 let voices_file = root.join(&voices_name);
352 let onnx_sha = manifest.artifact("onnx").and_then(|a| a.sha256.clone());
353 let voices_sha = manifest.artifact("voices").and_then(|a| a.sha256.clone());
354 let sample_rate = manifest.sample_rate_hz;
355 let catalogue_max = manifest.max_phoneme_tokens;
356 let model_id = manifest.model_id.clone();
357 let key = LoadKey::tts(
359 format!("local-pack:{}", manifest.model_id),
360 root.display().to_string(),
361 );
362 if let Some(pin) = self.pool.registry.get_and_pin(&key) {
363 return Ok((pin, manifest));
364 }
365 let key_for_load = key.clone();
366 let speed_priors = load_speed_priors_from_path(&root.join(&config_name));
367 let weight = TtsSessionPool::weight_for(&onnx);
368 let pool = Arc::clone(&self.pool);
369 let gov = Arc::clone(&self.governor);
370 let cache_root = self.cache_dir.clone();
371
372 let pin = tokio::task::spawn_blocking(move || {
373 pool.get_or_load_pin(key_for_load, weight, &gov, || {
374 let onnx_load =
377 stage_verified_for_load(&onnx, onnx_sha.as_deref(), &cache_root, &onnx_name)?;
378 let voices_load = stage_verified_for_load(
379 &voices_file,
380 voices_sha.as_deref(),
381 &cache_root,
382 &voices_name,
383 )?;
384 load_pack(
385 &onnx_load,
386 &voices_load,
387 sample_rate,
388 speed_priors,
389 catalogue_max,
390 &model_id,
391 )
392 })
393 })
394 .await
395 .map_err(|e| {
396 crate::error::TranscriptionError::internal(format!("TTS pack load join: {e}"))
397 })??;
398
399 Ok((pin, manifest))
400 }
401}
402
403#[allow(clippy::too_many_arguments)]
404fn synthesize_with_pack(
405 pack: &LoadedPack,
406 text: &str,
407 opts: &SynthesisOptions,
408 voice_internal: &str,
409 voice_canonical: &str,
410 model_canonical: &str,
411 text_chars: usize,
412 op: &OpContext,
413 adapter: &str,
414 trust: TrustMode,
415 provenance: &str,
416) -> Result<SynthesisResult> {
417 op.check()?;
418
419 let rate = validate_speaking_rate(opts.speaking_rate)?;
420 let sample_rate = resolve_sample_rate(opts.sample_rate_hz, pack.sample_rate_hz)?;
421
422 let voice_mat = pack
423 .voices
424 .get(voice_internal)
425 .ok_or_else(|| UserError::Other {
426 message: format!(
427 "voice embedding '{voice_internal}' missing from pack; available: {:?}",
428 pack.voices.keys().collect::<Vec<_>>()
429 ),
430 })?;
431
432 let pack_max = voice_mat.nrows.saturating_sub(1);
435 let max_tokens = pack_max.min(pack.max_phoneme_tokens);
436 if max_tokens <= 2 {
437 return Err(ProviderError::ModelLoad {
438 model: model_canonical.to_string(),
439 reason: "voice embedding has no usable sequence rows".into(),
440 }
441 .into());
442 }
443
444 let codec = phoneme_codec_for_adapter(adapter);
445 let chunks = prepare_tts_chunks_with(text, max_tokens, codec)?;
446 let chunk_count = chunks.len();
447 let effective_speed = rate
448 * pack
449 .speed_priors
450 .get(voice_internal)
451 .copied()
452 .unwrap_or(1.0);
453 let pause_samples = (sample_rate as u64)
454 .saturating_mul(CHUNK_PAUSE_MS)
455 .checked_div(1_000)
456 .unwrap_or(0) as usize;
457
458 let mut pcm_f32: Vec<f32> = Vec::new();
459 let trim_policy = TailTrimPolicy::default();
460
461 for (index, chunk) in chunks.iter().enumerate() {
462 op.check()?;
463 let chunk_audio =
464 synthesize_chunk_f32(pack, voice_mat, chunk, effective_speed).map_err(|err| {
465 ProviderError::Other {
466 message: format!(
467 "TTS chunk {}/{} failed near {:?}: {err}",
468 index + 1,
469 chunks.len(),
470 chunk.text.chars().take(80).collect::<String>()
471 ),
472 }
473 })?;
474 let trimmed = trim_trailing_silence(&chunk_audio, trim_policy);
475 if index > 0 {
476 pcm_f32.resize(pcm_f32.len().saturating_add(pause_samples), 0.0);
477 }
478 pcm_f32.extend_from_slice(trimmed);
479 }
480
481 validate_raw_pcm(&pcm_f32, sample_rate)?;
482 let pcm = peak_guard_f32_to_i16(&pcm_f32, PEAK_LIMIT);
483 if pcm.is_empty() {
484 return Err(ProviderError::Other {
485 message: "synthesis produced empty audio after validation".into(),
486 }
487 .into());
488 }
489
490 let duration_ms = duration_ms_from_pcm(pcm.len(), sample_rate);
491 let language = normalize_tts_language(&opts.language).unwrap_or_else(|_| "en".into());
492
493 Ok(SynthesisResult {
494 pcm_i16_mono: pcm,
495 sample_rate_hz: sample_rate,
496 channels: 1,
497 backend_kind: BackendKind::Local,
498 provider: "local".into(),
499 model: model_canonical.to_string(),
500 voice: voice_canonical.to_string(),
501 language,
502 duration_ms,
503 text_chars,
504 text_truncated: false,
505 chunk_count,
506 synthesized_chars: text_chars,
507 adapter: Some(adapter.into()),
508 trust: Some(trust.as_str().into()),
509 provenance: Some(provenance.into()),
510 })
511}
512
513fn phoneme_codec_for_adapter(adapter: &str) -> PhonemeCodec {
514 if adapter == ADAPTER_KOKORO_ONNX_V0 {
515 PhonemeCodec::Kokoro
516 } else {
517 PhonemeCodec::Kitten
518 }
519}
520
521fn synthesize_chunk_f32(
522 pack: &LoadedPack,
523 voice_mat: &super::npz::VoiceMatrix,
524 chunk: &TtsChunk,
525 effective_speed: f32,
526) -> Result<Vec<f32>> {
527 let seq_len = chunk.ids.len();
528 let style = voice_mat.style_row(seq_len).to_vec();
530 let style_dim = style.len();
531 let t_ids = Tensor::<i64>::from_array(([1usize, seq_len], chunk.ids.clone())).map_err(|e| {
532 ProviderError::Other {
533 message: format!("tokens/input_ids tensor: {e}"),
534 }
535 })?;
536 let t_style = Tensor::<f32>::from_array(([1usize, style_dim], style)).map_err(|e| {
537 ProviderError::Other {
538 message: format!("style tensor: {e}"),
539 }
540 })?;
541 let t_speed = Tensor::<f32>::from_array(([1usize], vec![effective_speed])).map_err(|e| {
542 ProviderError::Other {
543 message: format!("speed tensor: {e}"),
544 }
545 })?;
546
547 let mut session = pack
548 .session
549 .lock()
550 .map_err(|_| crate::error::TranscriptionError::internal("ORT session mutex poisoned"))?;
551 let outputs = session
553 .run(ort::inputs![t_ids, t_style, t_speed])
554 .map_err(|e| ProviderError::Other {
555 message: format!("ONNX inference failed: {e}"),
556 })?;
557 let (_shape, audio_data) =
558 outputs[0]
559 .try_extract_tensor::<f32>()
560 .map_err(|e| ProviderError::Other {
561 message: format!("extract audio tensor: {e}"),
562 })?;
563 Ok(audio_data.to_vec())
564}
565
566fn load_pack(
567 onnx: &Path,
568 voices_file: &Path,
569 sample_rate_hz: u32,
570 speed_priors: HashMap<String, f32>,
571 catalogue_max_tokens: usize,
572 model_label: &str,
573) -> Result<LoadedPack> {
574 let session = Session::builder()
575 .map_err(|e| ProviderError::ModelLoad {
576 model: model_label.to_string(),
577 reason: format!("ORT session builder: {e}"),
578 })?
579 .commit_from_file(onnx)
580 .map_err(|e| ProviderError::ModelLoad {
581 model: model_label.to_string(),
582 reason: format!("load ONNX: {e}"),
583 })?;
584 let voices = load_voices_npz(voices_file)?;
585 let matrix_max = voices
587 .values()
588 .map(|m| m.nrows.saturating_sub(1))
589 .min()
590 .unwrap_or(0);
591 let max_phoneme_tokens = if matrix_max == 0 {
592 catalogue_max_tokens
593 } else {
594 matrix_max.min(catalogue_max_tokens)
595 };
596 Ok(LoadedPack {
597 session: Mutex::new(session),
598 voices,
599 sample_rate_hz,
600 max_phoneme_tokens,
601 speed_priors,
602 })
603}
604
605fn load_speed_priors(
606 cache_dir: &Path,
607 info: &super::catalogue::TtsModelInfo,
608) -> HashMap<String, f32> {
609 load_speed_priors_from_path(&super::catalogue::config_path(cache_dir, info))
610}
611
612fn load_speed_priors_from_path(path: &Path) -> HashMap<String, f32> {
613 let Ok(bytes) = std::fs::read(path) else {
614 return HashMap::new();
615 };
616 let Ok(v) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
617 return HashMap::new();
618 };
619 let mut out = HashMap::new();
620 if let Some(map) = v.get("speed_priors").and_then(|x| x.as_object()) {
621 for (k, val) in map {
622 if let Some(f) = val.as_f64() {
623 out.insert(k.clone(), f as f32);
624 }
625 }
626 }
627 out
628}
629
630fn synthesize_fake_adapter(
632 opts: &SynthesisOptions,
633 model_id: &str,
634 trust: TrustMode,
635 provenance: &str,
636 text_chars: usize,
637) -> Result<SynthesisResult> {
638 let duration_ms = ((text_chars as u64).saturating_mul(40)).clamp(50, 2_000);
640 let pcm_f32 = synthesize_fake_sine_ms(duration_ms).map_err(|e| ProviderError::Other {
641 message: format!("fake-sine synth: {e}"),
642 })?;
643 validate_raw_pcm(&pcm_f32, 24_000)?;
644 let pcm = peak_guard_f32_to_i16(&pcm_f32, PEAK_LIMIT);
645 let sample_rate = resolve_sample_rate(opts.sample_rate_hz, 24_000)?;
646 let language = normalize_tts_language(&opts.language).unwrap_or_else(|_| "en".into());
647 let voice = if opts.voice.trim().is_empty() {
648 "Tone".into()
649 } else {
650 opts.voice.clone()
651 };
652 let out_duration_ms = duration_ms_from_pcm(pcm.len(), sample_rate);
653 Ok(SynthesisResult {
654 pcm_i16_mono: pcm,
655 sample_rate_hz: sample_rate,
656 channels: 1,
657 backend_kind: BackendKind::Local,
658 provider: "local".into(),
659 model: model_id.into(),
660 voice,
661 language,
662 duration_ms: out_duration_ms,
663 text_chars,
664 text_truncated: false,
665 chunk_count: 1,
666 synthesized_chars: text_chars,
667 adapter: Some(ADAPTER_FAKE_SINE_V1.into()),
668 trust: Some(trust.as_str().into()),
669 provenance: Some(provenance.into()),
670 })
671}
672
673fn resolve_pack_voice(
675 manifest: &super::adapter::ModelPackManifest,
676 requested: &str,
677) -> Result<(String, String)> {
678 let req = requested.trim();
679 let default_voice = if manifest.adapter_id == ADAPTER_KOKORO_ONNX_V0 {
680 super::catalogue::KOKORO_DEFAULT_VOICE
681 } else {
682 super::catalogue::DEFAULT_TTS_VOICE
683 };
684 let req = if req.is_empty() { default_voice } else { req };
685 if let Ok((_, v)) = resolve_voice_for_model(&manifest.model_id, req) {
686 return Ok((v.id.to_string(), v.internal_key.to_string()));
687 }
688 if let Ok((_, v)) = resolve_voice_for_model(super::catalogue::DEFAULT_TTS_MODEL, req) {
689 return Ok((v.id.to_string(), v.internal_key.to_string()));
690 }
691 if let Some(v) = manifest
692 .voices
693 .iter()
694 .find(|v| v.id.eq_ignore_ascii_case(req) || v.internal_key.eq_ignore_ascii_case(req))
695 {
696 return Ok((v.id.clone(), v.internal_key.clone()));
697 }
698 let available: Vec<_> = manifest.voices.iter().map(|v| v.id.as_str()).collect();
699 Err(UserError::Other {
700 message: format!(
701 "voice '{req}' not found in pack '{}'; available: {available:?}",
702 manifest.model_id
703 ),
704 }
705 .into())
706}
707
708#[async_trait]
709impl SynthesisProvider for LocalTtsProvider {
710 fn name(&self) -> &'static str {
711 "local"
712 }
713
714 async fn synthesize(&self, text: &str, opts: &SynthesisOptions) -> Result<SynthesisResult> {
715 let prepared = prepare_text(text, self.max_chars)?;
716 let mut opts = opts.clone();
717 opts.language = normalize_tts_language(&opts.language)?;
718
719 if let Some(pack_dir) = opts.pack_dir.clone() {
722 let (root, manifest) = load_pack_dir(&pack_dir, opts.allow_unverified)?;
723 let _ = root;
724 if manifest.adapter_id == ADAPTER_FAKE_SINE_V1 {
725 let trust = manifest.trust;
726 let model_id = if opts.model.trim().is_empty()
729 || opts.model == super::catalogue::DEFAULT_TTS_MODEL
730 {
731 manifest.model_id.clone()
732 } else {
733 opts.model.clone()
734 };
735 if opts.voice.trim().is_empty() || opts.voice == super::catalogue::DEFAULT_TTS_VOICE
736 {
737 if let Some(v) = manifest.voices.first() {
738 opts.voice = v.id.clone();
739 } else {
740 opts.voice = "Tone".into();
741 }
742 }
743 return synthesize_fake_adapter(
744 &opts,
745 &model_id,
746 trust,
747 "local_pack",
748 prepared.text_chars,
749 );
750 }
751 if manifest.adapter_id != ADAPTER_KITTEN_ONNX_V1
752 && manifest.adapter_id != ADAPTER_KOKORO_ONNX_V0
753 {
754 return Err(UserError::UnsupportedCapability {
755 provider: "tts".into(),
756 model: manifest.model_id,
757 reason: format!(
758 "adapter '{}' is not enabled for local pack synthesis",
759 manifest.adapter_id
760 ),
761 hint:
762 "use kitten-onnx-v1, kokoro-onnx-v0, or fake-sine-v1 packs; see `aurum tts adapters`"
763 .into(),
764 }
765 .into());
766 }
767 let model_canonical = if opts.model == super::catalogue::DEFAULT_TTS_MODEL
769 || opts.model.trim().is_empty()
770 {
771 manifest.model_id.clone()
772 } else {
773 opts.model.clone()
774 };
775 let (voice_canonical, voice_internal) = resolve_pack_voice(&manifest, &opts.voice)?;
776 validate_speaking_rate(opts.speaking_rate)?;
777 resolve_sample_rate(opts.sample_rate_hz, manifest.sample_rate_hz)?;
778 let trust = manifest.trust;
779 let adapter = manifest.adapter_id.clone();
780 let lease = self
782 .ensure_loaded_from_pack_pin(&pack_dir, opts.allow_unverified)
783 .await?
784 .0;
785
786 let timeout = Duration::from_millis(if opts.timeout_ms == 0 {
787 super::validate::DEFAULT_TIMEOUT_MS
788 } else {
789 opts.timeout_ms
790 });
791 let op = OpContext::from_optional_cancel(opts.cancel.clone())
795 .with_deadline_from_now(timeout);
796 op.check()?;
797 op.emit("tts", "pack_synth");
798 let text_owned = prepared.text.clone();
799 let text_chars = prepared.text_chars;
800 let opts_owned = opts.clone();
801 let op_for_worker = op.clone();
802 let gov = Arc::clone(&self.governor);
803 let join = tokio::task::spawn_blocking(move || {
804 let _lease = lease;
807 let _permit = gov.acquire_tts(0, Some(&op_for_worker))?;
808 op_for_worker.check()?;
809 synthesize_with_pack(
810 _lease.value().as_ref(),
811 &text_owned,
812 &opts_owned,
813 &voice_internal,
814 &voice_canonical,
815 &model_canonical,
816 text_chars,
817 &op_for_worker,
818 &adapter,
819 trust,
820 "local_pack",
821 )
822 });
823
824 return tokio::select! {
825 join_res = join => {
826 match join_res {
827 Ok(result) => result,
828 Err(e) => Err(crate::error::TranscriptionError::internal(format!(
829 "TTS pack synth join: {e}"
830 ))),
831 }
832 }
833 _ = tokio::time::sleep(timeout) => {
834 op.cancel.cancel();
838 Err(ProviderError::DeadlineExceeded.into())
839 }
840 };
841 }
842
843 let (model_info, voice_info) = resolve_voice_for_model(&opts.model, &opts.voice)?;
844 opts.model = model_info.id.to_string();
846 opts.voice = voice_info.id.to_string();
847 validate_speaking_rate(opts.speaking_rate)?;
848 resolve_sample_rate(opts.sample_rate_hz, model_info.sample_rate_hz)?;
849
850 let local_only = opts.local_only || self.local_only;
851 let lease = self.ensure_loaded_pin(&opts.model, local_only).await?;
852
853 let timeout = Duration::from_millis(if opts.timeout_ms == 0 {
854 super::validate::DEFAULT_TIMEOUT_MS
855 } else {
856 opts.timeout_ms
857 });
858 let op =
859 OpContext::from_optional_cancel(opts.cancel.clone()).with_deadline_from_now(timeout);
860 op.check()?;
861 op.emit("tts", "builtin_synth");
862
863 let text_owned = prepared.text.clone();
864 let text_chars = prepared.text_chars;
865 let opts_owned = opts.clone();
866 let voice_internal = voice_info.internal_key.to_string();
867 let voice_canonical = voice_info.id.to_string();
868 let model_canonical = model_info.id.to_string();
869 let op_for_worker = op.clone();
870 let adapter = model_info.adapter.to_string();
871 let gov = Arc::clone(&self.governor);
872
873 let join = tokio::task::spawn_blocking(move || {
877 let _lease = lease;
878 let _permit = gov.acquire_tts(0, Some(&op_for_worker))?;
879 op_for_worker.check()?;
880 synthesize_with_pack(
881 _lease.value().as_ref(),
882 &text_owned,
883 &opts_owned,
884 &voice_internal,
885 &voice_canonical,
886 &model_canonical,
887 text_chars,
888 &op_for_worker,
889 &adapter,
890 TrustMode::Builtin,
891 "builtin",
892 )
893 });
894
895 tokio::select! {
899 join_res = join => {
900 match join_res {
901 Ok(result) => result,
902 Err(e) => Err(crate::error::TranscriptionError::internal(format!(
903 "TTS synth join: {e}"
904 ))),
905 }
906 }
907 _ = tokio::time::sleep(timeout) => {
908 op.cancel.cancel();
909 Err(ProviderError::DeadlineExceeded.into())
910 }
911 }
912 }
913
914 async fn preload(&self, model: &str, voice: &str) -> Result<()> {
915 let (model_info, _) = resolve_voice_for_model(model, voice)?;
917 let pin = self
918 .ensure_loaded_pin(model_info.id, self.local_only)
919 .await?;
920 drop(pin);
921 Ok(())
922 }
923}
924
925pub async fn synthesize_local(
927 cache_dir: impl Into<PathBuf>,
928 text: &str,
929 opts: &SynthesisOptions,
930) -> Result<SynthesisResult> {
931 let provider = LocalTtsProvider::new(cache_dir.into()).with_progress(false);
932 provider.synthesize(text, opts).await
933}
934
935#[cfg(test)]
936mod tests {
937 use super::*;
938
939 #[tokio::test]
940 async fn empty_text_user_error() {
941 let dir = tempfile::tempdir().unwrap();
942 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
943 let err = p
944 .synthesize(" ", &SynthesisOptions::default())
945 .await
946 .unwrap_err();
947 assert_eq!(err.exit_code(), 2);
948 }
949
950 #[tokio::test]
951 async fn missing_pack_local_only() {
952 let dir = tempfile::tempdir().unwrap();
953 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
954 let err = p
955 .synthesize("Hello", &SynthesisOptions::default())
956 .await
957 .unwrap_err();
958 assert!(matches!(err.exit_code(), 2 | 4));
959 }
960
961 #[tokio::test]
962 async fn unsupported_language_user_error() {
963 let dir = tempfile::tempdir().unwrap();
964 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
965 let opts = SynthesisOptions {
966 language: "fr".into(),
967 ..Default::default()
968 };
969 let err = p.synthesize("Hello", &opts).await.unwrap_err();
970 assert_eq!(err.exit_code(), 2);
971 assert!(err.to_string().contains("unsupported TTS language"));
972 }
973
974 #[tokio::test]
975 async fn non_native_sample_rate_rejected() {
976 let dir = tempfile::tempdir().unwrap();
977 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
978 let opts = SynthesisOptions {
979 sample_rate_hz: Some(16_000),
980 ..Default::default()
981 };
982 let err = p.synthesize("Hello", &opts).await.unwrap_err();
983 assert_eq!(err.exit_code(), 2);
984 assert!(err.to_string().contains("sample rate"));
985 }
986
987 #[tokio::test]
988 async fn invalid_speaking_rate_rejected() {
989 let dir = tempfile::tempdir().unwrap();
990 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
991 let opts = SynthesisOptions {
992 speaking_rate: 9.0,
993 ..Default::default()
994 };
995 let err = p.synthesize("Hello", &opts).await.unwrap_err();
996 assert_eq!(err.exit_code(), 2);
997 }
998
999 #[tokio::test]
1000 async fn oversized_text_rejected_not_truncated() {
1001 let dir = tempfile::tempdir().unwrap();
1002 let p = LocalTtsProvider::new(dir.path().to_path_buf())
1003 .with_local_only(true)
1004 .with_max_chars(10);
1005 let err = p
1006 .synthesize(
1007 "this text is definitely longer than ten",
1008 &SynthesisOptions::default(),
1009 )
1010 .await
1011 .unwrap_err();
1012 assert_eq!(err.exit_code(), 2);
1013 assert!(err.to_string().contains("too long"));
1014 }
1015
1016 #[test]
1017 fn clear_sessions_is_safe_when_empty() {
1018 let dir = tempfile::tempdir().unwrap();
1019 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
1020 p.clear_sessions();
1021 p.clear_sessions();
1022 }
1023
1024 #[tokio::test]
1031 #[ignore]
1032 async fn kokoro_real_synth_from_cache() {
1033 if std::env::var("AURUM_KOKORO_INTEGRATION").ok().as_deref() != Some("1") {
1034 return;
1035 }
1036 let cache = std::env::var("AURUM_TTS_CACHE")
1037 .map(std::path::PathBuf::from)
1038 .unwrap_or_else(|_| {
1039 PathBuf::from(std::env::var("HOME").unwrap_or_else(|_| "/tmp".into()))
1040 .join(".cache/aurum")
1041 });
1042 let p = LocalTtsProvider::new(cache).with_local_only(true);
1043 let opts = SynthesisOptions {
1044 model: crate::tts::catalogue::KOKORO_TTS_MODEL.into(),
1045 voice: crate::tts::catalogue::KOKORO_DEFAULT_VOICE.into(),
1046 ..Default::default()
1047 };
1048 let r = p
1049 .synthesize("Hello from Kokoro.", &opts)
1050 .await
1051 .expect("kokoro synth");
1052 assert_eq!(r.sample_rate_hz, 24_000);
1053 assert!(!r.pcm_i16_mono.is_empty());
1054 assert_eq!(r.adapter.as_deref(), Some(ADAPTER_KOKORO_ONNX_V0));
1055 assert!(r.duration_ms > 0);
1056 }
1057}