c86e323d7a
Quality-first multi-pass audiobook generation from English manuscripts. Includes: - umbraaudio/ core package (pipeline, config, voices, web review UI) - scripts/ CLI utilities + model download script - tts_server/ remote TTS server - specs/ design docs, tests/, pyproject.toml + uv.lock - Qwen3-TTS git submodule (pinned 022e286) with local demo.py UI patch saved at scripts/qwen3_tts_demo_umbra_patch.patch - .env with runtime config/secrets Excluded (see .gitignore): models/ (downloaded via scripts/download_models.sh), output/, logs/, .venv, .umbra runtime data, and book projects (DealHunterShip_*), which are committed separately.
536 lines
19 KiB
Python
536 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""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 digital voice (a set of numbers) once.
|
|
|
|
Two capture modes (pick per pipeline; a speaker's cached prompt is used as-is
|
|
on every subsequent utterance):
|
|
|
|
* x_vector_only=0 (default, ICL/"8k matrix"): extracts the FULL codec-token
|
|
matrix (ref_code) + speaker embedding. Together those let the generator
|
|
re-create the exact reference recording -- and every utterance therefore
|
|
re-speaks the reference first, then trims it (slow + leaky).
|
|
* x_vector_only=1: captures ONLY the speaker embedding (the digital voice
|
|
variables). ref_code is None. Generation then feeds just those numbers +
|
|
text into the speech generator -- no reference re-hash, no clone re-run,
|
|
fast and output is exactly the text.
|
|
|
|
Idempotent: if `<slug>.pt` is already cached it returns without rebuilding,
|
|
unless `force=1` is passed (use force to switch an existing prompt mode).
|
|
"""
|
|
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),
|
|
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.
|
|
"""
|
|
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")
|
|
wavs, sr = model.generate_voice_clone(
|
|
text=text,
|
|
voice_clone_prompt=items,
|
|
language="english",
|
|
**gen_kwargs,
|
|
)
|
|
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="english",
|
|
**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
|
|
)
|
|
|
|
wavs, sr = model.generate_voice_clone(
|
|
text=text,
|
|
ref_audio=(ref_audio_array, sample_rate),
|
|
ref_text=None,
|
|
x_vector_only_mode=True,
|
|
language="english",
|
|
**gen_kwargs,
|
|
)
|
|
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="english",
|
|
**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)
|