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