speech/app/tts/workers/xtts_worker.py

from __future__ import annotations

import argparse
import asyncio
import gc
import os
import tempfile
import time
from pathlib import Path
from typing import Any, Dict

import sys
sys.path.insert(0, "/app")
from tts.config import TTS_MODEL_CACHE, ensure_runtime_dirs

# TTS.api를 import하기 전에 설정해야 적용된다.
os.environ.setdefault("TTS_HOME", str(TTS_MODEL_CACHE))
os.environ.setdefault("COQUI_TOS_AGREED", "1")  # XTTS(CPML 라이선스) 동의 프롬프트를 비대화식으로 통과
# 컨테이너는 983 유저로 실행되는데 HOME(/app)이 root 소유라 numba/matplotlib/torch가
# 각자 ~/.cache, ~/.config 아래에 쓰려다 실패한다. HOME을 쓰기 가능한 곳으로 돌린다.
os.environ["HOME"] = "/tmp"  # 컨테이너 기본 HOME=/app은 983 유저가 쓰기 불가 (setdefault로는 덮어쓰기 안 됨)
os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache")

import soundfile as sf
import uvicorn
from fastapi import FastAPI, Form, HTTPException
from fastapi.responses import Response

app = FastAPI(title="TTS XTTS Worker")

MODEL_NAME = "tts_models/multilingual/multi-dataset/xtts_v2"
_MODEL_CACHE: Dict[str, Any] = {}

# Idle-unload: the model is dropped after IDLE_UNLOAD_SECONDS of no synthesis
# requests, so this backend doesn't permanently hog GPU memory shared with
# the ASR workers.
IDLE_UNLOAD_SECONDS = 120
_last_used: float = 0.0
_active_requests: int = 0


def _log(msg: str) -> None:
    print(f"[xtts] {msg}", flush=True)


def _device() -> str:
    try:
        import torch
        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:
        return "cpu"


def _free_gpu() -> None:
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


def _unload_model() -> None:
    if not _MODEL_CACHE:
        return
    _log("unloading model (idle)")
    _MODEL_CACHE.clear()
    _free_gpu()


async def _idle_unload_loop() -> None:
    while True:
        await asyncio.sleep(30)
        if _active_requests == 0 and _last_used and (time.monotonic() - _last_used) >= IDLE_UNLOAD_SECONDS:
            _unload_model()


def _load_model() -> Any:
    if MODEL_NAME in _MODEL_CACHE:
        return _MODEL_CACHE[MODEL_NAME]

    from TTS.api import TTS

    _log(f"loading model {MODEL_NAME}")
    try:
        tts = TTS(MODEL_NAME).to(_device())
    except Exception as e:
        # A load interrupted partway (e.g. OOM) can leave CUDA memory
        # fragmented/leaked in ways gc.collect()+empty_cache() don't reliably
        # reclaim, and since _MODEL_CACHE never got populated the idle-unload
        # loop has nothing to clean up either. Restarting the whole process
        # is the only guaranteed way to get that memory back — supervisord's
        # autorestart=true respawns it immediately.
        _log(f"model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}")
        os._exit(1)
    _MODEL_CACHE[MODEL_NAME] = tts
    _log("model loaded")
    return tts


@app.on_event("startup")
async def startup() -> None:
    ensure_runtime_dirs()
    asyncio.create_task(_idle_unload_loop())


@app.get("/health")
def health() -> Dict[str, Any]:
    return {"status": "ok", "device": _device(), "loaded": MODEL_NAME in _MODEL_CACHE}


@app.post("/synthesize")
def synthesize(
    text: str = Form(...),
    language: str = Form("ko"),
    speaker_wav: str = Form(...),
) -> Response:
    speaker_path = Path(speaker_wav)
    if not speaker_path.exists():
        raise HTTPException(status_code=400, detail=f"speaker_wav not found: {speaker_wav}")

    global _active_requests, _last_used
    _active_requests += 1
    try:
        tts = _load_model()
    except Exception as e:
        _log(f"model load error: {type(e).__name__}: {e}")
        _active_requests -= 1
        _last_used = time.monotonic()
        raise HTTPException(status_code=500, detail=f"모델 로드 실패: {type(e).__name__}: {e}")

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        out_path = Path(tmp.name)

    try:
        _log(f"synthesizing language={language} chars={len(text)}")
        tts.tts_to_file(
            text=text,
            speaker_wav=str(speaker_path),
            language=language,
            file_path=str(out_path),
        )
        audio_bytes = out_path.read_bytes()
        duration = sf.info(str(out_path)).duration
        return Response(
            content=audio_bytes,
            media_type="audio/wav",
            headers={"x-audio-duration": str(round(duration, 3))},
        )
    except HTTPException:
        raise
    except Exception as e:
        _log(f"synthesize error: {type(e).__name__}: {e}")
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
    finally:
        out_path.unlink(missing_ok=True)
        _free_gpu()
        _active_requests -= 1
        _last_used = time.monotonic()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8005)
    args = parser.parse_args()
    ensure_runtime_dirs()
    _log(f"starting host={args.host} port={args.port} device={_device()}")
    uvicorn.run(app, host=args.host, port=args.port)