watch_sceen_ai/wsai/backends/melo.py

"""Real Korean TTS via MeloTTS, run as a persistent out-of-venv worker.

MeloTTS needs its own interpreter (melo311). Loading the model costs several
seconds, so we keep one worker process alive and stream synthesis requests to
it (see melo_worker.py for the protocol). Each `speak()` writes a wav to
`out_dir` and hands the path to a sink (default: log it). The Discord voice
integration later swaps the sink for "play this wav into the call".

Env:
    WSAI_MELO_PYTHON   interpreter with melo installed
                       (default: /home/claude/jarvis-tts/melo311/bin/python)
    WSAI_MELO_DEVICE   cpu | cuda | auto  (default auto: GPU if torch sees one,
                       else CPU; the worker falls back to CPU if CUDA fails)
    WSAI_TTS_OUT_DIR   where wavs are written (default ~/.cache/wsai/tts)
    WSAI_TTS_SPEED     synthesis speed multiplier (default 1.3)
"""

from __future__ import annotations

import asyncio
import collections
import json
import logging
import os
import time
from pathlib import Path
from typing import Awaitable, Callable

from ..interfaces import Reply

log = logging.getLogger("wsai.tts.melo")

_DEFAULT_PYTHON = "/home/claude/jarvis-tts/melo311/bin/python"

# A sink receives the finished wav path plus the reply it voices.
Sink = Callable[[str, Reply], Awaitable[None]]


async def _log_sink(path: str, reply: Reply) -> None:
    log.info("TTS wav ready: %s (%s)", path, reply.text[:40])


class MeloTTS:
    def __init__(
        self,
        *,
        python: str | None = None,
        device: str | None = None,
        out_dir: str | None = None,
        speed: float | None = None,
        sink: Sink | None = None,
    ) -> None:
        self.python = python or os.environ.get("WSAI_MELO_PYTHON", _DEFAULT_PYTHON)
        self.device = device or os.environ.get("WSAI_MELO_DEVICE", "auto")
        self.out_dir = Path(out_dir or os.environ.get("WSAI_TTS_OUT_DIR")
                            or (Path.home() / ".cache/wsai/tts"))
        self.speed = float(speed if speed is not None
                           else os.environ.get("WSAI_TTS_SPEED", "1.3"))
        self.sink = sink or _log_sink
        self._proc: asyncio.subprocess.Process | None = None
        self._lock = asyncio.Lock()
        self._n = 0
        self.load_ms: int | None = None
        # Keep the worker's most recent stderr lines so a crash reports its real
        # cause instead of a bare JSONDecodeError. Bounded so it can't grow.
        self._stderr_tail: collections.deque[str] = collections.deque(maxlen=40)
        self._stderr_task: asyncio.Task | None = None

    async def _drain_stderr(self, stream: asyncio.StreamReader) -> None:
        # The worker redirects fd1 -> fd2, so ALL library chatter lands on
        # stderr. If we PIPE stderr but never read it, the OS pipe buffer fills
        # and the worker blocks forever. So we continuously drain it and keep
        # only the last few lines for diagnostics.
        try:
            while True:
                line = await stream.readline()
                if not line:
                    return
                self._stderr_tail.append(line.decode(errors="replace").rstrip())
        except asyncio.CancelledError:
            raise
        except Exception:  # draining must never crash the caller
            return

    def _stderr_hint(self) -> str:
        tail = "\n".join(self._stderr_tail)
        return f" worker stderr tail:\n{tail}" if tail else " (worker produced no stderr)"

    async def warmup(self) -> None:
        """Start and load the worker now so the first real utterance is warm.

        Called at pipeline startup so users don't wait ~7s (CPU model load) for
        the very first spoken reply."""
        await self._ensure()

    async def _ensure(self) -> None:
        if self._proc is not None and self._proc.returncode is None:
            return
        self.out_dir.mkdir(parents=True, exist_ok=True)
        env = {**os.environ, "WSAI_MELO_DEVICE": self.device}
        # Run the worker module from the wsai source tree with the melo venv.
        repo_root = str(Path(__file__).resolve().parents[2])
        self._proc = await asyncio.create_subprocess_exec(
            self.python, "-m", "wsai.backends.melo_worker",
            cwd=repo_root, env=env,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        self._stderr_tail.clear()
        assert self._proc.stderr is not None
        self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr))
        ready = await self._proc.stdout.readline()
        if not ready:  # worker died before signalling ready
            await self._proc.wait()
            raise RuntimeError(
                f"melo worker exited before ready (code {self._proc.returncode})."
                f"{self._stderr_hint()}"
            )
        try:
            info = json.loads(ready.decode())
        except json.JSONDecodeError as exc:
            raise RuntimeError(
                f"melo worker sent invalid ready line {ready!r}: {exc}."
                f"{self._stderr_hint()}"
            ) from exc
        if not info.get("ready"):
            raise RuntimeError(
                f"melo worker failed to start: {info}.{self._stderr_hint()}"
            )
        self.load_ms = info.get("ms")
        log.info("melo worker ready in %s ms on %s", self.load_ms, info.get("device"))

    async def speak(self, reply: Reply) -> None:
        await self._ensure()
        self._n += 1
        out = str(self.out_dir / f"tts-{self._n:06d}.wav")
        req = json.dumps({"text": reply.text, "out": out, "speed": self.speed})
        s = time.monotonic()
        async with self._lock:
            assert self._proc and self._proc.stdin and self._proc.stdout
            self._proc.stdin.write((req + "\n").encode())
            await self._proc.stdin.drain()
            resp = await self._proc.stdout.readline()
        if not resp:
            raise RuntimeError(f"melo worker closed unexpectedly.{self._stderr_hint()}")
        res = json.loads(resp.decode())
        if not res.get("ok"):
            raise RuntimeError(f"melo synth failed: {res.get('error')}")
        log.debug("synth %d ms (worker %s ms)", int((time.monotonic() - s) * 1000), res.get("ms"))
        await self.sink(res["out"], reply)

    async def aclose(self) -> None:
        if self._proc is not None and self._proc.returncode is None:
            try:
                self._proc.terminate()
                await asyncio.wait_for(self._proc.wait(), timeout=5)
            except (ProcessLookupError, asyncio.TimeoutError):
                pass
        if self._stderr_task is not None:
            self._stderr_task.cancel()
            try:
                await self._stderr_task
            except (asyncio.CancelledError, Exception):
                pass
            self._stderr_task = None
        self._proc = None