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, reverify_artifact_before_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
274 let pin = tokio::task::spawn_blocking(move || {
275 pool.get_or_load_pin(key_for_load, weight, &gov, || {
276 load_pack(
277 &onnx,
278 &voices_file,
279 sample_rate,
280 speed_priors,
281 catalogue_max,
282 &model_id,
283 )
284 })
285 })
286 .await
287 .map_err(|e| crate::error::TranscriptionError::internal(format!("TTS load join: {e}")))??;
288
289 Ok(pin)
290 }
291
292 async fn ensure_loaded_from_pack_pin(
295 &self,
296 pack_dir: &Path,
297 allow_unverified: bool,
298 ) -> Result<(RegistryPin<LoadedPack>, super::adapter::ModelPackManifest)> {
299 let (root, manifest) = load_pack_dir(pack_dir, allow_unverified)?;
300 if manifest.adapter_id != ADAPTER_KITTEN_ONNX_V1
301 && manifest.adapter_id != ADAPTER_KOKORO_ONNX_V0
302 {
303 return Err(UserError::InvalidConfig {
304 reason: format!(
305 "local pack override synthesis currently supports adapters \
306 '{ADAPTER_KITTEN_ONNX_V1}' and '{ADAPTER_KOKORO_ONNX_V0}' \
307 (got '{}'); use `aurum tts inspect` / conformance for others",
308 manifest.adapter_id
309 ),
310 }
311 .into());
312 }
313 let onnx_name = manifest
314 .artifact("onnx")
315 .map(|a| a.filename.as_str())
316 .ok_or_else(|| UserError::InvalidConfig {
317 reason: "pack missing onnx artifact".into(),
318 })?;
319 let voices_name = manifest
320 .artifact("voices")
321 .map(|a| a.filename.as_str())
322 .ok_or_else(|| UserError::InvalidConfig {
323 reason: "pack missing voices artifact".into(),
324 })?;
325 let onnx = root.join(onnx_name);
326 let voices_file = root.join(voices_name);
327 let onnx_sha = manifest.artifact("onnx").and_then(|a| a.sha256.clone());
328 let voices_sha = manifest.artifact("voices").and_then(|a| a.sha256.clone());
329 let sample_rate = manifest.sample_rate_hz;
330 let catalogue_max = manifest.max_phoneme_tokens;
331 let model_id = manifest.model_id.clone();
332 let key = LoadKey::tts(
334 format!("local-pack:{}", manifest.model_id),
335 root.display().to_string(),
336 );
337 if let Some(pin) = self.pool.registry.get_and_pin(&key) {
338 return Ok((pin, manifest));
339 }
340 let key_for_load = key.clone();
341 let speed_priors = load_speed_priors_from_path(
342 &root.join(
343 manifest
344 .artifact("config")
345 .map(|a| a.filename.as_str())
346 .unwrap_or("config.json"),
347 ),
348 );
349 let weight = TtsSessionPool::weight_for(&onnx);
350 let pool = Arc::clone(&self.pool);
351 let gov = Arc::clone(&self.governor);
352
353 let pin = tokio::task::spawn_blocking(move || {
354 pool.get_or_load_pin(key_for_load, weight, &gov, || {
355 reverify_artifact_before_load(&onnx, onnx_sha.as_deref())?;
358 reverify_artifact_before_load(&voices_file, voices_sha.as_deref())?;
359 load_pack(
360 &onnx,
361 &voices_file,
362 sample_rate,
363 speed_priors,
364 catalogue_max,
365 &model_id,
366 )
367 })
368 })
369 .await
370 .map_err(|e| {
371 crate::error::TranscriptionError::internal(format!("TTS pack load join: {e}"))
372 })??;
373
374 Ok((pin, manifest))
375 }
376}
377
378#[allow(clippy::too_many_arguments)]
379fn synthesize_with_pack(
380 pack: &LoadedPack,
381 text: &str,
382 opts: &SynthesisOptions,
383 voice_internal: &str,
384 voice_canonical: &str,
385 model_canonical: &str,
386 text_chars: usize,
387 op: &OpContext,
388 adapter: &str,
389 trust: TrustMode,
390 provenance: &str,
391) -> Result<SynthesisResult> {
392 op.check()?;
393
394 let rate = validate_speaking_rate(opts.speaking_rate)?;
395 let sample_rate = resolve_sample_rate(opts.sample_rate_hz, pack.sample_rate_hz)?;
396
397 let voice_mat = pack
398 .voices
399 .get(voice_internal)
400 .ok_or_else(|| UserError::Other {
401 message: format!(
402 "voice embedding '{voice_internal}' missing from pack; available: {:?}",
403 pack.voices.keys().collect::<Vec<_>>()
404 ),
405 })?;
406
407 let pack_max = voice_mat.nrows.saturating_sub(1);
410 let max_tokens = pack_max.min(pack.max_phoneme_tokens);
411 if max_tokens <= 2 {
412 return Err(ProviderError::ModelLoad {
413 model: model_canonical.to_string(),
414 reason: "voice embedding has no usable sequence rows".into(),
415 }
416 .into());
417 }
418
419 let codec = phoneme_codec_for_adapter(adapter);
420 let chunks = prepare_tts_chunks_with(text, max_tokens, codec)?;
421 let chunk_count = chunks.len();
422 let effective_speed = rate
423 * pack
424 .speed_priors
425 .get(voice_internal)
426 .copied()
427 .unwrap_or(1.0);
428 let pause_samples = (sample_rate as u64)
429 .saturating_mul(CHUNK_PAUSE_MS)
430 .checked_div(1_000)
431 .unwrap_or(0) as usize;
432
433 let mut pcm_f32: Vec<f32> = Vec::new();
434 let trim_policy = TailTrimPolicy::default();
435
436 for (index, chunk) in chunks.iter().enumerate() {
437 op.check()?;
438 let chunk_audio =
439 synthesize_chunk_f32(pack, voice_mat, chunk, effective_speed).map_err(|err| {
440 ProviderError::Other {
441 message: format!(
442 "TTS chunk {}/{} failed near {:?}: {err}",
443 index + 1,
444 chunks.len(),
445 chunk.text.chars().take(80).collect::<String>()
446 ),
447 }
448 })?;
449 let trimmed = trim_trailing_silence(&chunk_audio, trim_policy);
450 if index > 0 {
451 pcm_f32.resize(pcm_f32.len().saturating_add(pause_samples), 0.0);
452 }
453 pcm_f32.extend_from_slice(trimmed);
454 }
455
456 validate_raw_pcm(&pcm_f32, sample_rate)?;
457 let pcm = peak_guard_f32_to_i16(&pcm_f32, PEAK_LIMIT);
458 if pcm.is_empty() {
459 return Err(ProviderError::Other {
460 message: "synthesis produced empty audio after validation".into(),
461 }
462 .into());
463 }
464
465 let duration_ms = duration_ms_from_pcm(pcm.len(), sample_rate);
466 let language = normalize_tts_language(&opts.language).unwrap_or_else(|_| "en".into());
467
468 Ok(SynthesisResult {
469 pcm_i16_mono: pcm,
470 sample_rate_hz: sample_rate,
471 channels: 1,
472 backend_kind: BackendKind::Local,
473 provider: "local".into(),
474 model: model_canonical.to_string(),
475 voice: voice_canonical.to_string(),
476 language,
477 duration_ms,
478 text_chars,
479 text_truncated: false,
480 chunk_count,
481 synthesized_chars: text_chars,
482 adapter: Some(adapter.into()),
483 trust: Some(trust.as_str().into()),
484 provenance: Some(provenance.into()),
485 })
486}
487
488fn phoneme_codec_for_adapter(adapter: &str) -> PhonemeCodec {
489 if adapter == ADAPTER_KOKORO_ONNX_V0 {
490 PhonemeCodec::Kokoro
491 } else {
492 PhonemeCodec::Kitten
493 }
494}
495
496fn synthesize_chunk_f32(
497 pack: &LoadedPack,
498 voice_mat: &super::npz::VoiceMatrix,
499 chunk: &TtsChunk,
500 effective_speed: f32,
501) -> Result<Vec<f32>> {
502 let seq_len = chunk.ids.len();
503 let style = voice_mat.style_row(seq_len).to_vec();
505 let style_dim = style.len();
506 let t_ids = Tensor::<i64>::from_array(([1usize, seq_len], chunk.ids.clone())).map_err(|e| {
507 ProviderError::Other {
508 message: format!("tokens/input_ids tensor: {e}"),
509 }
510 })?;
511 let t_style = Tensor::<f32>::from_array(([1usize, style_dim], style)).map_err(|e| {
512 ProviderError::Other {
513 message: format!("style tensor: {e}"),
514 }
515 })?;
516 let t_speed = Tensor::<f32>::from_array(([1usize], vec![effective_speed])).map_err(|e| {
517 ProviderError::Other {
518 message: format!("speed tensor: {e}"),
519 }
520 })?;
521
522 let mut session = pack
523 .session
524 .lock()
525 .map_err(|_| crate::error::TranscriptionError::internal("ORT session mutex poisoned"))?;
526 let outputs = session
528 .run(ort::inputs![t_ids, t_style, t_speed])
529 .map_err(|e| ProviderError::Other {
530 message: format!("ONNX inference failed: {e}"),
531 })?;
532 let (_shape, audio_data) =
533 outputs[0]
534 .try_extract_tensor::<f32>()
535 .map_err(|e| ProviderError::Other {
536 message: format!("extract audio tensor: {e}"),
537 })?;
538 Ok(audio_data.to_vec())
539}
540
541fn load_pack(
542 onnx: &Path,
543 voices_file: &Path,
544 sample_rate_hz: u32,
545 speed_priors: HashMap<String, f32>,
546 catalogue_max_tokens: usize,
547 model_label: &str,
548) -> Result<LoadedPack> {
549 let session = Session::builder()
550 .map_err(|e| ProviderError::ModelLoad {
551 model: model_label.to_string(),
552 reason: format!("ORT session builder: {e}"),
553 })?
554 .commit_from_file(onnx)
555 .map_err(|e| ProviderError::ModelLoad {
556 model: model_label.to_string(),
557 reason: format!("load ONNX: {e}"),
558 })?;
559 let voices = load_voices_npz(voices_file)?;
560 let matrix_max = voices
562 .values()
563 .map(|m| m.nrows.saturating_sub(1))
564 .min()
565 .unwrap_or(0);
566 let max_phoneme_tokens = if matrix_max == 0 {
567 catalogue_max_tokens
568 } else {
569 matrix_max.min(catalogue_max_tokens)
570 };
571 Ok(LoadedPack {
572 session: Mutex::new(session),
573 voices,
574 sample_rate_hz,
575 max_phoneme_tokens,
576 speed_priors,
577 })
578}
579
580fn load_speed_priors(
581 cache_dir: &Path,
582 info: &super::catalogue::TtsModelInfo,
583) -> HashMap<String, f32> {
584 load_speed_priors_from_path(&super::catalogue::config_path(cache_dir, info))
585}
586
587fn load_speed_priors_from_path(path: &Path) -> HashMap<String, f32> {
588 let Ok(bytes) = std::fs::read(path) else {
589 return HashMap::new();
590 };
591 let Ok(v) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
592 return HashMap::new();
593 };
594 let mut out = HashMap::new();
595 if let Some(map) = v.get("speed_priors").and_then(|x| x.as_object()) {
596 for (k, val) in map {
597 if let Some(f) = val.as_f64() {
598 out.insert(k.clone(), f as f32);
599 }
600 }
601 }
602 out
603}
604
605fn synthesize_fake_adapter(
607 opts: &SynthesisOptions,
608 model_id: &str,
609 trust: TrustMode,
610 provenance: &str,
611 text_chars: usize,
612) -> Result<SynthesisResult> {
613 let duration_ms = ((text_chars as u64).saturating_mul(40)).clamp(50, 2_000);
615 let pcm_f32 = synthesize_fake_sine_ms(duration_ms).map_err(|e| ProviderError::Other {
616 message: format!("fake-sine synth: {e}"),
617 })?;
618 validate_raw_pcm(&pcm_f32, 24_000)?;
619 let pcm = peak_guard_f32_to_i16(&pcm_f32, PEAK_LIMIT);
620 let sample_rate = resolve_sample_rate(opts.sample_rate_hz, 24_000)?;
621 let language = normalize_tts_language(&opts.language).unwrap_or_else(|_| "en".into());
622 let voice = if opts.voice.trim().is_empty() {
623 "Tone".into()
624 } else {
625 opts.voice.clone()
626 };
627 let out_duration_ms = duration_ms_from_pcm(pcm.len(), sample_rate);
628 Ok(SynthesisResult {
629 pcm_i16_mono: pcm,
630 sample_rate_hz: sample_rate,
631 channels: 1,
632 backend_kind: BackendKind::Local,
633 provider: "local".into(),
634 model: model_id.into(),
635 voice,
636 language,
637 duration_ms: out_duration_ms,
638 text_chars,
639 text_truncated: false,
640 chunk_count: 1,
641 synthesized_chars: text_chars,
642 adapter: Some(ADAPTER_FAKE_SINE_V1.into()),
643 trust: Some(trust.as_str().into()),
644 provenance: Some(provenance.into()),
645 })
646}
647
648fn resolve_pack_voice(
650 manifest: &super::adapter::ModelPackManifest,
651 requested: &str,
652) -> Result<(String, String)> {
653 let req = requested.trim();
654 let default_voice = if manifest.adapter_id == ADAPTER_KOKORO_ONNX_V0 {
655 super::catalogue::KOKORO_DEFAULT_VOICE
656 } else {
657 super::catalogue::DEFAULT_TTS_VOICE
658 };
659 let req = if req.is_empty() { default_voice } else { req };
660 if let Ok((_, v)) = resolve_voice_for_model(&manifest.model_id, req) {
661 return Ok((v.id.to_string(), v.internal_key.to_string()));
662 }
663 if let Ok((_, v)) = resolve_voice_for_model(super::catalogue::DEFAULT_TTS_MODEL, req) {
664 return Ok((v.id.to_string(), v.internal_key.to_string()));
665 }
666 if let Some(v) = manifest
667 .voices
668 .iter()
669 .find(|v| v.id.eq_ignore_ascii_case(req) || v.internal_key.eq_ignore_ascii_case(req))
670 {
671 return Ok((v.id.clone(), v.internal_key.clone()));
672 }
673 let available: Vec<_> = manifest.voices.iter().map(|v| v.id.as_str()).collect();
674 Err(UserError::Other {
675 message: format!(
676 "voice '{req}' not found in pack '{}'; available: {available:?}",
677 manifest.model_id
678 ),
679 }
680 .into())
681}
682
683#[async_trait]
684impl SynthesisProvider for LocalTtsProvider {
685 fn name(&self) -> &'static str {
686 "local"
687 }
688
689 async fn synthesize(&self, text: &str, opts: &SynthesisOptions) -> Result<SynthesisResult> {
690 let prepared = prepare_text(text, self.max_chars)?;
691 let mut opts = opts.clone();
692 opts.language = normalize_tts_language(&opts.language)?;
693
694 if let Some(pack_dir) = opts.pack_dir.clone() {
697 let (root, manifest) = load_pack_dir(&pack_dir, opts.allow_unverified)?;
698 let _ = root;
699 if manifest.adapter_id == ADAPTER_FAKE_SINE_V1 {
700 let trust = manifest.trust;
701 let model_id = if opts.model.trim().is_empty()
704 || opts.model == super::catalogue::DEFAULT_TTS_MODEL
705 {
706 manifest.model_id.clone()
707 } else {
708 opts.model.clone()
709 };
710 if opts.voice.trim().is_empty() || opts.voice == super::catalogue::DEFAULT_TTS_VOICE
711 {
712 if let Some(v) = manifest.voices.first() {
713 opts.voice = v.id.clone();
714 } else {
715 opts.voice = "Tone".into();
716 }
717 }
718 return synthesize_fake_adapter(
719 &opts,
720 &model_id,
721 trust,
722 "local_pack",
723 prepared.text_chars,
724 );
725 }
726 if manifest.adapter_id != ADAPTER_KITTEN_ONNX_V1
727 && manifest.adapter_id != ADAPTER_KOKORO_ONNX_V0
728 {
729 return Err(UserError::UnsupportedCapability {
730 provider: "tts".into(),
731 model: manifest.model_id,
732 reason: format!(
733 "adapter '{}' is not enabled for local pack synthesis",
734 manifest.adapter_id
735 ),
736 hint:
737 "use kitten-onnx-v1, kokoro-onnx-v0, or fake-sine-v1 packs; see `aurum tts adapters`"
738 .into(),
739 }
740 .into());
741 }
742 let model_canonical = if opts.model == super::catalogue::DEFAULT_TTS_MODEL
744 || opts.model.trim().is_empty()
745 {
746 manifest.model_id.clone()
747 } else {
748 opts.model.clone()
749 };
750 let (voice_canonical, voice_internal) = resolve_pack_voice(&manifest, &opts.voice)?;
751 validate_speaking_rate(opts.speaking_rate)?;
752 resolve_sample_rate(opts.sample_rate_hz, manifest.sample_rate_hz)?;
753 let trust = manifest.trust;
754 let adapter = manifest.adapter_id.clone();
755 let lease = self
757 .ensure_loaded_from_pack_pin(&pack_dir, opts.allow_unverified)
758 .await?
759 .0;
760
761 let timeout = Duration::from_millis(if opts.timeout_ms == 0 {
762 super::validate::DEFAULT_TIMEOUT_MS
763 } else {
764 opts.timeout_ms
765 });
766 let op = OpContext::from_optional_cancel(opts.cancel.clone())
770 .with_deadline_from_now(timeout);
771 op.check()?;
772 op.emit("tts", "pack_synth");
773 let text_owned = prepared.text.clone();
774 let text_chars = prepared.text_chars;
775 let opts_owned = opts.clone();
776 let op_for_worker = op.clone();
777 let gov = Arc::clone(&self.governor);
778 let join = tokio::task::spawn_blocking(move || {
779 let _lease = lease;
782 let _permit = gov.acquire_tts(0, Some(&op_for_worker))?;
783 op_for_worker.check()?;
784 synthesize_with_pack(
785 _lease.value().as_ref(),
786 &text_owned,
787 &opts_owned,
788 &voice_internal,
789 &voice_canonical,
790 &model_canonical,
791 text_chars,
792 &op_for_worker,
793 &adapter,
794 trust,
795 "local_pack",
796 )
797 });
798
799 return tokio::select! {
800 join_res = join => {
801 match join_res {
802 Ok(result) => result,
803 Err(e) => Err(crate::error::TranscriptionError::internal(format!(
804 "TTS pack synth join: {e}"
805 ))),
806 }
807 }
808 _ = tokio::time::sleep(timeout) => {
809 op.cancel.cancel();
813 Err(ProviderError::DeadlineExceeded.into())
814 }
815 };
816 }
817
818 let (model_info, voice_info) = resolve_voice_for_model(&opts.model, &opts.voice)?;
819 opts.model = model_info.id.to_string();
821 opts.voice = voice_info.id.to_string();
822 validate_speaking_rate(opts.speaking_rate)?;
823 resolve_sample_rate(opts.sample_rate_hz, model_info.sample_rate_hz)?;
824
825 let local_only = opts.local_only || self.local_only;
826 let lease = self.ensure_loaded_pin(&opts.model, local_only).await?;
827
828 let timeout = Duration::from_millis(if opts.timeout_ms == 0 {
829 super::validate::DEFAULT_TIMEOUT_MS
830 } else {
831 opts.timeout_ms
832 });
833 let op =
834 OpContext::from_optional_cancel(opts.cancel.clone()).with_deadline_from_now(timeout);
835 op.check()?;
836 op.emit("tts", "builtin_synth");
837
838 let text_owned = prepared.text.clone();
839 let text_chars = prepared.text_chars;
840 let opts_owned = opts.clone();
841 let voice_internal = voice_info.internal_key.to_string();
842 let voice_canonical = voice_info.id.to_string();
843 let model_canonical = model_info.id.to_string();
844 let op_for_worker = op.clone();
845 let adapter = model_info.adapter.to_string();
846 let gov = Arc::clone(&self.governor);
847
848 let join = tokio::task::spawn_blocking(move || {
852 let _lease = lease;
853 let _permit = gov.acquire_tts(0, Some(&op_for_worker))?;
854 op_for_worker.check()?;
855 synthesize_with_pack(
856 _lease.value().as_ref(),
857 &text_owned,
858 &opts_owned,
859 &voice_internal,
860 &voice_canonical,
861 &model_canonical,
862 text_chars,
863 &op_for_worker,
864 &adapter,
865 TrustMode::Builtin,
866 "builtin",
867 )
868 });
869
870 tokio::select! {
874 join_res = join => {
875 match join_res {
876 Ok(result) => result,
877 Err(e) => Err(crate::error::TranscriptionError::internal(format!(
878 "TTS synth join: {e}"
879 ))),
880 }
881 }
882 _ = tokio::time::sleep(timeout) => {
883 op.cancel.cancel();
884 Err(ProviderError::DeadlineExceeded.into())
885 }
886 }
887 }
888
889 async fn preload(&self, model: &str, voice: &str) -> Result<()> {
890 let (model_info, _) = resolve_voice_for_model(model, voice)?;
892 let pin = self
893 .ensure_loaded_pin(model_info.id, self.local_only)
894 .await?;
895 drop(pin);
896 Ok(())
897 }
898}
899
900pub async fn synthesize_local(
902 cache_dir: impl Into<PathBuf>,
903 text: &str,
904 opts: &SynthesisOptions,
905) -> Result<SynthesisResult> {
906 let provider = LocalTtsProvider::new(cache_dir.into()).with_progress(false);
907 provider.synthesize(text, opts).await
908}
909
910#[cfg(test)]
911mod tests {
912 use super::*;
913
914 #[tokio::test]
915 async fn empty_text_user_error() {
916 let dir = tempfile::tempdir().unwrap();
917 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
918 let err = p
919 .synthesize(" ", &SynthesisOptions::default())
920 .await
921 .unwrap_err();
922 assert_eq!(err.exit_code(), 2);
923 }
924
925 #[tokio::test]
926 async fn missing_pack_local_only() {
927 let dir = tempfile::tempdir().unwrap();
928 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
929 let err = p
930 .synthesize("Hello", &SynthesisOptions::default())
931 .await
932 .unwrap_err();
933 assert!(matches!(err.exit_code(), 2 | 4));
934 }
935
936 #[tokio::test]
937 async fn unsupported_language_user_error() {
938 let dir = tempfile::tempdir().unwrap();
939 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
940 let opts = SynthesisOptions {
941 language: "fr".into(),
942 ..Default::default()
943 };
944 let err = p.synthesize("Hello", &opts).await.unwrap_err();
945 assert_eq!(err.exit_code(), 2);
946 assert!(err.to_string().contains("unsupported TTS language"));
947 }
948
949 #[tokio::test]
950 async fn non_native_sample_rate_rejected() {
951 let dir = tempfile::tempdir().unwrap();
952 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
953 let opts = SynthesisOptions {
954 sample_rate_hz: Some(16_000),
955 ..Default::default()
956 };
957 let err = p.synthesize("Hello", &opts).await.unwrap_err();
958 assert_eq!(err.exit_code(), 2);
959 assert!(err.to_string().contains("sample rate"));
960 }
961
962 #[tokio::test]
963 async fn invalid_speaking_rate_rejected() {
964 let dir = tempfile::tempdir().unwrap();
965 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
966 let opts = SynthesisOptions {
967 speaking_rate: 9.0,
968 ..Default::default()
969 };
970 let err = p.synthesize("Hello", &opts).await.unwrap_err();
971 assert_eq!(err.exit_code(), 2);
972 }
973
974 #[tokio::test]
975 async fn oversized_text_rejected_not_truncated() {
976 let dir = tempfile::tempdir().unwrap();
977 let p = LocalTtsProvider::new(dir.path().to_path_buf())
978 .with_local_only(true)
979 .with_max_chars(10);
980 let err = p
981 .synthesize(
982 "this text is definitely longer than ten",
983 &SynthesisOptions::default(),
984 )
985 .await
986 .unwrap_err();
987 assert_eq!(err.exit_code(), 2);
988 assert!(err.to_string().contains("too long"));
989 }
990
991 #[test]
992 fn clear_sessions_is_safe_when_empty() {
993 let dir = tempfile::tempdir().unwrap();
994 let p = LocalTtsProvider::new(dir.path().to_path_buf()).with_local_only(true);
995 p.clear_sessions();
996 p.clear_sessions();
997 }
998
999 #[tokio::test]
1006 #[ignore]
1007 async fn kokoro_real_synth_from_cache() {
1008 if std::env::var("AURUM_KOKORO_INTEGRATION").ok().as_deref() != Some("1") {
1009 return;
1010 }
1011 let cache = std::env::var("AURUM_TTS_CACHE")
1012 .map(std::path::PathBuf::from)
1013 .unwrap_or_else(|_| {
1014 PathBuf::from(std::env::var("HOME").unwrap_or_else(|_| "/tmp".into()))
1015 .join(".cache/aurum")
1016 });
1017 let p = LocalTtsProvider::new(cache).with_local_only(true);
1018 let opts = SynthesisOptions {
1019 model: crate::tts::catalogue::KOKORO_TTS_MODEL.into(),
1020 voice: crate::tts::catalogue::KOKORO_DEFAULT_VOICE.into(),
1021 ..Default::default()
1022 };
1023 let r = p
1024 .synthesize("Hello from Kokoro.", &opts)
1025 .await
1026 .expect("kokoro synth");
1027 assert_eq!(r.sample_rate_hz, 24_000);
1028 assert!(!r.pcm_i16_mono.is_empty());
1029 assert_eq!(r.adapter.as_deref(), Some(ADAPTER_KOKORO_ONNX_V0));
1030 assert!(r.duration_ms > 0);
1031 }
1032}