564 lines
20 KiB
Python
564 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Qwen3-TTS FastAPI Server for UmbraAudio.
|
|
|
|
Provides REST endpoints for voice design and voice cloning.
|
|
Expected by umbraaudio.voices.qwen_tts.QwenTTSEngine.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
|
|
import anyio
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
from dataclasses import asdict
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
from fastapi.responses import Response
|
|
from pathlib import Path
|
|
from pydantic import BaseModel
|
|
from qwen_tts import Qwen3TTSModel
|
|
from qwen_tts.inference.qwen3_tts_model import VoiceClonePromptItem
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Directory where precomputed voice-clone prompts (the "8k matrix" ref_code +
|
|
# speaker embedding) are cached, one .pt file per speaker slug. Populated by
|
|
# POST /v1/voice-prompt and consumed by /v1/tts via the prompt_slug field so
|
|
# every utterance for a speaker reuses the exact same voice identity.
|
|
PROMPT_DIR = Path(
|
|
os.getenv(
|
|
"QWEN_PROMPT_DIR",
|
|
"/Volumes/1TB_APL_Storage/UmbraAudio/DealHunterShip_4.0e_whole-to-audio/.umbra/voice_prompts",
|
|
)
|
|
)
|
|
|
|
app = FastAPI(title="Qwen3-TTS Server", version="0.1.0")
|
|
|
|
|
|
async def _watchdog_loop() -> None:
|
|
"""Background task: if a generation has been stuck past MAX_GEN, restart."""
|
|
while True:
|
|
await asyncio.sleep(min(_MAX_GEN, 30))
|
|
s = _gen_start_at
|
|
if s is not None and (time.monotonic() - s) > _MAX_GEN:
|
|
_restart_server(
|
|
f"generation hung >{_MAX_GEN:.0f}s (poisoned CUDA context?)"
|
|
)
|
|
await asyncio.sleep(3600) # never reached after execv
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def _start_watchdog() -> None:
|
|
app.state._wd_task = asyncio.create_task(_watchdog_loop())
|
|
|
|
# Rainbow Passage — used as default ref_text for ICL voice cloning mode.
|
|
# All voice refs were generated with this text via Edge-TTS.
|
|
VOICE_REFERENCE_TEXT = (
|
|
"When the sunlight strikes raindrops in the air, they act as a prism and form a rainbow. "
|
|
"The rainbow is a division of white light into many beautiful colors. "
|
|
"These take the shape of a long round arch, with its path high above, "
|
|
"and its two ends apparently beyond the horizon. "
|
|
"There is, according to legend, a boiling pot of gold at one end. "
|
|
"People look, but no one ever finds it. "
|
|
"When a man looks for something beyond his reach, "
|
|
"his friends say he is looking for the pot of gold at the end of the rainbow. "
|
|
"When the sunlight strikes raindrops in the air, they act as a prism and form a rainbow. "
|
|
"The rainbow is a division of white light into many beautiful colors. "
|
|
"These take the shape of a long round arch, with its path high above, "
|
|
"and its two ends apparently beyond the horizon. "
|
|
"There is, according to legend, a boiling pot of gold at one end. "
|
|
"People look, but no one ever finds it. "
|
|
"When a man looks for something beyond his reach, "
|
|
"his friends say he is looking for the pot of gold at the end of the rainbow."
|
|
)
|
|
|
|
# Global model state
|
|
_vd_model: Qwen3TTSModel | None = None
|
|
_base_model: Qwen3TTSModel | None = None
|
|
_current_model: str | None = None # "voicedesign" or "base"
|
|
|
|
# Serializes all model access. Model loads and generations are blocking CPU/GPU
|
|
# calls; they run in worker threads (anyio.to_thread) under this lock so the
|
|
# event loop stays responsive and concurrent /v1/tts requests cannot wedge it.
|
|
_model_lock = asyncio.Lock()
|
|
|
|
# Hang watchdog. On CUDA, a kernel error (e.g. OOM) can poison the context and
|
|
# wedge the generation thread forever while the event loop stays responsive
|
|
# (health works, every /v1/tts hangs). The only safe recovery is to restart the
|
|
# process. We track when a generation started and, if it exceeds MAX_GEN_SECONDS,
|
|
# we relaunch ourselves via os.execv to get a clean CUDA context.
|
|
_MAX_GEN = float(os.getenv("QWEN_MAX_GEN_S", "1800"))
|
|
_gen_start_at: float | None = None # time.monotonic() when a generation began
|
|
|
|
|
|
def _mark_gen_started() -> None:
|
|
global _gen_start_at
|
|
_gen_start_at = time.monotonic()
|
|
|
|
|
|
def _mark_gen_done() -> None:
|
|
global _gen_start_at
|
|
_gen_start_at = None
|
|
|
|
|
|
def _restart_server(reason: str) -> None:
|
|
"""Relaunch the current process (os.execv) to clear a wedged GPU context."""
|
|
import sys as _sys
|
|
|
|
logger.warning("Restarting TTS server: %s", reason)
|
|
try:
|
|
for s in (_sys.stdout, _sys.stderr):
|
|
try:
|
|
s.flush()
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
try:
|
|
_sys.stdout = open(os.devnull, "w")
|
|
_sys.stderr = open(os.devnull, "w")
|
|
_sys.stdin = open(os.devnull)
|
|
os.execv(
|
|
_sys.executable,
|
|
[_sys.executable, os.path.abspath(__file__), *list(_sys.argv[1:])],
|
|
)
|
|
except Exception as e: # pragma: no cover
|
|
logger.error("os.execv failed: %s", e)
|
|
raise SystemExit(1)
|
|
|
|
# Model paths from environment
|
|
VOICEDESIGN_PATH = os.getenv(
|
|
"QWEN_VOICEDESIGN_MODEL",
|
|
"/Volumes/1TB_APL_Storage/UmbraAudio/models/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
|
|
)
|
|
BASE_PATH = os.getenv(
|
|
"QWEN_BASE_MODEL",
|
|
"/Volumes/1TB_APL_Storage/UmbraAudio/models/Qwen3-TTS-12Hz-1.7B-Base",
|
|
)
|
|
DEVICE = os.getenv("QWEN_DEVICE", "mps")
|
|
DTYPE = torch.bfloat16 if DEVICE in ("cuda", "mps") else torch.float32
|
|
# Attention implementation. "eager" is very slow for Qwen3-TTS autoregressive
|
|
# codec decoding but is the only one that reliably runs create_voice_clone_prompt
|
|
# (prompt extraction). "sdpa" (built into torch) is far faster for generation.
|
|
# Runtime-switchable via POST /attn (reloads the model); this is read at model
|
|
# load time, so switching reloads the active model with the new implementation.
|
|
_ATTN_IMPL: str = os.getenv("QWEN_ATTN", "sdpa")
|
|
|
|
|
|
class LoadRequest(BaseModel):
|
|
model_type: str # "voicedesign" or "base"
|
|
|
|
|
|
class GenerateRequest(BaseModel):
|
|
text: str
|
|
sample_rate: int = 24000
|
|
voice_description: str | None = None
|
|
|
|
|
|
def _load_model(model_type: str) -> Qwen3TTSModel:
|
|
"""Load a model, unloading the other if needed."""
|
|
global _vd_model, _base_model, _current_model
|
|
|
|
if _current_model == model_type:
|
|
return _vd_model if model_type == "voicedesign" else _base_model
|
|
|
|
# Unload current
|
|
if _current_model == "voicedesign" and _vd_model is not None:
|
|
del _vd_model
|
|
_vd_model = None
|
|
elif _current_model == "base" and _base_model is not None:
|
|
del _base_model
|
|
_base_model = None
|
|
|
|
if torch.mps.is_available() or torch.cuda.is_available():
|
|
torch.mps.empty_cache() if torch.mps.is_available() else torch.cuda.empty_cache()
|
|
|
|
# Load requested
|
|
path = VOICEDESIGN_PATH if model_type == "voicedesign" else BASE_PATH
|
|
logger.info(f"Loading {model_type} model from {path}")
|
|
|
|
model = Qwen3TTSModel.from_pretrained(
|
|
path,
|
|
device_map=DEVICE,
|
|
dtype=DTYPE,
|
|
attn_implementation=_ATTN_IMPL,
|
|
)
|
|
|
|
if model_type == "voicedesign":
|
|
_vd_model = model
|
|
else:
|
|
_base_model = model
|
|
|
|
_current_model = model_type
|
|
logger.info(f"{model_type} model loaded successfully")
|
|
return model
|
|
|
|
|
|
def _prompt_path(slug: str) -> Path:
|
|
safe = Path(slug).name
|
|
PROMPT_DIR.mkdir(parents=True, exist_ok=True)
|
|
return PROMPT_DIR / f"{safe}.pt"
|
|
|
|
|
|
def _save_prompt(slug: str, items: list[VoiceClonePromptItem]) -> Path:
|
|
path = _prompt_path(slug)
|
|
payload = {"items": [asdict(it) for it in items]}
|
|
torch.save(payload, str(path))
|
|
logger.info("Saved ICL voice prompt for %s to %s", slug, path)
|
|
return path
|
|
|
|
|
|
def _load_prompt_items(slug: str) -> list[VoiceClonePromptItem] | None:
|
|
path = _prompt_path(slug)
|
|
if not path.is_file():
|
|
return None
|
|
try:
|
|
payload = torch.load(str(path), map_location="cpu", weights_only=True)
|
|
except Exception as e:
|
|
logger.warning("Failed to load prompt %s: %s", slug, e)
|
|
return None
|
|
raw_items = payload.get("items", [])
|
|
items: list[VoiceClonePromptItem] = []
|
|
for d in raw_items:
|
|
ref_code = d.get("ref_code")
|
|
if ref_code is not None and not torch.is_tensor(ref_code):
|
|
ref_code = torch.tensor(ref_code)
|
|
ref_spk = d.get("ref_spk_embedding")
|
|
if ref_spk is None:
|
|
logger.warning("Prompt %s missing ref_spk_embedding", slug)
|
|
return None
|
|
if not torch.is_tensor(ref_spk):
|
|
ref_spk = torch.tensor(ref_spk)
|
|
items.append(
|
|
VoiceClonePromptItem(
|
|
ref_code=ref_code,
|
|
ref_spk_embedding=ref_spk,
|
|
x_vector_only_mode=bool(d.get("x_vector_only_mode", False)),
|
|
icl_mode=bool(d.get("icl_mode", not bool(d.get("x_vector_only_mode", False)))),
|
|
ref_text=d.get("ref_text"),
|
|
)
|
|
)
|
|
return items
|
|
|
|
|
|
@app.post("/v1/voice-prompt")
|
|
async def build_voice_prompt(
|
|
slug: str = Form(...),
|
|
ref_text: str | None = Form(None),
|
|
ref_audio: UploadFile | None = File(None),
|
|
x_vector_only: int = Form(0),
|
|
force: int = Form(0),
|
|
):
|
|
"""Build and cache a speaker's vmatrix (digital voice) once.
|
|
|
|
The vmatrix is the set of numbers inside Qwen3-TTS that IS the voice. Two capture
|
|
modes (a speaker's cached prompt is reused on every subsequent utterance via
|
|
prompt_slug):
|
|
|
|
* x_vector_only=0 (DEFAULT, canonical, "full vmatrix"): extracts the FULL
|
|
ICL codec-token matrix (ref_code, the "8k matrix") + speaker embedding.
|
|
Generation conditions on the exact reference identity, so every utterance
|
|
is faithful AND identical in identity across the whole pipeline.
|
|
* x_vector_only=1: captures ONLY the coarse speaker embedding (ref_code=None).
|
|
Legacy/experimental - caused the historical "voice changed between
|
|
chapters" inconsistency. Not used for production consistency.
|
|
|
|
Idempotent: if `<slug>.pt` is already cached it returns without rebuilding,
|
|
unless `force=1` is passed (use force to switch an existing prompt mode, e.g.
|
|
after an approved reference WAV is replaced).
|
|
"""
|
|
path = _prompt_path(slug)
|
|
if path.exists() and not force:
|
|
return {"ok": True, "slug": slug, "cached": True, "path": str(path)}
|
|
|
|
if ref_audio is None:
|
|
raise HTTPException(400, "ref_audio is required to build a voice prompt")
|
|
ref_bytes = await ref_audio.read()
|
|
|
|
ref_text = (ref_text or "").strip() or VOICE_REFERENCE_TEXT
|
|
xvec = bool(x_vector_only)
|
|
|
|
def _build() -> list[VoiceClonePromptItem]:
|
|
model = _load_model("base")
|
|
ref_audio_array, ref_sr = sf.read(io.BytesIO(ref_bytes), dtype="float32")
|
|
if ref_sr != 24000:
|
|
import librosa
|
|
ref_audio_array = librosa.resample(
|
|
ref_audio_array, orig_sr=ref_sr, target_sr=24000
|
|
)
|
|
items = model.create_voice_clone_prompt(
|
|
ref_audio=(ref_audio_array, 24000),
|
|
ref_text=ref_text,
|
|
x_vector_only_mode=xvec,
|
|
)
|
|
return items
|
|
|
|
_mark_gen_started()
|
|
async with _model_lock:
|
|
items = await anyio.to_thread.run_sync(_build)
|
|
_mark_gen_done()
|
|
|
|
_save_prompt(slug, items)
|
|
return {"ok": True, "slug": slug, "path": str(_prompt_path(slug))}
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {
|
|
"status": "ok",
|
|
"voicedesign_loaded": _vd_model is not None,
|
|
"base_loaded": _base_model is not None,
|
|
"current_model": _current_model,
|
|
"device": DEVICE,
|
|
}
|
|
|
|
|
|
@app.post("/load")
|
|
async def load_model(request: LoadRequest):
|
|
if request.model_type not in ("voicedesign", "base"):
|
|
raise HTTPException(400, "model_type must be 'voicedesign' or 'base'")
|
|
|
|
try:
|
|
async with _model_lock:
|
|
await anyio.to_thread.run_sync(_load_model, request.model_type)
|
|
return {"ok": True, "model": request.model_type}
|
|
except Exception as e:
|
|
logger.exception("Failed to load model")
|
|
raise HTTPException(500, f"Failed to load model: {e}") from e
|
|
|
|
|
|
def _unload_model(model_type: str) -> None:
|
|
global _vd_model, _base_model, _current_model
|
|
|
|
if model_type in ("voicedesign", "all") and _vd_model is not None:
|
|
del _vd_model
|
|
_vd_model = None
|
|
if model_type in ("base", "all") and _base_model is not None:
|
|
del _base_model
|
|
_base_model = None
|
|
|
|
if model_type == "all" or _current_model == model_type:
|
|
_current_model = None
|
|
|
|
if torch.mps.is_available() or torch.cuda.is_available():
|
|
torch.mps.empty_cache() if torch.mps.is_available() else torch.cuda.empty_cache()
|
|
|
|
|
|
@app.post("/unload")
|
|
async def unload_model(request: LoadRequest):
|
|
if request.model_type not in ("voicedesign", "base", "all"):
|
|
raise HTTPException(400, "model_type must be 'voicedesign', 'base', or 'all'")
|
|
|
|
async with _model_lock:
|
|
await anyio.to_thread.run_sync(_unload_model, request.model_type)
|
|
|
|
return {"ok": True, "unloaded": request.model_type}
|
|
|
|
|
|
class AttnRequest(BaseModel):
|
|
impl: str # "eager", "sdpa", "fluent", "flash_attention_2"
|
|
|
|
|
|
@app.get("/attn")
|
|
async def get_attn():
|
|
return {"ok": True, "impl": _ATTN_IMPL}
|
|
|
|
|
|
@app.post("/attn")
|
|
async def set_attn(request: AttnRequest):
|
|
"""Switch the attention implementation and reload the active model.
|
|
|
|
The attention implementation is fixed at model-load time, so this unloads
|
|
the current model(s) and reloads `base` with the new implementation. Callers
|
|
should switch to ``eager`` for prompt extraction (Phase 1) and ``sdpa`` for
|
|
generation (Phase 2) without a process restart.
|
|
"""
|
|
global _ATTN_IMPL
|
|
impl = (request.impl or "").strip().lower()
|
|
if impl not in ("eager", "sdpa", "fluent", "flash_attention_2"):
|
|
raise HTTPException(400, "impl must be eager|sdpa|fluent|flash_attention_2")
|
|
|
|
async with _model_lock:
|
|
try:
|
|
await anyio.to_thread.run_sync(_unload_model, "all")
|
|
_ATTN_IMPL = impl
|
|
# reload so subsequent use has the new attn; base is the main model
|
|
await anyio.to_thread.run_sync(_load_model, "base")
|
|
except Exception as e:
|
|
logger.exception("Failed to switch attn to %s", impl)
|
|
raise HTTPException(500, f"Failed to switch attn: {e}") from e
|
|
|
|
return {"ok": True, "impl": _ATTN_IMPL}
|
|
|
|
|
|
@app.post("/v1/tts")
|
|
async def generate_tts(
|
|
text: str = Form(...),
|
|
sample_rate: int = Form(24000),
|
|
# Was hardcoded to "english" in every branch, so German and French text was
|
|
# synthesised with English phonology and prosody - an English-accented
|
|
# foreign voice. Defaults to "english" so the existing English pipeline is
|
|
# byte-for-byte unchanged; multilingual callers pass their own language.
|
|
# "Auto" is also accepted and lets the model detect it.
|
|
language: str = Form("english"),
|
|
voice_description: str | None = Form(None),
|
|
ref_audio: UploadFile | None = File(None),
|
|
ref_text: str | None = Form(None),
|
|
speaker_embedding: str | None = Form(None),
|
|
prompt_slug: str | None = Form(None),
|
|
temperature: float = Form(0.9),
|
|
do_sample: bool = Form(True),
|
|
seed: int | None = Form(None),
|
|
):
|
|
"""Generate TTS audio.
|
|
|
|
- voice_description + text -> VoiceDesign generation
|
|
- ref_audio + text + ref_text -> VoiceClone generation
|
|
- speaker_embedding + text -> Direct embedding generation (fastest, most consistent)
|
|
- prompt_slug -> ICL generation from cached voice-clone prompt (most faithful &
|
|
consistent; the "8k matrix" built by POST /v1/voice-prompt)
|
|
|
|
prompt_slug takes priority: if the cached prompt exists it is used, giving
|
|
the speaker the identical reference voice on every utterance.
|
|
|
|
voice_description also combines with a clone/ICL prompt: when ref_audio or
|
|
prompt_slug is given, the instruction text is prepended as style direction
|
|
(emotion, pace, delivery) while the cloned speaker identity is preserved.
|
|
"""
|
|
gen_kwargs: dict = {"temperature": temperature, "do_sample": do_sample}
|
|
|
|
if seed is not None:
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
_mark_gen_started()
|
|
async with _model_lock:
|
|
if prompt_slug:
|
|
items = _load_prompt_items(prompt_slug)
|
|
if items is None:
|
|
raise HTTPException(
|
|
404, f"No cached voice prompt for slug '{prompt_slug}'"
|
|
)
|
|
|
|
def _icl() -> tuple[np.ndarray, int]:
|
|
model = _load_model("base")
|
|
kw = dict(gen_kwargs)
|
|
if voice_description:
|
|
kw["instruct_ids"] = [
|
|
model._tokenize_texts(
|
|
[model._build_instruct_text(voice_description)]
|
|
)[0]
|
|
]
|
|
wavs, sr = model.generate_voice_clone(
|
|
text=text,
|
|
voice_clone_prompt=items,
|
|
language=language,
|
|
**kw,
|
|
)
|
|
return wavs, sr
|
|
|
|
wavs, sr = await anyio.to_thread.run_sync(_icl)
|
|
elif speaker_embedding is not None:
|
|
emb_list = json.loads(speaker_embedding)
|
|
emb_np = np.array(emb_list, dtype=np.float32)
|
|
spk_tensor = torch.from_numpy(emb_np)
|
|
|
|
def _embed() -> tuple[np.ndarray, int]:
|
|
model = _load_model("base")
|
|
device = next(model.model.parameters()).device
|
|
dtype = next(model.model.parameters()).dtype
|
|
spk = spk_tensor.to(device).to(dtype)
|
|
|
|
wavs, sr = model.generate_voice_clone(
|
|
text=text,
|
|
voice_clone_prompt={
|
|
"ref_code": [None],
|
|
"ref_spk_embedding": [spk],
|
|
"x_vector_only_mode": [True],
|
|
"icl_mode": [False],
|
|
},
|
|
language=language,
|
|
**gen_kwargs,
|
|
)
|
|
return wavs, sr
|
|
|
|
wavs, sr = await anyio.to_thread.run_sync(_embed)
|
|
elif ref_audio is not None:
|
|
ref_bytes = await ref_audio.read()
|
|
|
|
def _clone() -> tuple[np.ndarray, int]:
|
|
model = _load_model("base")
|
|
ref_audio_array, ref_sr = sf.read(io.BytesIO(ref_bytes), dtype="float32")
|
|
|
|
if ref_sr != sample_rate:
|
|
import librosa
|
|
|
|
ref_audio_array = librosa.resample(
|
|
ref_audio_array, orig_sr=ref_sr, target_sr=sample_rate
|
|
)
|
|
|
|
kw = dict(gen_kwargs)
|
|
if voice_description:
|
|
kw["instruct_ids"] = [
|
|
model._tokenize_texts(
|
|
[model._build_instruct_text(voice_description)]
|
|
)[0]
|
|
]
|
|
wavs, sr = model.generate_voice_clone(
|
|
text=text,
|
|
ref_audio=(ref_audio_array, sample_rate),
|
|
ref_text=None,
|
|
x_vector_only_mode=True,
|
|
language=language,
|
|
**kw,
|
|
)
|
|
return wavs, sr
|
|
|
|
wavs, sr = await anyio.to_thread.run_sync(_clone)
|
|
elif voice_description:
|
|
def _design() -> tuple[np.ndarray, int]:
|
|
model = _load_model("voicedesign")
|
|
wavs, sr = model.generate_voice_design(
|
|
text=text,
|
|
instruct=voice_description,
|
|
language=language,
|
|
**gen_kwargs,
|
|
)
|
|
return wavs, sr
|
|
|
|
wavs, sr = await anyio.to_thread.run_sync(_design)
|
|
else:
|
|
raise HTTPException(400, "Either voice_description or ref_audio required")
|
|
|
|
_mark_gen_done()
|
|
if not wavs:
|
|
raise HTTPException(500, "Model returned no audio")
|
|
|
|
# Return WAV as binary
|
|
wav_bytes = io.BytesIO()
|
|
sf.write(wav_bytes, wavs[0], sr, format="WAV")
|
|
wav_bytes.seek(0)
|
|
|
|
return Response(
|
|
content=wav_bytes.read(),
|
|
media_type="audio/wav",
|
|
headers={"Content-Disposition": "attachment; filename=output.wav"},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8881)
|