FabledScribe/src/fabledassistant/services/tts.py

"""Text-to-speech service using Kokoro TTS (in-process)."""
import asyncio
import io
import logging
import os
import time
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from kokoro import KPipeline

logger = logging.getLogger(__name__)

_pipeline: "KPipeline | None" = None
_pipeline_lock = asyncio.Lock()
_load_error: str | None = None

# Repo identifier used for HuggingFace update checks
KOKORO_REPO = "hexgrad/Kokoro-82M"

# Persisted across restarts so we can detect model updates and run offline otherwise
_COMMIT_HASH_FILE = Path("/data/kokoro_commit_hash.txt")

# Static list of supported Kokoro voice IDs and display labels
_VOICES: list[dict] = [
    {"id": "af_heart", "label": "Heart (American Female, warm)"},
    {"id": "af_bella", "label": "Bella (American Female, expressive)"},
    {"id": "af_nicole", "label": "Nicole (American Female, intimate)"},
    {"id": "af_sarah", "label": "Sarah (American Female, clear)"},
    {"id": "af_sky", "label": "Sky (American Female, bright)"},
    {"id": "am_adam", "label": "Adam (American Male, neutral)"},
    {"id": "am_michael", "label": "Michael (American Male, deep)"},
    {"id": "bf_emma", "label": "Emma (British Female)"},
    {"id": "bf_isabella", "label": "Isabella (British Female, formal)"},
    {"id": "bm_george", "label": "George (British Male)"},
    {"id": "bm_lewis", "label": "Lewis (British Male, casual)"},
]


def _read_stored_commit() -> str | None:
    try:
        if _COMMIT_HASH_FILE.exists():
            return _COMMIT_HASH_FILE.read_text().strip() or None
    except Exception:
        pass
    return None


def _write_stored_commit(sha: str) -> None:
    try:
        _COMMIT_HASH_FILE.parent.mkdir(parents=True, exist_ok=True)
        _COMMIT_HASH_FILE.write_text(sha)
    except Exception:
        logger.warning("Failed to write Kokoro commit hash to %s", _COMMIT_HASH_FILE)


def _set_hf_offline(offline: bool) -> None:
    """Toggle HuggingFace offline mode in the current process environment."""
    if offline:
        os.environ["HF_HUB_OFFLINE"] = "1"
    else:
        os.environ.pop("HF_HUB_OFFLINE", None)


async def load_tts_model() -> None:
    """Load the Kokoro pipeline. Called once at startup when voice is enabled."""
    global _pipeline, _load_error
    from fabledassistant.services.voice_config import is_voice_enabled

    if not await is_voice_enabled():
        return

    async with _pipeline_lock:
        if _pipeline is not None:
            return
        try:
            from kokoro import KPipeline

            # If we have a stored commit hash the model files are already cached
            # locally — run offline to skip HuggingFace network validation entirely.
            already_cached = _read_stored_commit() is not None
            if already_cached:
                _set_hf_offline(True)
                logger.info("Kokoro model previously cached — loading in offline mode")

            logger.info("Loading Kokoro TTS pipeline...")
            loop = asyncio.get_running_loop()

            def _load_and_warm():
                p = KPipeline(lang_code="a")  # "a" = American English
                # Pre-load all voice tensors so synthesis never hits HuggingFace at request time
                for v in _VOICES:
                    try:
                        p.load_voice(v["id"])
                    except Exception:
                        pass
                return p

            _pipeline = await loop.run_in_executor(None, _load_and_warm)

            # Record the current commit hash on first successful load so future
            # restarts know the model is cached and can skip HF network checks.
            if not already_cached:
                asyncio.create_task(_record_initial_commit())

            logger.info("Kokoro TTS pipeline loaded and voices pre-warmed")
        except Exception:
            _load_error = "Failed to load Kokoro TTS pipeline"
            logger.exception(_load_error)


async def _record_initial_commit() -> None:
    """Fetch and store the current Kokoro commit hash after first download."""
    try:
        loop = asyncio.get_running_loop()

        def _fetch():
            from huggingface_hub import model_info
            return model_info(KOKORO_REPO).sha

        sha = await loop.run_in_executor(None, _fetch)
        if sha:
            _write_stored_commit(sha)
            # Now that files are cached, switch to offline mode for subsequent runs
            _set_hf_offline(True)
            logger.info("Kokoro commit hash stored (%s) — future restarts will use offline mode", sha[:8])
    except Exception:
        logger.warning("Could not record Kokoro commit hash", exc_info=True)


async def check_for_kokoro_updates() -> None:
    """Check HuggingFace for Kokoro model updates.

    Intended to be called on a daily schedule. Temporarily lifts offline mode
    to query the HF API, compares the latest commit SHA against the locally
    stored one, and reloads the pipeline only if the model has changed.
    """
    from fabledassistant.services.voice_config import is_voice_enabled
    if not await is_voice_enabled():
        return

    stored_sha = _read_stored_commit()

    try:
        loop = asyncio.get_running_loop()

        def _fetch_latest_sha():
            # Temporarily lift offline mode to reach the HF API
            _set_hf_offline(False)
            try:
                from huggingface_hub import model_info
                return model_info(KOKORO_REPO).sha
            finally:
                # Restore offline mode regardless of outcome
                _set_hf_offline(True)

        latest_sha = await loop.run_in_executor(None, _fetch_latest_sha)
    except Exception:
        logger.warning("Kokoro update check failed — will retry tomorrow", exc_info=True)
        # Ensure offline mode is restored if the executor raised before the finally
        _set_hf_offline(bool(stored_sha))
        return

    if latest_sha == stored_sha:
        logger.debug("Kokoro model is up to date (sha: %s)", (latest_sha or "")[:8])
        return

    logger.info(
        "Kokoro model update detected (%s → %s), reloading pipeline",
        (stored_sha or "none")[:8],
        (latest_sha or "")[:8],
    )
    # Go online for the reload so Kokoro can download the updated files
    _set_hf_offline(False)
    try:
        await reload_tts_model()
        if latest_sha:
            _write_stored_commit(latest_sha)
        logger.info("Kokoro model updated and reloaded successfully")
    except Exception:
        logger.exception("Kokoro pipeline reload after update failed")
    finally:
        _set_hf_offline(True)


async def reload_tts_model() -> None:
    """Unload and reload the TTS pipeline. Safe to call at runtime."""
    global _pipeline, _load_error
    async with _pipeline_lock:
        _pipeline = None
        _load_error = None
    await load_tts_model()


def tts_available() -> bool:
    return _pipeline is not None


def list_voices() -> list[dict]:
    return _VOICES


async def synthesise(
    text: str,
    voice: str = "af_heart",
    speed: float = 1.0,
    voice_blend: list[dict] | None = None,
) -> bytes:
    """Synthesise text to WAV bytes (24kHz, 16-bit mono). Runs in executor.

    voice_blend is a list of {"voice": str, "weight": float} dicts.
    When provided with 2+ entries the voice style tensors are merged as a
    weighted average before synthesis. Weights are normalised automatically.
    """
    if _pipeline is None:
        raise RuntimeError("TTS pipeline not loaded")

    speed = max(0.7, min(1.3, speed))

    def _build_voice_param():
        """Return either a blended style tensor or a single voice ID string."""
        if not voice_blend or len(voice_blend) < 2:
            return voice_blend[0]["voice"] if voice_blend else voice
        import numpy as np
        total_w = sum(max(0.0, e.get("weight", 1.0)) for e in voice_blend) or 1.0
        blended = None
        for entry in voice_blend:
            vid = entry.get("voice", "af_heart")
            w = max(0.0, entry.get("weight", 1.0)) / total_w
            vt = _pipeline.load_voice(vid)  # type: ignore[union-attr]
            blended = vt * w if blended is None else blended + vt * w
        return blended

    def _run() -> bytes:
        import numpy as np
        import soundfile as sf

        voice_param = _build_voice_param()
        t0 = time.monotonic()
        audio_chunks: list = []
        try:
            for _, _, audio in _pipeline(text, voice=voice_param, speed=speed):  # type: ignore[misc]
                if audio is not None:
                    audio_chunks.append(audio)
        except Exception:
            logger.exception(
                "Kokoro pipeline error during synthesis: %d chars, preview=%r",
                len(text), text[:80],
            )
            raise

        if not audio_chunks:
            logger.warning(
                "Kokoro produced no audio chunks: %d chars, preview=%r",
                len(text), text[:80],
            )
            return b""

        combined = np.concatenate(audio_chunks)
        buf = io.BytesIO()
        sf.write(buf, combined, samplerate=24000, format="WAV", subtype="PCM_16")
        elapsed = time.monotonic() - t0
        logger.info(
            "Kokoro synthesis: %d chars → %d samples (%.2fs, %.0f chars/s)",
            len(text), len(combined), elapsed, len(text) / elapsed if elapsed > 0 else 0,
        )
        return buf.getvalue()

    loop = asyncio.get_running_loop()
    return await loop.run_in_executor(None, _run)