handler.py

"""HF Inference Endpoint handler — Qwen3-TTS-12Hz-1.7B-CustomVoice.

Runs on an HF Inference Endpoint with a GPU instance (L4 24GB is plenty for a
1.9B TTS model). Loads Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice via the official
`qwen-tts` package, which wires the talker LM + code predictor + the separate
12Hz speech tokenizer (vocoder) into one `generate_custom_voice` call.

Why an endpoint and not fal/ElevenLabs: this is OUR infra — unfiltered, fixed
per-hour cost (scale-to-zero when idle), with natural-language style control
(`instruct`) and best-effort zero-shot cloning that hosted TTS doesn't expose
cheaply.

Request shape (POST to /):
    {
        "inputs": "the text to speak",
        "parameters": {
            "speaker": "Ryan",          # built-in voice; see SPEAKERS below
            "language": "English",      # or "Auto" for auto-detect
            "instruct": "calm, warm",   # optional natural-language style control
            "max_new_tokens": 2048,
            "temperature": 0.9,
            "top_p": 0.9,
            "top_k": 50,
            "seed": 42                  # optional, for reproducibility
        }
    }

Zero-shot clone (omit `speaker`, pass a reference instead). Best-effort on the
CustomVoice checkpoint; the *-Base model specializes in cloning:
    {
        "inputs": "text in the cloned voice",
        "parameters": {
            "language": "English",
            "ref_audio_b64": "<base64 of a wav/mp3 reference clip>",
            "ref_text": "exact transcript of the reference clip"
        }
    }

Response shape (audio is always 24kHz mono WAV; transcode client-side if needed):
    [{"audio": "<base64 wav>", "format": "wav", "sample_rate": 24000,
      "duration_s": <float>, "speaker": "...", "language": "..."}]

The model is Apache-2.0 and ungated — no license-acceptance or token gymnastics
needed (unlike the Flux endpoint). HF_TOKEN is read if present, purely to raise
Hub download rate limits during cold start.
"""
from __future__ import annotations

import base64
import io
import os
import sys
from typing import Any

import soundfile as sf
import torch

# The qwen_tts package is VENDORED next to this file (./qwen_tts/) rather than
# pip-installed — see requirements.txt for why (gradio→starlette would crash
# HF's serving wrapper). Ensure this directory is importable regardless of the
# toolkit's CWD so `from qwen_tts import ...` resolves to the vendored copy.
_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

# Built-in speakers shipped with Qwen3-TTS-12Hz-1.7B-CustomVoice, with their
# native language. Each can speak any supported language, but native-language
# pairing gives the best quality (per the model card).
SPEAKERS = {
    "Vivian":   "Chinese",   # bright, slightly edgy young female
    "Serena":   "Chinese",   # warm, gentle young female
    "Uncle_Fu": "Chinese",   # seasoned male, low mellow timbre
    "Dylan":    "Chinese",   # youthful Beijing male
    "Eric":     "Chinese",   # lively Chengdu (Sichuan) male
    "Ryan":     "English",   # dynamic male, strong rhythm
    "Aiden":    "English",   # sunny American male, clear midrange
    "Ono_Anna": "Japanese",  # playful Japanese female
    "Sohee":    "Korean",    # warm Korean female, rich emotion
}
DEFAULT_SPEAKER = "Ryan"     # English narration default


class EndpointHandler:
    def __init__(self, path: str = ""):
        """Cold-start: load the CustomVoice model + 12Hz vocoder to VRAM.

        We pull the model by Hub id (not from the local `path`) — the `qwen-tts`
        package resolves both the 1.9B model and the 12Hz speech tokenizer it
        needs, the way it's designed to. Mirrors the Flux handler, which
        downloads its base model by id at init.
        """
        from qwen_tts import Qwen3TTSModel

        token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_API_TOKEN")
        if token:
            os.environ.setdefault("HF_TOKEN", token)  # higher Hub rate limits

        model_id = os.environ.get(
            "QWEN_TTS_MODEL_ID", "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
        )

        # sdpa, NOT flash_attention_2: flash-attn is an optional extra that
        # compiles from source and routinely breaks endpoint builds. sdpa is
        # built into torch and fast enough for a 1.9B model.
        self.model = Qwen3TTSModel.from_pretrained(
            model_id,
            device_map="cuda:0",
            dtype=torch.bfloat16,
            attn_implementation="sdpa",
        )

        # Cache the speaker list the loaded checkpoint actually supports so we
        # can validate requests instead of failing deep in generate(). Falls
        # back to our static map if the API differs.
        try:
            self.supported_speakers = set(self.model.get_supported_speakers())
        except Exception:  # noqa: BLE001
            self.supported_speakers = set(SPEAKERS)

    def __call__(self, data: dict[str, Any]) -> list[dict[str, Any]]:
        text = data.get("inputs")
        if not text or not isinstance(text, str):
            return [{"error": "missing 'inputs' (text to speak) in request body"}]

        params = data.get("parameters") or {}
        language = params.get("language") or "Auto"
        instruct = params.get("instruct") or ""

        # Generation knobs (pass-through to HF generate). Defaults come from the
        # checkpoint's generation_config.json; we only override when asked.
        gen_kwargs: dict[str, Any] = {}
        for key in ("max_new_tokens", "top_k"):
            if key in params:
                gen_kwargs[key] = int(params[key])
        for key in ("temperature", "top_p", "repetition_penalty"):
            if key in params:
                gen_kwargs[key] = float(params[key])

        seed = params.get("seed")
        if seed is not None:
            torch.manual_seed(int(seed))

        ref_audio_b64 = params.get("ref_audio_b64")

        try:
            with torch.inference_mode():
                if ref_audio_b64:
                    ref_text = params.get("ref_text") or ""
                    ref_bytes = base64.b64decode(ref_audio_b64)
                    ref_wav, ref_sr = sf.read(io.BytesIO(ref_bytes), dtype="float32")
                    wavs, sr = self.model.generate_voice_clone(
                        text=text, language=language,
                        ref_audio=(ref_wav, ref_sr), ref_text=ref_text,
                        **gen_kwargs,
                    )
                    chosen_speaker = "<clone>"
                else:
                    speaker = params.get("speaker") or DEFAULT_SPEAKER
                    if self.supported_speakers and speaker not in self.supported_speakers:
                        return [{
                            "error": f"unknown speaker '{speaker}'. "
                                     f"supported: {sorted(self.supported_speakers)}"
                        }]
                    wavs, sr = self.model.generate_custom_voice(
                        text=text, language=language, speaker=speaker,
                        instruct=instruct, **gen_kwargs,
                    )
                    chosen_speaker = speaker
        except Exception as e:  # noqa: BLE001 — surface generate errors to caller
            return [{"error": f"{type(e).__name__}: {e}"}]

        # qwen-tts returns a batch; we sent one line, so take wavs[0].
        wav = wavs[0]
        if hasattr(wav, "detach"):  # tensor -> numpy float32 in [-1, 1]
            wav = wav.detach().to(torch.float32).cpu().numpy()

        buf = io.BytesIO()
        sf.write(buf, wav, int(sr), format="WAV")
        audio_bytes = buf.getvalue()
        duration_s = float(len(wav)) / float(sr) if sr else 0.0
        return [{
            "audio": base64.b64encode(audio_bytes).decode("ascii"),
            "format": "wav",
            "sample_rate": int(sr),
            "duration_s": round(duration_s, 3),
            "speaker": chosen_speaker,
            "language": language,
        }]