mcspeak/src/tts_mcp/engines/orpheus.py

"""Orpheus TTS via Ollama + SNAC decoder.

Sends text to Ollama's Orpheus model, parses <custom_token_N> responses,
and decodes through SNAC to 24kHz WAV. SNAC is lazy-loaded on first use
and runs on CPU (RTX 5070 SM 120 not yet supported by PyTorch 2.6).
"""

import asyncio
import os
import re
import sys
import time

import numpy as np
import requests

from ..audio import wav_duration, write_wav
from ..settings import settings
from .base import TTSEngine, TTSResult

SAMPLE_RATE = 24000
SNAC_CODEBOOK_SIZE = 4096
TOKEN_PATTERN = re.compile(r"<custom_token_(\d+)>")

ALL_VOICES = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"]


def _turn_token_into_id(token_str: str, index: int) -> int:
    """Convert a <custom_token_N> string to a SNAC codebook ID.

    Applies position-dependent offset: each of the 7 tokens per frame
    maps to a different SNAC codebook layer at a different offset.
    """
    match = TOKEN_PATTERN.search(token_str)
    if not match:
        return -1
    raw_id = int(match.group(1))
    return raw_id - 10 - ((index % 7) * SNAC_CODEBOOK_SIZE)


def _tokens_to_audio(token_strings: list[str], snac_model) -> np.ndarray | None:
    """Convert Orpheus custom token strings to audio via SNAC.

    Redistributes the flat token stream into SNAC's 3 codebook layers:
      - codes_0: 1 per frame (positions 0)     -> coarse
      - codes_1: 2 per frame (positions 1, 4)  -> mid
      - codes_2: 4 per frame (positions 2,3,5,6) -> fine
    """
    import torch

    num_frames = len(token_strings) // 7
    if num_frames == 0:
        return None

    token_strings = token_strings[: num_frames * 7]
    ids = [_turn_token_into_id(t, i) for i, t in enumerate(token_strings)]

    invalid = sum(1 for x in ids if x < 0 or x >= SNAC_CODEBOOK_SIZE)
    if invalid > 0:
        print(f"  Warning: {invalid}/{len(ids)} invalid token IDs", file=sys.stderr)

    codes_0, codes_1, codes_2 = [], [], []
    for i in range(num_frames):
        b = i * 7
        codes_0.append(ids[b + 0])
        codes_1.append(ids[b + 1])
        codes_2.append(ids[b + 2])
        codes_2.append(ids[b + 3])
        codes_1.append(ids[b + 4])
        codes_2.append(ids[b + 5])
        codes_2.append(ids[b + 6])

    def clamp(lst):
        return [max(0, min(SNAC_CODEBOOK_SIZE - 1, x)) for x in lst]

    device = "cpu"
    codes = [
        torch.tensor(clamp(codes_0), dtype=torch.long).unsqueeze(0).to(device),
        torch.tensor(clamp(codes_1), dtype=torch.long).unsqueeze(0).to(device),
        torch.tensor(clamp(codes_2), dtype=torch.long).unsqueeze(0).to(device),
    ]

    with torch.no_grad():
        audio = snac_model.decode(codes)

    return audio.squeeze().cpu().numpy()


def _load_snac():
    """Load SNAC model on CPU. Called once, lazily."""
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    from snac import SNAC

    return SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to("cpu")


class OrpheusEngine(TTSEngine):
    """Orpheus TTS via Ollama's completions API + SNAC audio decoding.

    SNAC is lazy-loaded on first synthesize() call to avoid holding
    ~200MB of RAM when the engine isn't being used.
    """

    name = "orpheus"
    default_voice = "tara"

    def __init__(self, ollama_url: str, model_name: str) -> None:
        self._snac = None
        self._snac_lock = asyncio.Lock()
        self._ollama_url = ollama_url
        self._model = model_name

    async def _get_snac(self):
        """Lazy-load SNAC on first use."""
        if self._snac is not None:
            return self._snac
        async with self._snac_lock:
            if self._snac is not None:
                return self._snac
            print("Loading SNAC decoder on CPU (first Orpheus call)...", file=sys.stderr)
            t0 = time.time()
            loop = asyncio.get_running_loop()
            self._snac = await loop.run_in_executor(None, _load_snac)
            print(f"  SNAC ready in {time.time() - t0:.1f}s", file=sys.stderr)
            return self._snac

    async def synthesize(self, text: str, voice: str | None = None) -> TTSResult:
        voice = voice or self.default_voice
        prompt = f"<|audio|>{voice}: {text}<|eot_id|>"

        loop = asyncio.get_running_loop()

        # Ollama API call (blocking HTTP)
        t0 = time.time()

        def _call_ollama():
            resp = requests.post(
                f"{self._ollama_url}/v1/completions",
                json={
                    "model": self._model,
                    "prompt": prompt,
                    "max_tokens": 8192,
                    "temperature": 0.6,
                    "top_p": 0.9,
                    "stream": False,
                },
                timeout=600,
            )
            resp.raise_for_status()
            return resp.json()

        result = await loop.run_in_executor(None, _call_ollama)
        gen_time = time.time() - t0

        resp_text = result.get("choices", [{}])[0].get("text", "")

        # Extract <custom_token_N> strings
        token_strings = TOKEN_PATTERN.findall(resp_text)
        token_strings = [f"<custom_token_{t}>" for t in token_strings]

        # Skip leading special tokens (values < 10)
        skip = 0
        for ts in token_strings:
            m = TOKEN_PATTERN.search(ts)
            if m and int(m.group(1)) < 10:
                skip += 1
            else:
                break
        if skip > 0:
            token_strings = token_strings[skip:]

        print(
            f"  Orpheus: {len(token_strings)} tokens "
            f"({len(token_strings) // 7} frames) in {gen_time:.1f}s",
            file=sys.stderr,
        )

        if len(token_strings) < 7:
            raise RuntimeError(
                f"Orpheus returned insufficient tokens ({len(token_strings)}). "
                f"Response preview: {resp_text[:200]}"
            )

        # Lazy-load SNAC, then decode (CPU-bound)
        snac = await self._get_snac()
        audio = await loop.run_in_executor(
            None, _tokens_to_audio, token_strings, snac
        )
        if audio is None:
            raise RuntimeError("SNAC decoding produced no audio")

        path = write_wav(audio, SAMPLE_RATE, prefix="orpheus-")

        return TTSResult(
            audio_path=path,
            sample_rate=SAMPLE_RATE,
            duration_seconds=wav_duration(path),
            engine=self.name,
            voice=voice,
        )

    async def list_voices(self) -> list[str]:
        blacklist = settings.blacklisted_voices
        return sorted(v for v in ALL_VOICES if v.lower() not in blacklist)

    async def check_health(self) -> dict:
        try:
            resp = requests.get(f"{self._ollama_url}/api/tags", timeout=5)
            resp.raise_for_status()
            models = [m["name"] for m in resp.json().get("models", [])]
            has_orpheus = any("orpheus" in m.lower() for m in models)
            return {
                "status": "healthy" if has_orpheus else "degraded",
                "engine": self.name,
                "model_loaded": has_orpheus,
                "ollama_models": len(models),
            }
        except Exception as e:
            return {"status": "unhealthy", "engine": self.name, "error": str(e)}