javis_bot/bridge/melo_worker.py

"""
MeloTTS worker
==============

A tiny, dependency-light HTTP service that keeps a MeloTTS voice warm and
synthesises speech on demand. It runs in its OWN Python venv (``/opt/melo`` in
the container) so the heavy MeloTTS/torch/transformers stack stays isolated
from the slim brain-bridge venv (which pins ``numpy<2`` for faster-whisper).

The bridge's ``synthesize()`` POSTs ``{"text": "..."}`` here and gets back a
16-bit PCM WAV. The MeloTTS model is loaded once at startup and reused, so each
request only pays inference cost, not model-load cost.

Config (env):
    MELO_WORKER_HOST   bind host          (default 127.0.0.1)
    MELO_WORKER_PORT   bind port          (default 8770)
    MELO_LANGUAGE      MeloTTS language   (default KR)
    MELO_SPEED         speaking rate      (default 1.5  -> the approved "150")
    MELO_DEVICE        torch device       (default cpu)

Run:
    /opt/melo/bin/python -m bridge.melo_worker
"""

from __future__ import annotations

import io
import json
import os
import sys
import tempfile
import threading
import wave
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

HOST = os.environ.get("MELO_WORKER_HOST", "127.0.0.1")
PORT = int(os.environ.get("MELO_WORKER_PORT", "8770"))
LANGUAGE = os.environ.get("MELO_LANGUAGE", "KR")


def _resolve_speed() -> float:
    """Speaking rate: the settings-UI value (runtime config JSON) wins, else the
    MELO_SPEED env, else 1.5. Read at startup; the settings UI restarts this
    worker on apply so a new value takes effect."""
    try:
        cp = os.environ.get("JARVIS_CONFIG_PATH", "/app/config/jarvis.json")
        v = json.loads(open(cp, encoding="utf-8").read()).get("melo_speed")
        if v is not None:
            return float(v)
    except Exception:
        pass
    try:
        return float(os.environ.get("MELO_SPEED", "1.5"))
    except ValueError:
        return 1.5


SPEED = _resolve_speed()
DEVICE = os.environ.get("MELO_DEVICE", "cpu")

# Model + speaker id are loaded once, guarded by a lock because MeloTTS
# inference is not guaranteed thread-safe.
_model = None
_speaker_id = None
_model_lock = threading.Lock()
_load_error: str | None = None


def _ensure_model() -> None:
    global _model, _speaker_id, _load_error
    if _model is not None or _load_error is not None:
        return
    with _model_lock:
        if _model is not None or _load_error is not None:
            return
        try:
            from melo.api import TTS  # type: ignore

            model = TTS(language=LANGUAGE, device=DEVICE)
            # spk2id is a melo HParams object (dict-like, supports __getitem__,
            # __contains__, keys) but NOT .get(). The KR model exposes a single
            # 'KR' speaker; fall back to the first id for other languages.
            spk_map = model.hps.data.spk2id
            keys = list(spk_map.keys())
            speaker_id = spk_map[LANGUAGE] if LANGUAGE in spk_map else spk_map[keys[0]]
            _model = model
            _speaker_id = speaker_id
            # Warm the GPU once at load: the first CUDA synth pays a one-off
            # kernel-init cost (~5s) that would otherwise land on the user's
            # first reply. A throwaway synth here moves it to startup. No-op
            # cost on CPU.
            try:
                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as _wt:
                    _wp = _wt.name
                model.tts_to_file("워밍업", speaker_id, _wp, speed=SPEED)
                try:
                    os.unlink(_wp)
                except OSError:
                    pass
            except Exception as _we:  # pragma: no cover
                print(f"[melo-worker] warmup synth skipped: {_we}", flush=True)
            print(
                f"[melo-worker] ready (lang={LANGUAGE} speed={SPEED} "
                f"device={DEVICE} speakers={list(spk_map.keys())})",
                flush=True,
            )
        except Exception as e:  # pragma: no cover - depends on local model files
            _load_error = f"{type(e).__name__}: {e}"
            print(f"[melo-worker] model load FAILED: {_load_error}", flush=True)


def _synthesize(text: str) -> bytes:
    """Synthesise ``text`` to a 16-bit PCM WAV (bytes)."""
    _ensure_model()
    if _model is None:
        raise RuntimeError(_load_error or "melo model unavailable")
    # MeloTTS writes to a file via soundfile; render to a container-disk temp
    # file (NOT tmpfs), read it back, then drop it.
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        tmp_path = tmp.name
    try:
        with _model_lock:
            _model.tts_to_file(text, _speaker_id, tmp_path, speed=SPEED)
        with open(tmp_path, "rb") as f:
            raw = f.read()
    finally:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass
    return _ensure_pcm16_wav(raw)


def _ensure_pcm16_wav(raw: bytes) -> bytes:
    """Guarantee a 16-bit PCM WAV. MeloTTS/soundfile usually emit float WAVs;
    the Discord playback path (ffmpeg) tolerates both, but we normalise to
    PCM16 so the contract matches the previous Piper output."""
    try:
        with wave.open(io.BytesIO(raw), "rb") as wf:
            if wf.getsampwidth() == 2:
                return raw  # already PCM16
    except wave.Error:
        pass
    # Non-PCM16 (e.g. float) — convert with soundfile if available.
    try:
        import numpy as np
        import soundfile as sf

        data, sr = sf.read(io.BytesIO(raw), dtype="float32")
        if data.ndim > 1:
            data = data.mean(axis=1)  # mono
        pcm = np.clip(data, -1.0, 1.0)
        pcm = (pcm * 32767.0).astype("<i2").tobytes()
        buf = io.BytesIO()
        with wave.open(buf, "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(int(sr))
            wf.writeframes(pcm)
        return buf.getvalue()
    except Exception:
        return raw  # last resort: hand back whatever MeloTTS produced


class _Handler(BaseHTTPRequestHandler):
    def _json(self, code: int, payload: dict) -> None:
        body = json.dumps(payload).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):  # noqa: N802
        if self.path == "/health":
            _ensure_model()
            ok = _model is not None
            self._json(200 if ok else 503, {"ok": ok, "error": _load_error})
        else:
            self._json(404, {"error": "not found"})

    def do_POST(self):  # noqa: N802
        if self.path != "/synth":
            self._json(404, {"error": "not found"})
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
            data = json.loads(self.rfile.read(length) or b"{}")
            text = (data.get("text") or "").strip()
        except Exception as e:
            self._json(400, {"error": f"bad request: {e}"})
            return
        if not text:
            self._json(400, {"error": "missing 'text'"})
            return
        try:
            wav = _synthesize(text)
        except Exception as e:
            self._json(503, {"error": f"{type(e).__name__}: {e}"})
            return
        self.send_response(200)
        self.send_header("Content-Type", "audio/wav")
        self.send_header("Content-Length", str(len(wav)))
        self.end_headers()
        self.wfile.write(wav)

    def log_message(self, *args):  # silence default request logging
        return


def main() -> int:
    # Warm the model at startup so the first Discord turn isn't slow.
    _ensure_model()
    server = ThreadingHTTPServer((HOST, PORT), _Handler)
    print(f"[melo-worker] listening on http://{HOST}:{PORT}", flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    return 0


if __name__ == "__main__":
    sys.exit(main())