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

_FENCE_RE = re.compile(r"```[^\n`]*\n?(.*?)```", re.DOTALL)
_LINK_RE = re.compile(r"\[([^\]]+)\]\([^)]*\)")
_INLINE_CODE_RE = re.compile(r"`+([^`]*)`+")


def normalize_for_speech(text: str) -> str:
    """Flatten Claude's markdown/code formatting into plain prose before TTS.

    MeloTTS's Korean text normaliser has no dictionary entry for characters
    like the backtick and dies with ``KeyError: '`'`` — which crashes the whole
    voice turn the moment the model mentions a command or shows a code block.
    Code/markdown also reads terribly aloud. So strip the formatting and keep
    the words. Every spoken path (dashboard voice turn and the Discord speak()
    bridge) funnels through ``synth()``, so normalising there covers them both.
    """
    if not text:
        return text
    # Fenced code block -> keep its inner text as spoken words, drop the fences.
    text = _FENCE_RE.sub(lambda m: " " + m.group(1) + " ", text)
    # [label](url) -> label
    text = _LINK_RE.sub(r"\1", text)
    # `code` -> code
    text = _INLINE_CODE_RE.sub(r"\1", text)
    # Any stray/unbalanced backtick that survived -> gone. This is the exact
    # character that crashes MeloTTS, so guarantee none remain.
    text = text.replace("`", "")
    # Markdown structure markers -> plain text.
    text = re.sub(r"(?m)^\s{0,3}#{1,6}\s*", "", text)   # ATX headings
    text = re.sub(r"(?m)^\s{0,3}>\s?", "", text)         # blockquotes
    text = re.sub(r"(?m)^\s{0,3}[-*+]\s+", "", text)     # bullet list markers
    text = re.sub(r"[*_]{1,3}", "", text)                # bold/italic emphasis
    # Collapse the whitespace the stripping leaves behind.
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{2,}", "\n", text)
    return text.strip()

# 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 synth(self, text: str) -> str:
        """Synthesize `text` to a wav and return its path (no sink). Reusable by
        callers that want the wav directly (e.g. the Discord voice bridge)."""
        await self._ensure()
        text = normalize_for_speech(text)
        self._n += 1
        out = str(self.out_dir / f"tts-{self._n:06d}.wav")
        req = json.dumps({"text": 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"))
        return res["out"]

    async def speak(self, reply: Reply) -> None:
        out = await self.synth(reply.text)
        await self.sink(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