Add asynchronous model download handling for TTS and Whisper models

- Implemented background tasks for downloading TTS and Whisper models to prevent blocking HTTP responses.
- Enhanced error handling during model downloads, providing users with real-time feedback on download status.
- Updated HTTP responses to indicate when models are being downloaded, improving user experience during model initialization.
This commit is contained in:
Jamie Pine
2026-01-29 02:55:17 -08:00
parent 8d730621bc
commit 99ddd5a0b4
+62
View File
@@ -11,6 +11,7 @@ from fastapi.staticfiles import StaticFiles
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
from typing import List, Optional from typing import List, Optional
from datetime import datetime from datetime import datetime
import asyncio
import uvicorn import uvicorn
import argparse import argparse
import torch import torch
@@ -451,6 +452,36 @@ async def generate_speech(
tts_model = tts.get_tts_model() tts_model = tts.get_tts_model()
# Load the requested model size if different from current (async to not block) # Load the requested model size if different from current (async to not block)
model_size = data.model_size or "1.7B" model_size = data.model_size or "1.7B"
# Check if model needs to be downloaded first
model_path = tts_model._get_model_path(model_size)
if model_path.startswith("Qwen/"):
# Model not cached - check if it exists remotely or needs download
from huggingface_hub import constants as hf_constants
repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_path.replace("/", "--"))
if not repo_cache.exists():
# Start download in background
model_name = f"qwen-tts-{model_size}"
async def download_model_background():
try:
await tts_model.load_model_async(model_size)
except Exception as e:
task_manager.error_download(model_name, str(e))
task_manager.start_download(model_name)
asyncio.create_task(download_model_background())
# Return 202 Accepted with download info
raise HTTPException(
status_code=202,
detail={
"message": f"Model {model_size} is being downloaded. Please wait and try again.",
"model_name": model_name,
"downloading": True
}
)
await tts_model.load_model_async(model_size) await tts_model.load_model_async(model_size)
audio, sample_rate = await tts_model.generate( audio, sample_rate = await tts_model.generate(
data.text, data.text,
@@ -684,6 +715,37 @@ async def transcribe_audio(
# Transcribe # Transcribe
whisper_model = transcribe.get_whisper_model() whisper_model = transcribe.get_whisper_model()
# Check if Whisper model is downloaded (uses default size "base")
model_size = whisper_model.model_size
model_name = f"openai/whisper-{model_size}"
# Check if model is cached
from huggingface_hub import constants as hf_constants
repo_cache = Path(hf_constants.HF_HUB_CACHE) / ("models--" + model_name.replace("/", "--"))
if not repo_cache.exists():
# Start download in background
progress_model_name = f"whisper-{model_size}"
async def download_whisper_background():
try:
await whisper_model.load_model_async(model_size)
except Exception as e:
get_task_manager().error_download(progress_model_name, str(e))
get_task_manager().start_download(progress_model_name)
asyncio.create_task(download_whisper_background())
# Return 202 Accepted
raise HTTPException(
status_code=202,
detail={
"message": f"Whisper model {model_size} is being downloaded. Please wait and try again.",
"model_name": progress_model_name,
"downloading": True
}
)
text = await whisper_model.transcribe(tmp_path, language) text = await whisper_model.transcribe(tmp_path, language)
return models.TranscriptionResponse( return models.TranscriptionResponse(