196 строки
6.3 KiB
Python
196 строки
6.3 KiB
Python
"""
|
|
Minimal FastAPI server exposing Qwen3-TTS generation over HTTP.
|
|
Models are loaded once at startup on GPU if the installed PyTorch build
|
|
reports one available (e.g. a ROCm image), otherwise on 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"]
|
|
)
|
|
|
|
# Used whenever a request doesn't pass its own `instruct`. Qwen3-TTS has no
|
|
# numeric "speed"/"style" knob, so tone is steered with a natural-language
|
|
# instruction instead.
|
|
DEFAULT_INSTRUCT = os.environ.get(
|
|
"QWEN_TTS_DEFAULT_INSTRUCT",
|
|
"This is a voice message between two close friends on WhatsApp. "
|
|
"Speak in a casual, warm, friendly tone, not formal or professional-sounding.",
|
|
).strip()
|
|
|
|
# ROCm PyTorch builds expose the GPU through the same torch.cuda API as CUDA.
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
# bfloat16/float16 are not well supported on CPU; use full precision there.
|
|
DTYPE = torch.float32 if DEVICE == "cpu" else torch.bfloat16
|
|
|
|
app = FastAPI(title="Qwen3-TTS 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 {DEVICE}...")
|
|
models["customvoice"] = Qwen3TTSModel.from_pretrained(
|
|
CUSTOMVOICE_MODEL_NAME,
|
|
device_map=DEVICE,
|
|
dtype=DTYPE,
|
|
)
|
|
print("CustomVoice model loaded.")
|
|
|
|
if "voiceclone" in LOAD_MODES:
|
|
print(f"Loading Base (voice-clone) model {VOICECLONE_MODEL_NAME} on {DEVICE}...")
|
|
models["voiceclone"] = Qwen3TTSModel.from_pretrained(
|
|
VOICECLONE_MODEL_NAME,
|
|
device_map=DEVICE,
|
|
dtype=DTYPE,
|
|
)
|
|
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 DEFAULT_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")
|