Spaces:
Sleeping
Sleeping
| """ | |
| Text-to-speech using Kokoro-82M (https://huggingface.co/hexgrad/Kokoro-82M). | |
| Two interchangeable backends, selected by TTS_BACKEND (auto | onnx | torch): | |
| - "onnx" — kokoro-onnx running the **fp32** ONNX export (bit-identical weights | |
| to the torch release, so no quality loss) + misaki for English G2P (the same | |
| phonemizer the torch pipeline uses, so pronunciation is unchanged). 2-4x | |
| faster than torch on CPU and keeps torch out of the image entirely. | |
| Needs the model files (see Dockerfile): | |
| KOKORO_ONNX_MODEL (default models/kokoro-v1.0.onnx) | |
| KOKORO_ONNX_VOICES (default models/voices-v1.0.bin) | |
| - "torch" — the original `kokoro` KPipeline (pip install kokoro soundfile). | |
| "auto" prefers ONNX when its deps + model files are present, else torch, else | |
| the service reports unavailable and callers fall back to browser speech. | |
| Everything loads lazily on first use (call `warmup()` at startup to pre-pay | |
| that cost off the request path). Synthesis results are LRU-cached per | |
| (text, voice, format) — interview questions are static text, so replays and | |
| re-entered rounds are served instantly. A semaphore keeps concurrent syntheses | |
| from thrashing each other's CPU threads. | |
| """ | |
| import io | |
| import os | |
| import re | |
| import threading | |
| from collections import OrderedDict | |
| from pathlib import Path | |
| # American English by default. The interviewer persona ("Folio") uses a male | |
| # voice — am_michael is the highest-rated US male in Kokoro v1.0 (am_fenrir and | |
| # am_puck are the alternatives; af_heart is the best female if ever needed). | |
| DEFAULT_VOICE = "am_michael" | |
| SAMPLE_RATE = 24000 | |
| _SERVICE_ROOT = Path(__file__).resolve().parent.parent | |
| _backend = None | |
| _backend_error: str | None = None | |
| _init_lock = threading.Lock() | |
| # Only TTS_MAX_CONCURRENCY syntheses run at once; both backends parallelize a | |
| # single utterance across cores, so more concurrency usually means all requests | |
| # get slower, not faster. | |
| _synth_sem = threading.BoundedSemaphore(int(os.getenv("TTS_MAX_CONCURRENCY", "1"))) | |
| _cache: "OrderedDict[tuple, tuple[bytes, str]]" = OrderedDict() | |
| _cache_lock = threading.Lock() | |
| _CACHE_MAX = int(os.getenv("TTS_CACHE_SIZE", "256")) | |
| # Second cache level: raw float32 samples per (text, voice). Synthesis is the | |
| # expensive step (seconds); encoding is milliseconds — so a request for the | |
| # same text in a different container (e.g. Safari asking for wav after the | |
| # cache was warmed with ogg) re-encodes instead of re-synthesizing. | |
| _sample_cache: "OrderedDict[tuple, object]" = OrderedDict() | |
| _SAMPLE_CACHE_MAX = int(os.getenv("TTS_SAMPLE_CACHE_SIZE", "128")) | |
| def _onnx_paths() -> tuple[Path, Path]: | |
| model = Path(os.getenv("KOKORO_ONNX_MODEL", "models/kokoro-v1.0.onnx")) | |
| voices = Path(os.getenv("KOKORO_ONNX_VOICES", "models/voices-v1.0.bin")) | |
| if not model.is_absolute(): | |
| model = _SERVICE_ROOT / model | |
| if not voices.is_absolute(): | |
| voices = _SERVICE_ROOT / voices | |
| return model, voices | |
| def _split_chunks(text: str, max_len: int = 300) -> list[str]: | |
| """Sentence-ish chunks under max_len chars (stays well inside Kokoro's | |
| phoneme-sequence limit). The torch pipeline splits internally; the ONNX | |
| path needs us to do it.""" | |
| sentences = re.findall(r"[^.?!]+[.?!]+[\"')\]]*\s*|[^.?!]+$", text) | |
| chunks: list[str] = [] | |
| current = "" | |
| for s in sentences: | |
| s = s.strip() | |
| if not s: | |
| continue | |
| if current and len(current) + len(s) + 1 > max_len: | |
| chunks.append(current) | |
| current = s | |
| else: | |
| current = f"{current} {s}".strip() | |
| while len(current) > max_len: # single overlong sentence: hard split | |
| chunks.append(current[:max_len]) | |
| current = current[max_len:].strip() | |
| if current: | |
| chunks.append(current) | |
| return chunks or [text] | |
| class _OnnxBackend: | |
| """kokoro-onnx (fp32 weights) + misaki G2P — torch-free, CPU-fast.""" | |
| name = "onnx" | |
| def __init__(self): | |
| from kokoro_onnx import Kokoro | |
| from misaki import en | |
| model, voices = _onnx_paths() | |
| if not model.exists() or not voices.exists(): | |
| raise FileNotFoundError(f"Kokoro ONNX model files missing ({model}, {voices})") | |
| # Same G2P stack as the torch pipeline: misaki dictionary first, | |
| # espeak-ng for out-of-dictionary words when available. | |
| fallback = None | |
| try: | |
| from misaki import espeak | |
| fallback = espeak.EspeakFallback(british=False) | |
| except Exception: | |
| pass # espeak-ng not installed — dictionary-only G2P still works | |
| self._g2p = en.G2P(trf=False, british=False, fallback=fallback) | |
| self._kokoro = Kokoro(str(model), str(voices)) | |
| def segments(self, text: str, voice: str): | |
| import numpy as np | |
| for chunk in _split_chunks(text): | |
| phonemes, _ = self._g2p(chunk) | |
| if not phonemes: | |
| continue | |
| samples, _sr = self._kokoro.create(phonemes, voice=voice, is_phonemes=True) | |
| yield np.asarray(samples, dtype="float32") | |
| class _TorchBackend: | |
| """The original kokoro KPipeline (heavy: pulls in torch).""" | |
| name = "torch" | |
| def __init__(self): | |
| from kokoro import KPipeline | |
| threads = os.getenv("TTS_TORCH_THREADS") | |
| if threads: | |
| import torch | |
| torch.set_num_threads(max(1, int(threads))) | |
| self._pipeline = KPipeline(lang_code="a") # 'a' = American English | |
| def segments(self, text: str, voice: str): | |
| import numpy as np | |
| for _, _, audio in self._pipeline(text, voice=voice): | |
| # Kokoro yields a torch tensor (or ndarray) of float samples per segment. | |
| if hasattr(audio, "detach"): | |
| audio = audio.detach().cpu().numpy() | |
| yield np.asarray(audio, dtype="float32") | |
| def _get_backend(): | |
| """Lazily build (and cache) the configured backend.""" | |
| global _backend, _backend_error | |
| if _backend is not None: | |
| return _backend | |
| with _init_lock: | |
| if _backend is not None: | |
| return _backend | |
| choice = os.getenv("TTS_BACKEND", "auto").lower() | |
| errors = [] | |
| if choice in ("auto", "onnx"): | |
| try: | |
| _backend = _OnnxBackend() | |
| except Exception as e: | |
| errors.append(f"onnx: {e}") | |
| if choice == "onnx": | |
| _backend_error = "; ".join(errors) | |
| raise | |
| if _backend is None and choice in ("auto", "torch"): | |
| try: | |
| _backend = _TorchBackend() | |
| except Exception as e: | |
| errors.append(f"torch: {e}") | |
| if _backend is None: | |
| _backend_error = "; ".join(errors) or f"unknown TTS_BACKEND {choice!r}" | |
| raise RuntimeError(f"No TTS backend available ({_backend_error})") | |
| _backend_error = None | |
| return _backend | |
| def is_available() -> bool: | |
| """True if some backend can plausibly run (cheap check, no model load).""" | |
| if _backend is not None: | |
| return True | |
| try: | |
| import kokoro_onnx # noqa: F401 | |
| import misaki # noqa: F401 | |
| import soundfile # noqa: F401 | |
| model, voices = _onnx_paths() | |
| if model.exists() and voices.exists(): | |
| return True | |
| except Exception: | |
| pass | |
| try: | |
| import kokoro # noqa: F401 | |
| import soundfile # noqa: F401 | |
| return True | |
| except Exception: | |
| return False | |
| def backend_status() -> dict: | |
| """For /tts/health: which backend is configured/loaded.""" | |
| return { | |
| "configured": os.getenv("TTS_BACKEND", "auto"), | |
| "loaded": _backend.name if _backend is not None else None, | |
| "error": _backend_error, | |
| } | |
| def _cache_get(key: tuple) -> tuple[bytes, str] | None: | |
| with _cache_lock: | |
| hit = _cache.get(key) | |
| if hit is not None: | |
| _cache.move_to_end(key) | |
| return hit | |
| def _cache_put(key: tuple, value: tuple[bytes, str]) -> None: | |
| with _cache_lock: | |
| _cache[key] = value | |
| _cache.move_to_end(key) | |
| while len(_cache) > _CACHE_MAX: | |
| _cache.popitem(last=False) | |
| def _samples_get(key: tuple): | |
| with _cache_lock: | |
| hit = _sample_cache.get(key) | |
| if hit is not None: | |
| _sample_cache.move_to_end(key) | |
| return hit | |
| def _samples_put(key: tuple, samples) -> None: | |
| with _cache_lock: | |
| _sample_cache[key] = samples | |
| _sample_cache.move_to_end(key) | |
| while len(_sample_cache) > _SAMPLE_CACHE_MAX: | |
| _sample_cache.popitem(last=False) | |
| def _encode(samples, fmt: str) -> tuple[bytes, str]: | |
| """Encode float32 samples → (bytes, mime). OGG/Opus is ~10x smaller than | |
| WAV for speech; falls back to WAV if libsndfile lacks Opus support.""" | |
| import soundfile as sf | |
| buf = io.BytesIO() | |
| if fmt == "ogg": | |
| try: | |
| sf.write(buf, samples, SAMPLE_RATE, format="OGG", subtype="OPUS") | |
| return buf.getvalue(), "audio/ogg" | |
| except Exception: | |
| buf = io.BytesIO() | |
| sf.write(buf, samples, SAMPLE_RATE, format="WAV") | |
| return buf.getvalue(), "audio/wav" | |
| def synthesize(text: str, voice: str | None = None, fmt: str = "wav") -> tuple[bytes, str]: | |
| """ | |
| Synthesize `text` → (audio bytes, mime type). Raises on any failure so the | |
| router can return 503 and the client can fall back to browser TTS. | |
| """ | |
| import numpy as np | |
| voice = voice or DEFAULT_VOICE | |
| fmt = "ogg" if str(fmt).lower() in ("ogg", "opus") else "wav" | |
| key = (text, voice, fmt) | |
| skey = (text, voice) | |
| cached = _cache_get(key) | |
| if cached is not None: | |
| return cached | |
| samples = _samples_get(skey) | |
| if samples is None: | |
| with _synth_sem: | |
| cached = _cache_get(key) # may have been filled while we waited | |
| if cached is not None: | |
| return cached | |
| samples = _samples_get(skey) | |
| if samples is None: | |
| backend = _get_backend() | |
| chunks = list(backend.segments(text, voice)) | |
| if not chunks: | |
| raise RuntimeError("Kokoro produced no audio") | |
| samples = np.concatenate(chunks) if len(chunks) > 1 else chunks[0] | |
| _samples_put(skey, samples) | |
| result = _encode(samples, fmt) | |
| _cache_put(key, result) | |
| return result | |
| def warmup() -> bool: | |
| """Build the backend and run a tiny synthesis so the first real request | |
| doesn't pay model-load + kernel-JIT cost. Safe to call in a thread; never | |
| raises (unavailability is already handled per-request).""" | |
| try: | |
| synthesize("Hi.", fmt="wav") | |
| return True | |
| except Exception: | |
| return False | |