bdi_podman_serverconf/containers/qwen-tts/server.py

"""
Minimal FastAPI server exposing Qwen3-TTS generation over HTTP.
CPU-only: models are loaded once at startup with device_map="cpu".

Loads BOTH Qwen3-TTS model families at once, so a single container/process
can serve:
- /speech       -> CustomVoice model, 9 built-in preset speakers
                   (generate_custom_voice)
- /speech/clone -> Base model, clone a voice from a short reference audio
                   file (generate_voice_clone), 10-20s recommended,
                   3s min, 60s max

Set QWEN_TTS_LOAD=customvoice,voiceclone (default: both) to control which
model(s) are loaded, useful if you want a lighter single-purpose instance
for memory-constrained setups. Requesting an endpoint whose model wasn't
loaded returns a clear 400 error instead of a generic failure.
"""
import io
import os
import tempfile

import soundfile as sf
import torch
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from qwen_tts import Qwen3TTSModel

_DEFAULT_MODELS = {
    "customvoice": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
    "voiceclone": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
}

_load_env = os.environ.get("QWEN_TTS_LOAD", "customvoice,voiceclone")
LOAD_MODES = [m.strip() for m in _load_env.split(",") if m.strip()]
for m in LOAD_MODES:
    if m not in _DEFAULT_MODELS:
        raise ValueError(
            f"QWEN_TTS_LOAD entries must be 'customvoice' and/or 'voiceclone', got: {m}"
        )

CUSTOMVOICE_MODEL_NAME = os.environ.get(
    "QWEN_TTS_CUSTOMVOICE_MODEL", _DEFAULT_MODELS["customvoice"]
)
VOICECLONE_MODEL_NAME = os.environ.get(
    "QWEN_TTS_VOICECLONE_MODEL", _DEFAULT_MODELS["voiceclone"]
)

app = FastAPI(title="Qwen3-TTS CPU server")

# Populated at startup, keyed by mode ("customvoice" / "voiceclone")
models: dict[str, Qwen3TTSModel] = {}


class SpeechRequest(BaseModel):
    text: str
    language: str = "Auto"
    speaker: str = "Vivian"
    instruct: str = ""


def _require_model(mode: str) -> Qwen3TTSModel:
    model = models.get(mode)
    if model is None:
        raise HTTPException(
            status_code=400,
            detail=f"The '{mode}' model is not loaded on this instance "
                   f"(QWEN_TTS_LOAD={','.join(LOAD_MODES)}). "
                   f"Restart the container with QWEN_TTS_LOAD including "
                   f"'{mode}' to use this endpoint.",
        )
    return model


@app.on_event("startup")
def load_models():
    if "customvoice" in LOAD_MODES:
        print(f"Loading CustomVoice model {CUSTOMVOICE_MODEL_NAME} on CPU...")
        models["customvoice"] = Qwen3TTSModel.from_pretrained(
            CUSTOMVOICE_MODEL_NAME,
            device_map="cpu",
            dtype=torch.float32,  # bfloat16/float16 are not well supported on CPU
        )
        print("CustomVoice model loaded.")

    if "voiceclone" in LOAD_MODES:
        print(f"Loading Base (voice-clone) model {VOICECLONE_MODEL_NAME} on CPU...")
        models["voiceclone"] = Qwen3TTSModel.from_pretrained(
            VOICECLONE_MODEL_NAME,
            device_map="cpu",
            dtype=torch.float32,
        )
        print("Base (voice-clone) model loaded.")


@app.get("/health")
def health():
    return {
        "status": "ok",
        "loaded_models": {
            mode: (CUSTOMVOICE_MODEL_NAME if mode == "customvoice" else VOICECLONE_MODEL_NAME)
            for mode in models
        },
    }


@app.get("/speakers")
def speakers():
    model = _require_model("customvoice")
    return {"speakers": model.get_supported_speakers()}


@app.post("/speech")
def speech(req: SpeechRequest):
    model = _require_model("customvoice")

    wavs, sr = model.generate_custom_voice(
        text=req.text,
        language=req.language,
        speaker=req.speaker,
        instruct=req.instruct or None,
    )

    buf = io.BytesIO()
    sf.write(buf, wavs[0], sr, format="WAV")
    buf.seek(0)
    return StreamingResponse(buf, media_type="audio/wav")


@app.post("/speech/clone")
def speech_clone(
    text: str = Form(...),
    language: str = Form("Auto"),
    ref_text: str = Form(""),
    x_vector_only_mode: bool = Form(False),
    ref_audio: UploadFile = File(...),
):
    """
    Clone a voice from a short reference audio file and synthesize `text`
    with it.

    ref_audio: 10-20s of clear, continuous speech recommended (min 3s, max 60s).
    ref_text: transcript of what is said in ref_audio. Required unless
              x_vector_only_mode=true.
    x_vector_only_mode: if true, clones only from the speaker embedding
              (x-vector) without needing ref_text. Faster, but generally
              lower fidelity than the default in-context-learning mode.
    """
    model = _require_model("voiceclone")

    if not x_vector_only_mode and not ref_text:
        raise HTTPException(
            status_code=422,
            detail="ref_text is required unless x_vector_only_mode=true.",
        )

    ref_bytes = ref_audio.file.read()

    # generate_voice_clone does not accept an in-memory BytesIO object,
    # it expects a file path. Write the upload to a temp file instead.
    suffix = os.path.splitext(ref_audio.filename or "")[1] or ".wav"
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(ref_bytes)
        tmp_path = tmp.name

    try:
        wavs, sr = model.generate_voice_clone(
            text=text,
            language=language,
            ref_audio=tmp_path,
            ref_text=ref_text or None,
            x_vector_only_mode=x_vector_only_mode,
        )
    finally:
        os.remove(tmp_path)

    buf = io.BytesIO()
    sf.write(buf, wavs[0], sr, format="WAV")
    buf.seek(0)
    return StreamingResponse(buf, media_type="audio/wav")