docker/rp_handler.py

"""RunPod serverless handler for BlueTTS ONNX (GPU).

- Lazy-loads models on first job
- Async + asyncio.to_thread so concurrency_modifier is real
- Returns OGG/Opus by default (ffmpeg); WAV optional via input.format
- Stage timings: g2p_s, inference_s, encode_s, total_s
- Concurrency via env MAX_CONCURRENCY (default 2)
"""
from __future__ import annotations

import asyncio
import base64
import io
import os
import subprocess
import sys
import tempfile
import threading
import time
import traceback
from typing import Any, Dict, Optional, Tuple

# Must be set before importing app (skips Gradio Blocks construction).
os.environ.setdefault("BLUE_DOCKER", "1")
os.environ.setdefault("BLUE_SERVE_MODE", "serverless")
os.environ.setdefault("BLUE_SKIP_GRADIO", "1")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")

import numpy as np
import runpod
import soundfile as sf

_blue = None
_load_error: Optional[str] = None
# ORT InferenceSession is not safely shared across concurrent Run() calls.
_synth_lock = threading.Lock()


def _max_concurrency() -> int:
    raw = os.environ.get("MAX_CONCURRENCY", "2").strip() or "2"
    try:
        n = int(raw)
    except ValueError:
        n = 2
    return max(1, min(n, 16))


def _concurrency_modifier(_current: int) -> int:
    return _max_concurrency()


def _ensure_blue():
    """Import app + ONNX sessions once (first job / preload)."""
    global _blue, _load_error
    if _blue is not None:
        return _blue
    if _load_error is not None:
        raise RuntimeError(_load_error)
    print("[rp_handler] loading BlueTTS ONNX …", flush=True)
    t0 = time.time()
    try:
        import app as blue  # noqa: WPS433 — deferred on purpose
        _blue = blue
        print(
            f"[rp_handler] ready in {time.time() - t0:.1f}s "
            f"voices={list(blue.VOICE_STYLES.keys())} "
            f"langs={list(blue.AVAILABLE_LANGS)} "
            f"onnx_gpu={blue._want_onnx_gpu()} "
            f"max_concurrency={_max_concurrency()}",
            flush=True,
        )
        return _blue
    except Exception as exc:  # noqa: BLE001
        _load_error = f"{type(exc).__name__}: {exc}"
        print(f"[rp_handler] LOAD FAILED: {_load_error}", flush=True)
        traceback.print_exc()
        raise


def _resolve_voice(blue, name: Optional[str]):
    """Accept a voice name ("Daniel"), or the underlying file stem.

    The stem fallback keeps older clients working: the files are now named after
    the voice, so "libri_male_6209" resolves through VOICE_ALIASES to Daniel.
    Neither "Female" nor "Male" has a counterpart in the v2.5 bundle — Maya, the
    one "Female" used to reach, is withdrawn — so both raise with the current
    names rather than silently substituting a different speaker.
    """
    if not blue.VOICE_STYLES:
        raise RuntimeError("No voices installed in the image")
    if not name:
        key = next(iter(blue.VOICE_STYLES))
        return key, blue.VOICE_STYLES[key]

    raw = str(name).strip()
    if raw in blue.VOICE_STYLES:
        return raw, blue.VOICE_STYLES[raw]

    needle = raw.lower().replace("_", " ").replace("-", " ")
    for key, style in blue.VOICE_STYLES.items():
        if key.lower().replace("_", " ") == needle:
            return key, style
    for key, style in blue.VOICE_STYLES.items():
        if needle in key.lower():
            return key, style

    for key, path in blue.VOICES.items():
        base = os.path.splitext(os.path.basename(path))[0].lower()
        if base == raw.lower() or base.replace("_", "") == raw.lower().replace("_", ""):
            return key, blue.VOICE_STYLES[key]

    alias = blue.VOICE_ALIASES.get(raw.lower())
    if alias in blue.VOICE_STYLES:
        return alias, blue.VOICE_STYLES[alias]

    available = ", ".join(blue.VOICE_STYLES.keys())
    raise ValueError(f"Unknown voice {name!r}. Available: {available}")


def _synthesize_locked(
    blue,
    prepared: str,
    lang: str,
    style,
    steps: int,
    speed: float,
    cfg_scale: float,
) -> Tuple[np.ndarray, int]:
    with _synth_lock:
        wav, sr = blue.TTS.synthesize(
            prepared,
            lang=lang,
            style=style,
            total_step=steps,
            speed=speed,
            cfg_scale=cfg_scale,
        )
        wav = blue.normalize_generated_audio(np.asarray(wav).squeeze())
        return wav, int(sr)


def _encode_audio(wav: np.ndarray, sr: int, fmt: str) -> Tuple[str, str, int]:
    """Return (base64, format, encoded_bytes). Default OGG/Opus via ffmpeg."""
    fmt = (fmt or "ogg").strip().lower()
    if fmt in {"wav", "wave"}:
        buf = io.BytesIO()
        sf.write(buf, np.asarray(wav, dtype=np.float32), int(sr), format="WAV")
        raw = buf.getvalue()
        return base64.b64encode(raw).decode("ascii"), "wav", len(raw)

    # OGG Opus (voip-oriented) — much smaller than WAV for the client.
    with tempfile.TemporaryDirectory(prefix="blue_tts_") as td:
        wav_path = os.path.join(td, "in.wav")
        ogg_path = os.path.join(td, "out.ogg")
        sf.write(wav_path, np.asarray(wav, dtype=np.float32), int(sr), format="WAV")
        cmd = [
            "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
            "-i", wav_path,
            "-c:a", "libopus",
            "-application", "voip",
            "-b:a", "32k",
            "-ar", "48000",
            "-ac", "1",
            ogg_path,
        ]
        subprocess.run(cmd, check=True)
        raw = open(ogg_path, "rb").read()
        return base64.b64encode(raw).decode("ascii"), "ogg", len(raw)


def _process_job(job: Dict[str, Any]) -> Dict[str, Any]:
    t_total = time.time()
    blue = _ensure_blue()
    inp = job.get("input") or {}
    if not isinstance(inp, dict):
        return {"error": "input must be a JSON object"}

    text = inp.get("prompt") or inp.get("text") or inp.get("tts_text") or ""
    text = str(text).strip()
    if not text:
        return {
            "error": "Missing text. Pass input.prompt or input.text",
            "voices": list(blue.VOICE_STYLES.keys()),
            "languages": list(blue.AVAILABLE_LANGS),
        }

    lang = str(inp.get("lang") or inp.get("language") or "en")
    if blue._canonical_lang(lang) not in blue.AVAILABLE_LANGS:
        return {
            "error": f"Unsupported lang {lang!r}",
            "languages": list(blue.AVAILABLE_LANGS),
        }
    steps = int(inp.get("steps") or inp.get("total_step") or 8)
    speed = float(inp.get("speed") or 0.95)
    cfg_scale = float(inp.get("cfg_scale") or 4.0)
    fmt = str(inp.get("format") or inp.get("audio_format") or "ogg")
    voice_name, style = _resolve_voice(blue, inp.get("voice") or inp.get("speaker"))

    t0 = time.time()
    prepared = blue.prepare_text_for_synthesis(text, lang=lang)
    g2p_s = time.time() - t0

    t0 = time.time()
    wav, sr = _synthesize_locked(
        blue, prepared, lang, style, steps, speed, cfg_scale,
    )
    inference_s = time.time() - t0

    if wav.size == 0:
        return {
            "error": "No speakable text detected",
            "voice": voice_name,
            "lang": lang,
            "g2p_s": round(g2p_s, 3),
            "inference_s": round(inference_s, 3),
            "encode_s": 0.0,
            "total_s": round(time.time() - t_total, 3),
        }

    t0 = time.time()
    audio_b64, out_fmt, nbytes = _encode_audio(wav, sr, fmt)
    encode_s = time.time() - t0

    duration_s = float(len(wav) / sr) if sr else 0.0
    total_s = time.time() - t_total
    rtf = (inference_s / duration_s) if duration_s > 0 else 0.0

    return {
        "audio_base64": audio_b64,
        "sample_rate": int(sr) if out_fmt == "wav" else 48000,
        "format": out_fmt,
        "audio_bytes": nbytes,
        "duration_s": round(duration_s, 3),
        "g2p_s": round(g2p_s, 3),
        "inference_s": round(inference_s, 3),
        "encode_s": round(encode_s, 3),
        "total_s": round(total_s, 3),
        "rtf": round(rtf, 3),
        "voice": voice_name,
        "lang": lang,
        "steps": steps,
        "speed": speed,
        "providers": list(__import__("onnxruntime").get_available_providers()),
    }


async def handler(job: Dict[str, Any]) -> Dict[str, Any]:
    try:
        # Offload blocking synth/ffmpeg so concurrency_modifier can overlap jobs.
        return await asyncio.to_thread(_process_job, job)
    except Exception as exc:  # noqa: BLE001 — surface clean JSON to RunPod
        return {
            "error": f"{type(exc).__name__}: {exc}",
            "traceback": traceback.format_exc(),
        }


if __name__ == "__main__":
    print(
        f"[rp_handler] registering RunPod serverless worker "
        f"(MAX_CONCURRENCY={_max_concurrency()}) …",
        flush=True,
    )
    sys.stdout.flush()
    sys.stderr.flush()
    if os.environ.get("BLUE_PRELOAD", "").strip().lower() in {"1", "true", "yes", "on"}:
        try:
            _ensure_blue()
        except Exception:
            print("[rp_handler] preload failed; will retry on first job", flush=True)
    runpod.serverless.start({
        "handler": handler,
        "concurrency_modifier": _concurrency_modifier,
    })