DGXAI/driftcall
0
1"""Cell 09 — Audio pipeline (Kokoro-82M TTS + faster-whisper-small ASR).2 3Implements docs/modules/audio.md: TTS and ASR engines exposed at the env4boundary. Training never imports this module (docs/modules/audio.md §6.3).5Heavy deps (``kokoro``, ``faster_whisper``, ``torchaudio``, ``soundfile``)6are loaded lazily inside ``_load_*`` helpers so this cell imports cleanly7in environments where those optional packages are absent, and so tests can8monkeypatch the loaders to return fakes without ever touching the network.9"""10 11from __future__ import annotations12 13import hashlib14import io15import logging16import math17import struct18import threading19import time20import unicodedata21import wave22from collections.abc import Callable23from dataclasses import dataclass24from datetime import datetime, timedelta, timezone25from typing import Any, Literal, cast26 27import numpy as np28from cachetools import LRUCache29 30logger = logging.getLogger(__name__)31 32 33# ---------------------------------------------------------------------------34# Public literal types (audio.md §2.1, §2.2)35# ---------------------------------------------------------------------------36 37LanguageCode = Literal["hi", "ta", "kn", "en", "hinglish"]38VoicePack = Literal[39 "hi_female_1",40 "hi_male_1",41 "ta_female_1",42 "kn_male_1",43 "en_indian_female_1",44]45 46_LANGUAGE_CODES: frozenset[str] = frozenset({"hi", "ta", "kn", "en", "hinglish"})47_VOICE_PACKS_SET: frozenset[str] = frozenset(48 {49 "hi_female_1",50 "hi_male_1",51 "ta_female_1",52 "kn_male_1",53 "en_indian_female_1",54 }55)56 57 58# ---------------------------------------------------------------------------59# Errors (audio.md §2.3)60# ---------------------------------------------------------------------------61 62 63class AudioError(Exception):64 """Base class for all audio-module errors."""65 66 67class ModelLoadError(AudioError):68 """Raised when Kokoro or faster-whisper cannot be instantiated."""69 70 71class UnsupportedLanguageError(AudioError):72 """Raised when a non-registered language code is passed to synthesize()."""73 74 75class UnsupportedVoicePackError(AudioError):76 """Raised when a voice pack is not in VOICE_PACKS[lang].allowed."""77 78 79class AudioDecodeError(AudioError):80 """Raised when transcribe() cannot decode the input bytes."""81 82 83class AudioTooLongError(AudioError):84 """Raised when transcribe() receives audio longer than max_duration_s in strict mode."""85 86 87class TTSOutOfMemoryError(AudioError):88 """Raised when TTS synthesis exhausts memory mid-call."""89 90 91# ---------------------------------------------------------------------------92# Data records (audio.md §2.1, §2.2, §2.2a, §4.1, §4.2)93# ---------------------------------------------------------------------------94 95 96@dataclass(frozen=True)97class VoicePackMapping:98 """Per-language default + allowed voice packs. audio.md §4.3."""99 100 language: LanguageCode101 default: VoicePack102 allowed: tuple[VoicePack, ...]103 104 105VOICE_PACKS: dict[LanguageCode, VoicePackMapping] = {106 "hi": VoicePackMapping(107 language="hi",108 default="hi_female_1",109 allowed=("hi_female_1", "hi_male_1"),110 ),111 "ta": VoicePackMapping(112 language="ta",113 default="ta_female_1",114 allowed=("ta_female_1",),115 ),116 "kn": VoicePackMapping(117 language="kn",118 default="kn_male_1",119 allowed=("kn_male_1",),120 ),121 "en": VoicePackMapping(122 language="en",123 default="en_indian_female_1",124 allowed=("en_indian_female_1",),125 ),126 "hinglish": VoicePackMapping(127 language="hinglish",128 default="en_indian_female_1",129 allowed=("en_indian_female_1", "hi_female_1"),130 ),131}132 133 134@dataclass(frozen=True)135class TranscriptResult:136 """ASR output surfaced to the env observation builder. audio.md §4.1."""137 138 text: str139 language_detected: LanguageCode | Literal["unknown"]140 confidence: float141 duration_s: float142 143 144@dataclass(frozen=True)145class AudioTrace:146 """Per-call diagnostic record emitted via the configured trace sink.147 148 audio.md §2.2a, §3.8.149 """150 151 op: Literal["synthesize", "transcribe"]152 input_hash: str153 language: str154 duration_s: float155 latency_ms: int156 confidence: float | None157 cache_hit: bool158 degraded: bool159 ts_ist: str160 161 162TraceSink = Callable[[AudioTrace], None]163 164 165# ---------------------------------------------------------------------------166# Lazy dep loaders — patched by tests to inject fakes.167# ---------------------------------------------------------------------------168 169 170def _load_kokoro() -> Any:171 """Return the ``kokoro`` module. Patched in tests."""172 173 import kokoro174 175 return kokoro176 177 178def _load_faster_whisper() -> Any:179 """Return the ``faster_whisper`` module. Patched in tests."""180 181 import faster_whisper182 183 return faster_whisper184 185 186def _load_torchaudio_functional() -> Any:187 """Return ``torchaudio.functional``. Patched in tests."""188 189 import torchaudio.functional as F190 191 return F192 193 194def _load_torchaudio() -> Any:195 """Return the top-level ``torchaudio`` module. Patched in tests."""196 197 import torchaudio198 199 return torchaudio200 201 202def _load_soundfile() -> Any:203 """Return the ``soundfile`` module. Patched in tests."""204 205 import soundfile206 207 return soundfile208 209 210def _load_torch() -> Any:211 """Return the ``torch`` module. Patched in tests."""212 213 import torch214 215 return torch216 217 218# ---------------------------------------------------------------------------219# Helpers220# ---------------------------------------------------------------------------221 222 223_IST_TZ = timezone(timedelta(hours=5, minutes=30))224 225 226def _ts_ist_now() -> str:227 return datetime.now(tz=_IST_TZ).isoformat(timespec="milliseconds")228 229 230def _input_hash(payload: bytes) -> str:231 return hashlib.blake2b(payload, digest_size=16).hexdigest()232 233 234def _logprob_to_confidence(avg_logprob: float) -> float:235 """Map faster-whisper ``avg_logprob`` into [0, 1] per audio.md §3.5."""236 237 clamped = max(-1.5, min(0.0, float(avg_logprob)))238 return round(math.exp(clamped), 3)239 240 241def _riff_header_sample_rate(audio_bytes: bytes) -> int | None:242 """Return the sample-rate field from a RIFF header, or None if not RIFF."""243 244 if len(audio_bytes) < 28:245 return None246 if audio_bytes[0:4] != b"RIFF" or audio_bytes[8:12] != b"WAVE":247 return None248 return int(struct.unpack_from("<I", audio_bytes, 24)[0])249 250 251def _pcm16_silence_wav(duration_s: float, sample_rate_hz: int = 16000) -> bytes:252 """Build a 16-bit mono PCM WAV of pure silence for warmup / fallback."""253 254 n_samples = max(1, int(duration_s * sample_rate_hz))255 buf = io.BytesIO()256 with wave.open(buf, "wb") as w:257 w.setnchannels(1)258 w.setsampwidth(2)259 w.setframerate(sample_rate_hz)260 w.writeframes(b"\x00\x00" * n_samples)261 return buf.getvalue()262 263 264def _np_to_wav_bytes(pcm: np.ndarray, sample_rate_hz: int) -> bytes:265 """Encode a float32 mono numpy array as 16-bit PCM RIFF WAV bytes.266 267 Used when torchaudio is unavailable or mocked — the fallback path268 produces the same byte-level contract (RIFF header + 16 kHz mono 16-bit).269 """270 271 if pcm.dtype != np.int16:272 clipped = np.clip(pcm.astype(np.float32), -1.0, 1.0)273 pcm_i16 = (clipped * 32767.0).astype(np.int16)274 else:275 pcm_i16 = pcm276 buf = io.BytesIO()277 with wave.open(buf, "wb") as w:278 w.setnchannels(1)279 w.setsampwidth(2)280 w.setframerate(sample_rate_hz)281 w.writeframes(pcm_i16.tobytes())282 return buf.getvalue()283 284 285# ---------------------------------------------------------------------------286# TTS287# ---------------------------------------------------------------------------288 289 290_TTS_CACHE_MAX_BYTES: int = 64 * 1024 * 1024291_TTS_CACHE_MAX_ENTRIES: int = 256292 293 294def _available_voice_packs(kokoro_module: Any) -> set[str]:295 """Probe the installed Kokoro bundle for shipped voice-pack names.296 297 Looks for ``AVAILABLE_VOICES``, ``list_voices()``, or ``VOICES``. A fresh298 install typically exposes at least one of these. If none is present we299 fall back to the full canonical set (best-effort; runtime per-call300 fallback in ``_resolve_voice_pack`` still protects against missing packs).301 """302 303 candidates: set[str] = set()304 for attr in ("AVAILABLE_VOICES", "VOICES"):305 value = getattr(kokoro_module, attr, None)306 if isinstance(value, (list, tuple, set, frozenset)):307 candidates.update(str(v) for v in value)308 list_voices = getattr(kokoro_module, "list_voices", None)309 if callable(list_voices):310 try:311 value = list_voices()312 if isinstance(value, (list, tuple, set, frozenset)):313 candidates.update(str(v) for v in value)314 except Exception: # pragma: no cover — defensive315 pass316 if not candidates:317 return set(_VOICE_PACKS_SET)318 return candidates319 320 321_FALLBACK_CHAIN: dict[str, str] = {322 "ta_female_1": "hi_female_1",323 "kn_male_1": "hi_female_1",324 "hi_male_1": "hi_female_1",325 "hi_female_1": "en_indian_female_1",326}327 328 329class TTSEngine:330 """Kokoro-82M wrapper. Constructed via ``get_tts_engine()``.331 332 One instance per process. All heavy deps are imported lazily.333 """334 335 def __init__(336 self,337 *,338 model_id: str = "hexgrad/Kokoro-82M",339 trace_sink: TraceSink | None = None,340 ) -> None:341 self._model_id = model_id342 self._trace_sink = trace_sink343 self._lock = threading.Lock()344 self._cache: LRUCache[tuple[Any, ...], bytes] = LRUCache(345 maxsize=_TTS_CACHE_MAX_BYTES, getsizeof=len346 )347 self._numpy_cache: LRUCache[tuple[Any, ...], np.ndarray] = LRUCache(348 maxsize=_TTS_CACHE_MAX_BYTES, getsizeof=lambda a: int(a.nbytes)349 )350 self._fallback_used: dict[str, str] = {}351 try:352 kokoro = _load_kokoro()353 except Exception as exc: # network / disk / import failure354 raise ModelLoadError(f"failed to load kokoro: {exc}") from exc355 self._kokoro = kokoro356 try:357 pipeline_cls = getattr(kokoro, "KPipeline", None)358 if pipeline_cls is None:359 raise AttributeError("kokoro.KPipeline missing")360 self._pipeline = pipeline_cls(model_id=model_id)361 except Exception as exc:362 raise ModelLoadError(f"failed to construct KPipeline: {exc}") from exc363 self._available_packs = _available_voice_packs(kokoro)364 self._verify_critical_packs()365 366 def _verify_critical_packs(self) -> None:367 if (368 "en_indian_female_1" not in self._available_packs369 and "hi_female_1" not in self._available_packs370 ):371 raise ModelLoadError("no usable voice pack for hi or en")372 373 def _resolve_voice_pack(self, requested: VoicePack) -> tuple[VoicePack, bool, str | None]:374 """Walk the fallback chain until an available pack is found.375 376 Returns ``(resolved_pack, degraded, fallback_from)``.377 """378 379 current = requested380 original = requested381 degraded = False382 fallback_from: str | None = None383 visited: set[str] = set()384 while current not in self._available_packs:385 if current in visited:386 break387 visited.add(current)388 successor = _FALLBACK_CHAIN.get(current)389 if successor is None:390 raise ModelLoadError(391 f"no usable voice pack; chain exhausted from {original!r}"392 )393 fallback_from = original394 current = cast("VoicePack", successor)395 degraded = True396 if degraded:397 self._fallback_used[original] = current398 return current, degraded, fallback_from399 400 def _emit_trace(self, trace: AudioTrace) -> None:401 if self._trace_sink is None:402 return403 try:404 self._trace_sink(trace)405 except Exception: # telemetry must never break production406 logger.debug("trace sink raised; swallowed", exc_info=True)407 408 def _render_pcm(self, text: str, voice_pack: VoicePack, seed: int) -> np.ndarray:409 """Invoke Kokoro inside a forked RNG context and return 24 kHz float32 PCM."""410 411 torch = _load_torch()412 with torch.random.fork_rng(devices=[]):413 torch.manual_seed(seed)414 try:415 result = self._pipeline(text, voice=voice_pack)416 except MemoryError as exc:417 raise TTSOutOfMemoryError(f"TTS OOM: {exc}") from exc418 except RuntimeError as exc:419 msg = str(exc).lower()420 if "out of memory" in msg or "alloc" in msg:421 raise TTSOutOfMemoryError(f"TTS OOM: {exc}") from exc422 raise423 return _coerce_to_float32_mono(result)424 425 def _resample_to_16k(self, pcm_24k: np.ndarray) -> np.ndarray:426 """Downsample 24 kHz → 16 kHz via torchaudio.functional.resample."""427 428 try:429 F = _load_torchaudio_functional()430 except Exception as exc: # pragma: no cover — hard runtime failure431 raise ModelLoadError(f"torchaudio.functional missing: {exc}") from exc432 torch = _load_torch()433 tensor = torch.from_numpy(pcm_24k.astype(np.float32)).unsqueeze(0)434 resampled = F.resample(435 tensor, orig_freq=24000, new_freq=16000, lowpass_filter_width=64436 )437 out = resampled.squeeze(0).cpu().numpy().astype(np.float32)438 return cast("np.ndarray", out)439 440 def _encode_wav(self, pcm_16k: np.ndarray, sample_rate_hz: int) -> bytes:441 """Encode the 16 kHz float32 PCM into 16-bit mono RIFF WAV bytes."""442 443 try:444 torchaudio = _load_torchaudio()445 torch = _load_torch()446 tensor = torch.from_numpy(pcm_16k.astype(np.float32)).unsqueeze(0)447 buf = io.BytesIO()448 torchaudio.save(449 buf,450 tensor,451 sample_rate=sample_rate_hz,452 bits_per_sample=16,453 format="wav",454 encoding="PCM_S",455 )456 return buf.getvalue()457 except Exception:458 # Fall back to stdlib wave encoder so the byte contract still holds459 # even when torchaudio is unavailable.460 return _np_to_wav_bytes(pcm_16k, sample_rate_hz)461 462 def synthesize(463 self,464 text: str,465 language_code: LanguageCode,466 voice_pack: VoicePack | None = None,467 *,468 seed: int = 0,469 sample_rate_hz: int = 16000,470 ) -> bytes:471 """Return 16-bit PCM mono WAV bytes. audio.md §2.1, §4.4."""472 473 if sample_rate_hz != 16000:474 raise UnsupportedLanguageError(475 f"sample_rate_hz={sample_rate_hz} unsupported; only 16000 allowed in v1"476 )477 if language_code not in _LANGUAGE_CODES:478 raise UnsupportedLanguageError(f"language_code={language_code!r} unsupported")479 mapping = VOICE_PACKS[language_code]480 if voice_pack is None:481 voice_pack = mapping.default482 if voice_pack not in mapping.allowed:483 raise UnsupportedVoicePackError(484 f"voice_pack={voice_pack!r} not allowed for language={language_code!r}"485 )486 text_hash = _input_hash(text.encode("utf-8"))487 cache_key = (text_hash, voice_pack, seed, sample_rate_hz, "bytes")488 start = time.perf_counter()489 with self._lock:490 cached = self._cache.get(cache_key)491 if cached is not None:492 latency_ms = int((time.perf_counter() - start) * 1000)493 duration_s = _wav_duration_s(cached)494 self._emit_trace(495 AudioTrace(496 op="synthesize",497 input_hash=text_hash,498 language=language_code,499 duration_s=duration_s,500 latency_ms=latency_ms,501 confidence=None,502 cache_hit=True,503 degraded=False,504 ts_ist=_ts_ist_now(),505 )506 )507 return cached508 resolved_pack, degraded, _ = self._resolve_voice_pack(voice_pack)509 pcm_24k = self._render_pcm(text, resolved_pack, seed)510 pcm_16k = self._resample_to_16k(pcm_24k)511 wav_bytes = self._encode_wav(pcm_16k, sample_rate_hz)512 with self._lock:513 self._cache[cache_key] = wav_bytes514 latency_ms = int((time.perf_counter() - start) * 1000)515 duration_s = _wav_duration_s(wav_bytes)516 self._emit_trace(517 AudioTrace(518 op="synthesize",519 input_hash=text_hash,520 language=language_code,521 duration_s=duration_s,522 latency_ms=latency_ms,523 confidence=None,524 cache_hit=False,525 degraded=degraded,526 ts_ist=_ts_ist_now(),527 )528 )529 return wav_bytes530 531 def synthesize_to_gradio(532 self,533 text: str,534 language_hint: LanguageCode,535 voice_pack: VoicePack | None = None,536 *,537 seed: int = 0,538 ) -> tuple[int, np.ndarray]:539 """Return ``(sample_rate, float32 mono ndarray)``. audio.md §2.1."""540 541 if language_hint not in _LANGUAGE_CODES:542 raise UnsupportedLanguageError(f"language_hint={language_hint!r} unsupported")543 mapping = VOICE_PACKS[language_hint]544 if voice_pack is None:545 voice_pack = mapping.default546 if voice_pack not in mapping.allowed:547 raise UnsupportedVoicePackError(548 f"voice_pack={voice_pack!r} not allowed for language={language_hint!r}"549 )550 text_hash = _input_hash(text.encode("utf-8"))551 sample_rate_hz = 16000552 cache_key = (text_hash, voice_pack, seed, sample_rate_hz, "numpy")553 start = time.perf_counter()554 with self._lock:555 cached = self._numpy_cache.get(cache_key)556 if cached is not None:557 self._emit_trace(558 AudioTrace(559 op="synthesize",560 input_hash=text_hash,561 language=language_hint,562 duration_s=float(len(cached)) / sample_rate_hz,563 latency_ms=int((time.perf_counter() - start) * 1000),564 confidence=None,565 cache_hit=True,566 degraded=False,567 ts_ist=_ts_ist_now(),568 )569 )570 return sample_rate_hz, cached.copy()571 resolved_pack, degraded, _ = self._resolve_voice_pack(voice_pack)572 pcm_24k = self._render_pcm(text, resolved_pack, seed)573 pcm_16k = self._resample_to_16k(pcm_24k)574 with self._lock:575 self._numpy_cache[cache_key] = pcm_16k576 self._emit_trace(577 AudioTrace(578 op="synthesize",579 input_hash=text_hash,580 language=language_hint,581 duration_s=float(len(pcm_16k)) / sample_rate_hz,582 latency_ms=int((time.perf_counter() - start) * 1000),583 confidence=None,584 cache_hit=False,585 degraded=degraded,586 ts_ist=_ts_ist_now(),587 )588 )589 return sample_rate_hz, pcm_16k.copy()590 591 def warmup(self) -> None:592 """Probe each voice pack; log WARN on missing Indic packs. audio.md §4.3.1."""593 594 for lang, mapping in VOICE_PACKS.items():595 for pack in mapping.allowed:596 if pack not in self._available_packs:597 logger.warning(598 "voice pack %r missing from bundle (language=%s); will fall back at synth time",599 pack,600 lang,601 )602 try:603 self.synthesize("warmup", "en")604 except Exception: # pragma: no cover — warmup best-effort605 logger.debug("warmup synthesize failed; continuing", exc_info=True)606 607 608def _coerce_to_float32_mono(result: Any) -> np.ndarray:609 """Turn whatever Kokoro returned into a 1-D float32 numpy array."""610 611 torch = _load_torch()612 if hasattr(result, "cpu") and hasattr(result, "numpy"):613 arr = result.detach().cpu().numpy()614 elif isinstance(result, tuple):615 audio_like = result[0]616 if hasattr(audio_like, "cpu") and hasattr(audio_like, "numpy"):617 arr = audio_like.detach().cpu().numpy()618 else:619 arr = np.asarray(audio_like)620 elif isinstance(result, np.ndarray):621 arr = result622 else:623 try:624 tensor = torch.as_tensor(result)625 arr = tensor.detach().cpu().numpy()626 except Exception as exc: # pragma: no cover — defensive627 raise TTSOutOfMemoryError(f"unexpected Kokoro return type: {type(result)!r}: {exc}") from exc628 arr = np.asarray(arr, dtype=np.float32).reshape(-1)629 return arr630 631 632def _wav_duration_s(wav_bytes: bytes) -> float:633 """Return the duration in seconds for a RIFF WAV payload (best-effort)."""634 635 try:636 with wave.open(io.BytesIO(wav_bytes), "rb") as w:637 frames = w.getnframes()638 rate = w.getframerate()639 if rate <= 0:640 return 0.0641 return round(frames / rate, 3)642 except Exception:643 return 0.0644 645 646# ---------------------------------------------------------------------------647# ASR648# ---------------------------------------------------------------------------649 650 651def _map_language(code: str | None) -> LanguageCode | Literal["unknown"]:652 if code in _LANGUAGE_CODES:653 return cast("LanguageCode", code)654 return "unknown"655 656 657def _nfc(text: str) -> str:658 return unicodedata.normalize("NFC", text).strip()659 660 661class ASREngine:662 """faster-whisper-small wrapper. Constructed via ``get_asr_engine()``.663 664 audio.md §2.2. Heavy deps loaded lazily.665 """666 667 def __init__(668 self,669 *,670 model_id: str = "Systran/faster-whisper-small",671 compute_type: Literal["int8", "int8_float16"] = "int8",672 trace_sink: TraceSink | None = None,673 ) -> None:674 self._model_id = model_id675 self._compute_type = compute_type676 self._trace_sink = trace_sink677 self._lock = threading.Lock()678 try:679 fw = _load_faster_whisper()680 except Exception as exc:681 raise ModelLoadError(f"failed to load faster_whisper: {exc}") from exc682 model_cls = getattr(fw, "WhisperModel", None)683 if model_cls is None:684 raise ModelLoadError("faster_whisper.WhisperModel missing")685 try:686 self._model = model_cls(model_id, compute_type=compute_type, device="cpu")687 except Exception as exc:688 raise ModelLoadError(f"failed to construct WhisperModel: {exc}") from exc689 690 def _emit_trace(self, trace: AudioTrace) -> None:691 if self._trace_sink is None:692 return693 try:694 self._trace_sink(trace)695 except Exception:696 logger.debug("trace sink raised; swallowed", exc_info=True)697 698 def transcribe(699 self,700 audio_bytes: bytes,701 language_hint: LanguageCode | None,702 *,703 beam_size: int = 1,704 vad_filter: bool = True,705 max_duration_s: float = 30.0,706 ) -> TranscriptResult:707 """Decode WAV/PCM bytes. audio.md §2.2, §3.5, §4.4."""708 709 start = time.perf_counter()710 pcm, clip_duration = self._decode_input(audio_bytes)711 if clip_duration > max_duration_s:712 pcm = pcm[: int(max_duration_s * 16000)]713 clip_duration = max_duration_s714 language_for_whisper: str | None715 if language_hint == "hinglish":716 language_for_whisper = "hi"717 elif language_hint is None:718 language_for_whisper = None719 else:720 language_for_whisper = language_hint721 segments, info = self._run_whisper(722 pcm,723 language=language_for_whisper,724 beam_size=beam_size,725 vad_filter=vad_filter,726 )727 segments_list = list(segments)728 detected_code = _map_language(getattr(info, "language", None))729 vad_dropped_all = getattr(info, "vad_dropped_all_segments", None)730 if vad_dropped_all is None:731 vad_dropped_all = len(segments_list) == 0 and vad_filter732 combined_text = _nfc("".join(getattr(s, "text", "") for s in segments_list))733 duration_s = round(min(float(clip_duration), float(max_duration_s)), 3)734 degraded = False735 if combined_text == "":736 confidence = 0.0737 if vad_dropped_all:738 detected: LanguageCode | Literal["unknown"] = "unknown"739 else:740 detected = detected_code741 degraded = True742 else:743 confidence = _duration_weighted_confidence(segments_list)744 detected = _infer_hinglish(detected_code, combined_text, language_hint)745 result = TranscriptResult(746 text=combined_text,747 language_detected=detected,748 confidence=confidence,749 duration_s=duration_s,750 )751 latency_ms = int((time.perf_counter() - start) * 1000)752 self._emit_trace(753 AudioTrace(754 op="transcribe",755 input_hash=_input_hash(audio_bytes),756 language=language_hint or "unknown",757 duration_s=duration_s,758 latency_ms=latency_ms,759 confidence=confidence,760 cache_hit=False,761 degraded=degraded,762 ts_ist=_ts_ist_now(),763 )764 )765 return result766 767 def _decode_input(self, audio_bytes: bytes) -> tuple[np.ndarray, float]:768 """Return (float32 mono @ 16 kHz, duration_s); raise AudioDecodeError on mismatch."""769 770 if len(audio_bytes) >= 3 and audio_bytes[:3] == b"ID3":771 raise AudioDecodeError("MP3 / ID3-tagged inputs are not supported (no ffmpeg in image)")772 rate = _riff_header_sample_rate(audio_bytes)773 if rate is not None:774 if rate != 16000:775 raise AudioDecodeError("input must be 16 kHz mono; caller must pre-resample")776 try:777 sf = _load_soundfile()778 data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32", always_2d=False)779 except Exception as exc:780 raise AudioDecodeError(f"soundfile failed to decode RIFF WAV: {exc}") from exc781 if sr != 16000:782 raise AudioDecodeError("input must be 16 kHz mono; caller must pre-resample")783 arr = np.asarray(data, dtype=np.float32).reshape(-1)784 duration = float(len(arr)) / 16000.0785 return arr, duration786 # Raw float32 PCM path (demo mic input). 16 kHz assumed. We only accept787 # payloads that look like plausible audio — ≥ 0.25 s of float32 samples788 # (4000 × 4 = 16000 bytes) whose magnitudes fit inside the normalized789 # [-1, 1] range that Gradio emits. Short / out-of-range payloads are790 # rejected so arbitrary random bytes do not slip through.791 min_raw_pcm_bytes = 4000 * 4792 if len(audio_bytes) >= min_raw_pcm_bytes and len(audio_bytes) % 4 == 0:793 pcm = np.frombuffer(audio_bytes, dtype=np.float32).copy()794 if pcm.size and np.all(np.isfinite(pcm)) and np.max(np.abs(pcm)) <= 2.0:795 duration = float(pcm.size) / 16000.0796 return pcm, duration797 raise AudioDecodeError("input is not a valid 16 kHz RIFF WAV or float32 PCM payload")798 799 def _run_whisper(800 self,801 pcm: np.ndarray,802 *,803 language: str | None,804 beam_size: int,805 vad_filter: bool,806 ) -> tuple[Any, Any]:807 try:808 segments, info = self._model.transcribe(809 pcm,810 language=language,811 beam_size=beam_size,812 vad_filter=vad_filter,813 )814 except Exception as exc:815 raise AudioDecodeError(f"whisper decode failed: {exc}") from exc816 return segments, info817 818 def warmup(self) -> None:819 """Run one transcribe() on 0.5 s of silence to force load. audio.md §2.2."""820 821 silence = _pcm16_silence_wav(0.5)822 try:823 self.transcribe(silence, "en")824 except Exception: # pragma: no cover — warmup best-effort825 logger.debug("warmup transcribe failed; continuing", exc_info=True)826 827 828def _duration_weighted_confidence(segments: list[Any]) -> float:829 if not segments:830 return 0.0831 total_dur = 0.0832 weighted = 0.0833 for seg in segments:834 start = float(getattr(seg, "start", 0.0) or 0.0)835 end = float(getattr(seg, "end", 0.0) or 0.0)836 dur = max(0.0, end - start)837 avg_logprob = float(getattr(seg, "avg_logprob", 0.0) or 0.0)838 confidence = _logprob_to_confidence(avg_logprob)839 if dur == 0.0:840 total_dur += 1.0841 weighted += confidence842 else:843 total_dur += dur844 weighted += confidence * dur845 if total_dur == 0.0:846 return 0.0847 return round(weighted / total_dur, 3)848 849 850def _infer_hinglish(851 detected: LanguageCode | Literal["unknown"],852 text: str,853 hint: LanguageCode | None,854) -> LanguageCode | Literal["unknown"]:855 """Downgrade ``hi`` to ``hinglish`` when the decoded text is code-mixed.856 857 Heuristic per audio.md §3.6: ≥ 2 ASCII words intermixed with Devanagari.858 """859 860 if hint != "hinglish":861 return detected862 if detected != "hi":863 return detected864 ascii_words = [tok for tok in text.split() if tok.isascii() and tok.isalpha()]865 has_devanagari = any("ऀ" <= ch <= "ॿ" for ch in text)866 if len(ascii_words) >= 2 and has_devanagari:867 return "hinglish"868 return detected869 870 871# ---------------------------------------------------------------------------872# Singletons873# ---------------------------------------------------------------------------874 875 876_tts_engine: TTSEngine | None = None877_asr_engine: ASREngine | None = None878_tts_lock = threading.Lock()879_asr_lock = threading.Lock()880 881 882def get_tts_engine(883 *, trace_sink: TraceSink | None = None, model_id: str = "hexgrad/Kokoro-82M"884) -> TTSEngine:885 """Return the process-wide TTSEngine singleton. audio.md §3.2, §3.8."""886 887 global _tts_engine888 with _tts_lock:889 if _tts_engine is None:890 _tts_engine = TTSEngine(model_id=model_id, trace_sink=trace_sink)891 elif trace_sink is not None and trace_sink is not _tts_engine._trace_sink:892 logger.warning("get_tts_engine: different sink passed after construction; ignoring")893 return _tts_engine894 895 896def get_asr_engine(897 *,898 trace_sink: TraceSink | None = None,899 model_id: str = "Systran/faster-whisper-small",900 compute_type: Literal["int8", "int8_float16"] = "int8",901) -> ASREngine:902 """Return the process-wide ASREngine singleton. audio.md §3.2, §3.8."""903 904 global _asr_engine905 with _asr_lock:906 if _asr_engine is None:907 _asr_engine = ASREngine(908 model_id=model_id, compute_type=compute_type, trace_sink=trace_sink909 )910 elif trace_sink is not None and trace_sink is not _asr_engine._trace_sink:911 logger.warning("get_asr_engine: different sink passed after construction; ignoring")912 return _asr_engine913 914 915def _reset_singletons_for_tests() -> None:916 """Tear down singletons. Tests only. audio.md §3.2 "Unload. Never." exemption."""917 918 global _tts_engine, _asr_engine919 with _tts_lock:920 _tts_engine = None921 with _asr_lock:922 _asr_engine = None923 924 925__all__ = [926 "AudioDecodeError",927 "AudioError",928 "AudioTooLongError",929 "AudioTrace",930 "ASREngine",931 "LanguageCode",932 "ModelLoadError",933 "TTSEngine",934 "TTSOutOfMemoryError",935 "TranscriptResult",936 "TraceSink",937 "UnsupportedLanguageError",938 "UnsupportedVoicePackError",939 "VOICE_PACKS",940 "VoicePack",941 "VoicePackMapping",942 "get_asr_engine",943 "get_tts_engine",944]945