From d04fa88853b7598ee7f0583bdb76252beef9d9be Mon Sep 17 00:00:00 2001 From: mARTin-B78 Date: Tue, 26 May 2026 23:07:51 +0200 Subject: [PATCH] refactor: load models asynchronously and add CUDA warmup for VoiceDesign and CustomVoice servers --- config/run_customvoice_server.py | 62 ++++++++++++++++++++++++-------- config/run_voicedesign_server.py | 61 ++++++++++++++++++++++++------- 2 files changed, 96 insertions(+), 27 deletions(-) diff --git a/config/run_customvoice_server.py b/config/run_customvoice_server.py index 1d5d5f5..ec2469b 100644 --- a/config/run_customvoice_server.py +++ b/config/run_customvoice_server.py @@ -21,6 +21,7 @@ import uvicorn from fastapi import FastAPI, HTTPException from fastapi.responses import JSONResponse, Response, StreamingResponse from pydantic import BaseModel +from contextlib import asynccontextmanager sys.path.append("/app") from faster_qwen3_tts.model import FasterQwen3TTS @@ -35,6 +36,51 @@ default_voice: str = None SAMPLE_RATE = 24000 DEFAULT_MAX_NEW_TOKENS = 2048 _model_lock = threading.Lock() +_load_model_kwargs = None + + +@asynccontextmanager +async def lifespan(app: FastAPI): + loop = asyncio.get_event_loop() + loop.run_in_executor(None, _do_load_and_warmup) + yield + + +def _do_load_and_warmup(): + global tts_model, SAMPLE_RATE + import torch + args = _load_model_kwargs + try: + logger.info("Loading CustomVoice model %s ...", args.model) + model = FasterQwen3TTS.from_pretrained( + args.model, + device=args.device, + dtype=torch.bfloat16, + attn_implementation="sdpa", + max_seq_len=args.max_seq_len, + ) + SAMPLE_RATE = model.sample_rate + logger.info("Model ready. Sample rate: %d Hz", SAMPLE_RATE) + + # Warmup + logger.info("Warming up CUDA graphs (first request will be fast)...") + try: + for _ in model.generate_custom_voice_streaming( + text="Warmup.", + speaker="Ryan", + language="English" + ): + pass + logger.info("CUDA warmup complete — server ready.") + except Exception as exc: + logger.warning("Warmup failed (non-fatal): %s", exc) + + tts_model = model + except Exception as exc: + logger.error("Failed to load model: %s", exc) + + +app = FastAPI(lifespan=lifespan) class SpeechRequest(BaseModel): @@ -196,7 +242,7 @@ async def options_handler(path: str): def main(): - global tts_model, voices, default_voice, SAMPLE_RATE, DEFAULT_MAX_NEW_TOKENS + global voices, default_voice, SAMPLE_RATE, DEFAULT_MAX_NEW_TOKENS, _load_model_kwargs parser = argparse.ArgumentParser() parser.add_argument("--model", default="/models/Qwen3-TTS-CustomVoice") @@ -207,25 +253,13 @@ def main(): parser.add_argument("--max-seq-len", type=int, default=2048) args = parser.parse_args() DEFAULT_MAX_NEW_TOKENS = args.max_seq_len + _load_model_kwargs = args with open(args.voices) as f: voices = json.load(f) default_voice = next(iter(voices), None) _build_voice_list() - import torch - - logger.info("Loading CustomVoice model %s ...", args.model) - tts_model = FasterQwen3TTS.from_pretrained( - args.model, - device=args.device, - dtype=torch.bfloat16, - attn_implementation="sdpa", - max_seq_len=args.max_seq_len, - ) - SAMPLE_RATE = tts_model.sample_rate - logger.info("Model ready. Sample rate: %d Hz", SAMPLE_RATE) - uvicorn.run(app, host=args.host, port=args.port, log_level="info") diff --git a/config/run_voicedesign_server.py b/config/run_voicedesign_server.py index 47e98f2..d17aa34 100644 --- a/config/run_voicedesign_server.py +++ b/config/run_voicedesign_server.py @@ -20,6 +20,7 @@ import uvicorn from fastapi import FastAPI, HTTPException from fastapi.responses import Response, StreamingResponse, JSONResponse from pydantic import BaseModel +from contextlib import asynccontextmanager sys.path.append("/app") from faster_qwen3_tts.model import FasterQwen3TTS @@ -34,6 +35,51 @@ default_voice: str = None SAMPLE_RATE = 24000 DEFAULT_MAX_NEW_TOKENS = 2048 _model_lock = threading.Lock() +_load_model_kwargs = None + + +@asynccontextmanager +async def lifespan(app: FastAPI): + loop = asyncio.get_event_loop() + loop.run_in_executor(None, _do_load_and_warmup) + yield + + +def _do_load_and_warmup(): + global tts_model, SAMPLE_RATE + import torch + args = _load_model_kwargs + try: + logger.info("Loading VoiceDesign model %s …", args.model) + model = FasterQwen3TTS.from_pretrained( + args.model, + device=args.device, + dtype=torch.bfloat16, + attn_implementation="sdpa", + max_seq_len=args.max_seq_len, + ) + SAMPLE_RATE = model.sample_rate + logger.info("Model ready. Sample rate: %d Hz", SAMPLE_RATE) + + # Warmup + logger.info("Warming up CUDA graphs (first request will be fast)...") + try: + for _ in model.generate_voice_design_streaming( + text="Warmup.", + instruct="Warmup.", + language="English" + ): + pass + logger.info("CUDA warmup complete — server ready.") + except Exception as exc: + logger.warning("Warmup failed (non-fatal): %s", exc) + + tts_model = model + except Exception as exc: + logger.error("Failed to load model: %s", exc) + + +app = FastAPI(lifespan=lifespan) # --------------------------------------------------------------------------- @@ -208,7 +254,7 @@ async def options_handler(path: str): # --------------------------------------------------------------------------- def main(): - global tts_model, voices, default_voice, SAMPLE_RATE, DEFAULT_MAX_NEW_TOKENS + global voices, default_voice, SAMPLE_RATE, DEFAULT_MAX_NEW_TOKENS, _load_model_kwargs parser = argparse.ArgumentParser() parser.add_argument("--model", default="/models/Qwen3-TTS-VoiceDesign") @@ -219,24 +265,13 @@ def main(): parser.add_argument("--max-seq-len", type=int, default=2048) args = parser.parse_args() DEFAULT_MAX_NEW_TOKENS = args.max_seq_len + _load_model_kwargs = args with open(args.voices) as f: voices = json.load(f) default_voice = next(iter(voices), None) _build_voice_list() - import torch - logger.info("Loading VoiceDesign model %s …", args.model) - tts_model = FasterQwen3TTS.from_pretrained( - args.model, - device=args.device, - dtype=torch.bfloat16, - attn_implementation="sdpa", - max_seq_len=args.max_seq_len, - ) - SAMPLE_RATE = tts_model.sample_rate - logger.info("Model ready. Sample rate: %d Hz", SAMPLE_RATE) - uvicorn.run(app, host=args.host, port=args.port, log_level="info")