javis_bot/bridge/xtts_worker.py

"""
XTTS worker
===========

A tiny HTTP service that keeps a Coqui XTTS-v2 voice warm and synthesises
speech on demand. It mirrors ``melo_worker.py`` (same ``/synth`` + ``/health``
contract, same PCM16 WAV output) so the bridge can talk to either worker the
same way.

XTTS-v2 is a natural, multilingual neural voice. The default speaker is the
built-in female studio voice "Ana Florence" speaking Korean — the voice this
deployment uses in place of MeloTTS. No reference WAV is needed for the
built-in studio speakers.

It runs in its OWN Python venv (``/opt/xtts`` in the container) so the heavy
Coqui TTS / torch stack stays isolated from the slim brain-bridge venv.

Config (env):
    XTTS_WORKER_HOST   bind host            (default 127.0.0.1)
    XTTS_WORKER_PORT   bind port            (default 8771)
    XTTS_MODEL         Coqui model id       (default tts_models/multilingual/multi-dataset/xtts_v2)
    XTTS_SPEAKER       built-in speaker     (default "Ana Florence")
    XTTS_LANGUAGE      synthesis language   (default ko)
    XTTS_DEVICE        torch device         (default cpu; compose sets cuda)

Run:
    /opt/xtts/bin/python -m bridge.xtts_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

# XTTS-v2 is gated behind a one-time license prompt; agreeing here keeps the
# load non-interactive in a container. XTTS-v2 is non-commercial (CPML).
os.environ.setdefault("COQUI_TOS_AGREED", "1")

HOST = os.environ.get("XTTS_WORKER_HOST", "127.0.0.1")
PORT = int(os.environ.get("XTTS_WORKER_PORT", "8771"))
MODEL = os.environ.get("XTTS_MODEL", "tts_models/multilingual/multi-dataset/xtts_v2")
SPEAKER = os.environ.get("XTTS_SPEAKER", "Ana Florence")
LANGUAGE = os.environ.get("XTTS_LANGUAGE", "ko")
DEVICE = os.environ.get("XTTS_DEVICE", "cpu")

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


def _ensure_model() -> None:
    global _model, _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 TTS.api import TTS  # type: ignore

            model = TTS(MODEL).to(DEVICE)
            _model = model
            # Warm once: the first GPU synth pays a one-off kernel-init cost
            # that would otherwise land on the user's first reply.
            try:
                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as _wt:
                    _wp = _wt.name
                model.tts_to_file(
                    text="워밍업", speaker=SPEAKER, language=LANGUAGE, file_path=_wp
                )
                try:
                    os.unlink(_wp)
                except OSError:
                    pass
            except Exception as _we:  # pragma: no cover
                print(f"[xtts-worker] warmup synth skipped: {_we}", flush=True)
            print(
                f"[xtts-worker] ready (model={MODEL} speaker={SPEAKER!r} "
                f"language={LANGUAGE} device={DEVICE})",
                flush=True,
            )
        except Exception as e:  # pragma: no cover - depends on local model files
            _load_error = f"{type(e).__name__}: {e}"
            print(f"[xtts-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 "xtts model unavailable")
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        tmp_path = tmp.name
    try:
        with _model_lock:
            _model.tts_to_file(
                text=text, speaker=SPEAKER, language=LANGUAGE, file_path=tmp_path
            )
        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. Coqui writes float/other WAVs; the Discord
    playback path tolerates both, but we normalise to PCM16 so the contract
    matches the previous Melo/Piper output (mono, file's own sample rate)."""
    try:
        with wave.open(io.BytesIO(raw), "rb") as wf:
            if wf.getsampwidth() == 2:
                return raw  # already PCM16
    except wave.Error:
        pass
    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 XTTS 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"[xtts-worker] listening on http://{HOST}:{PORT}", flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    return 0


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