visualnovel/tts.py

"""The Voice — text to speech.

Synthesises NPC dialogue using Kokoro-82M via ONNX (~82M params, no transformers dependency).

Design notes:
- Voice selection reads Character.tts_voice_description (frozen at creation) and scores every
  Kokoro preset by keyword overlap — the best-scoring voice wins.  sprite_seed breaks ties so
  two characters with the same description still get different voices.
- `tts_voice_description` is frozen at character creation (see state.apply_directives and
  orchestrator.init_world), ensuring voice identity is stable for the whole session.
- Results are cached in memory by SHA-256(text + voice_id + seed) so repeated lines are free.

Model files (~337 MB total) live in models/kokoro/ and are downloaded once:
  curl -L -o models/kokoro/kokoro-v1.0.onnx \\
    https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/kokoro-v1.0.onnx
  curl -L -o models/kokoro/voices-v1.0.bin \\
    https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/voices-v1.0.bin
"""

from __future__ import annotations

import hashlib
import io
from typing import Protocol

from . import config

# ---------------------------------------------------------------------------
# Voice profiles — (voice_id, gender, keyword_set)
# gender: "f" | "m"
# keyword_set: words in a voice descriptor that suggest this voice
# ---------------------------------------------------------------------------
_VOICE_PROFILES: list[tuple[str, str, frozenset[str]]] = [
    # ── American Female ────────────────────────────────────────────────────
    (
        "af_heart",
        "f",
        frozenset({"warm", "bright", "breathless", "tender", "sweet", "laughs", "light"}),
    ),
    (
        "af_bella",
        "f",
        frozenset({"soft", "gentle", "quiet", "shy", "delicate", "whisper", "whispered"}),
    ),
    ("af_sarah", "f", frozenset({"clear", "professional", "calm", "steady", "cool", "composed"})),
    ("af_sky", "f", frozenset({"airy", "bubbly", "cheerful", "energetic", "lively", "peppy"})),
    (
        "af_jessica",
        "f",
        frozenset({"bold", "confident", "sharp", "spirited", "assertive", "direct"}),
    ),
    ("af_nicole", "f", frozenset({"rich", "smooth", "soothing", "mature", "warm", "velvet"})),
    ("af_nova", "f", frozenset({"crisp", "precise", "intelligent", "focused", "analytical"})),
    ("af_river", "f", frozenset({"flowing", "peaceful", "serene", "dreamy", "drifts", "trails"})),
    ("af_kore", "f", frozenset({"strong", "powerful", "serious", "stern", "commanding"})),
    ("af_aoede", "f", frozenset({"melodic", "musical", "lyrical", "poetic", "ethereal", "sings"})),
    ("af_alloy", "f", frozenset({"neutral", "balanced", "modern", "even"})),
    # ── British Female ─────────────────────────────────────────────────────
    ("bf_alice", "f", frozenset({"british", "english", "uk", "posh", "accent", "clipped"})),
    ("bf_emma", "f", frozenset({"british", "english", "uk", "homely", "friendly", "warm"})),
    (
        "bf_isabella",
        "f",
        frozenset({"british", "english", "uk", "refined", "elegant", "sophisticated"}),
    ),
    ("bf_lily", "f", frozenset({"british", "english", "uk", "gentle", "soft", "sweet"})),
    # ── American Male ─────────────────────────────────────────────────────
    ("am_adam", "m", frozenset({"confident", "deep", "steady", "bass", "strong", "rich"})),
    ("am_echo", "m", frozenset({"resonant", "full", "commanding", "measured"})),
    ("am_eric", "m", frozenset({"friendly", "casual", "warm", "approachable", "cheerful"})),
    (
        "am_fenrir",
        "m",
        frozenset({"gruff", "rough", "intense", "dark", "rugged", "growl", "harsh"}),
    ),
    ("am_liam", "m", frozenset({"young", "energetic", "bright", "enthusiastic", "lively"})),
    ("am_michael", "m", frozenset({"mature", "calm", "authoritative", "professorial", "quiet"})),
    ("am_onyx", "m", frozenset({"smooth", "velvet", "baritone", "low", "silky", "suave"})),
    ("am_puck", "m", frozenset({"playful", "mischievous", "quick", "witty", "impish", "teasing"})),
    # ── British Male ──────────────────────────────────────────────────────
    ("bm_daniel", "m", frozenset({"british", "english", "uk", "smooth", "refined", "crisp"})),
    (
        "bm_fable",
        "m",
        frozenset({"british", "english", "uk", "storyteller", "dramatic", "expressive"}),
    ),
    ("bm_george", "m", frozenset({"british", "english", "uk", "distinguished", "formal", "noble"})),
    ("bm_lewis", "m", frozenset({"british", "english", "uk", "young", "casual", "easy"})),
]

_MALE_GENDER_HINTS = {"man", "male", "his", "him", "he", "boy", "lad", "gentleman"}
_FEMALE_GENDER_HINTS = {"woman", "female", "her", "she", "girl", "lady"}

_MODEL_PATH = config.MODELS_DIR / "kokoro" / "kokoro-v1.0.onnx"
_VOICES_PATH = config.MODELS_DIR / "kokoro" / "voices-v1.0.bin"
_KOKORO_BASE_URL = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0"


def _ensure_kokoro_models() -> None:
    """Download Kokoro model files if absent (handles HF Space ephemeral filesystem)."""
    files = {
        _MODEL_PATH: f"{_KOKORO_BASE_URL}/kokoro-v1.0.onnx",
        _VOICES_PATH: f"{_KOKORO_BASE_URL}/voices-v1.0.bin",
    }
    missing = {p: u for p, u in files.items() if not p.exists()}
    if not missing:
        return
    import urllib.request

    _MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    for path, url in missing.items():
        print(f"[tts] Downloading {path.name} …", flush=True)
        urllib.request.urlretrieve(url, path)
        print(f"[tts] {path.name} ready ({path.stat().st_size // 1_000_000} MB)", flush=True)


def _pick_voice(voice_description: str, sprite_seed: int) -> str:
    """Score every Kokoro voice against the description; use sprite_seed to break ties."""
    words = set(voice_description.lower().split())

    # Detect gender from description
    is_male = bool(words & _MALE_GENDER_HINTS)
    is_female = bool(words & _FEMALE_GENDER_HINTS)
    # If ambiguous, default to female (most anime VN characters)
    target_gender = "m" if is_male and not is_female else "f"

    candidates = [(vid, kws) for vid, g, kws in _VOICE_PROFILES if g == target_gender]

    # Score: number of keyword matches, then deterministic tie-break by seed
    scored = [(len(words & kws), vid) for vid, kws in candidates]
    scored.sort(key=lambda x: x[0], reverse=True)

    # Among voices tied at the top score, pick by sprite_seed
    top_score = scored[0][0]
    top_voices = [vid for score, vid in scored if score == top_score]
    return top_voices[sprite_seed % len(top_voices)]


class TTSBackend(Protocol):
    def synthesize(self, text: str, voice_description: str, seed: int) -> bytes | None: ...


class MockTTS:
    def synthesize(self, text: str, voice_description: str, seed: int) -> bytes | None:
        return None


class KokoroTTS:
    def __init__(self) -> None:
        import logging  # noqa: PLC0415

        from kokoro_onnx import Kokoro  # noqa: PLC0415

        # phonemizer resets its logger level on use, so setLevel from app.py doesn't
        # stick — a filter survives ("words count mismatch" fires on nearly every line).
        logging.getLogger("phonemizer").addFilter(lambda r: r.levelno >= logging.ERROR)

        _ensure_kokoro_models()
        self._kokoro = Kokoro(str(_MODEL_PATH), str(_VOICES_PATH))
        # Filter profiles to voices actually present in the loaded model
        available = set(self._kokoro.get_voices())
        self._available_voices = available
        self._cache: dict[str, bytes] = {}

    def synthesize(self, text: str, voice_description: str, seed: int) -> bytes | None:
        import soundfile as sf  # noqa: PLC0415

        voice_id = _pick_voice(voice_description, seed)
        # Fallback if the chosen voice isn't in this model build
        if voice_id not in self._available_voices:
            fallback = [
                vid for vid, g, _ in _VOICE_PROFILES if g == "f" and vid in self._available_voices
            ]
            voice_id = (
                fallback[seed % len(fallback)] if fallback else next(iter(self._available_voices))
            )

        text = text.strip()
        if not text:
            return None

        cache_key = hashlib.sha256(f"{text}\x00{voice_id}\x00{seed}".encode()).hexdigest()
        if cache_key in self._cache:
            return self._cache[cache_key]

        try:
            samples, sample_rate = self._kokoro.create(
                text, voice=voice_id, speed=1.0, lang="en-us"
            )
        except (ValueError, Exception):
            return None

        if samples is None or (hasattr(samples, "__len__") and len(samples) == 0):
            return None

        buf = io.BytesIO()
        sf.write(buf, samples, sample_rate, format="WAV")
        wav_bytes = buf.getvalue()

        self._cache[cache_key] = wav_bytes
        return wav_bytes


def get_tts() -> TTSBackend:
    if config.TTS_BACKEND == "kokoro":
        return KokoroTTS()
    return MockTTS()