Blame for services/tts_service.py

"""
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