mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-16 21:30:39 -07:00
fix: address review feedback — race condition, GPU safety, task GC
- Add threading lock to get_tts_backend_for_engine() to prevent race condition where concurrent requests could create duplicate backend instances (double-checked locking pattern) - Fix LuxTTS generate: call .detach().cpu() before .numpy() so it works on GPU/MPS devices, not just CPU - Store background download tasks in a module-level set to prevent garbage collection before completion (asyncio.create_task fire-and- forget pattern) - Deduplicate cache_key computation in LuxTTS create_voice_prompt - Prefix unused sr variable with underscore
This commit is contained in:
@@ -4,6 +4,7 @@ Backend abstraction layer for TTS and STT.
|
||||
Provides a unified interface for MLX and PyTorch backends.
|
||||
"""
|
||||
|
||||
import threading
|
||||
from typing import Protocol, Optional, Tuple, List
|
||||
from typing_extensions import runtime_checkable
|
||||
import numpy as np
|
||||
@@ -113,6 +114,7 @@ class STTBackend(Protocol):
|
||||
# Global backend instances
|
||||
_tts_backend: Optional[TTSBackend] = None
|
||||
_tts_backends: dict[str, TTSBackend] = {}
|
||||
_tts_backends_lock = threading.Lock()
|
||||
_stt_backend: Optional[STTBackend] = None
|
||||
|
||||
# Supported TTS engines
|
||||
@@ -144,25 +146,32 @@ def get_tts_backend_for_engine(engine: str) -> TTSBackend:
|
||||
"""
|
||||
global _tts_backends
|
||||
|
||||
# Fast path: check without lock
|
||||
if engine in _tts_backends:
|
||||
return _tts_backends[engine]
|
||||
|
||||
if engine == "qwen":
|
||||
backend_type = get_backend_type()
|
||||
if backend_type == "mlx":
|
||||
from .mlx_backend import MLXTTSBackend
|
||||
backend = MLXTTSBackend()
|
||||
# Slow path: create with lock to avoid duplicate instantiation
|
||||
with _tts_backends_lock:
|
||||
# Double-check after acquiring lock
|
||||
if engine in _tts_backends:
|
||||
return _tts_backends[engine]
|
||||
|
||||
if engine == "qwen":
|
||||
backend_type = get_backend_type()
|
||||
if backend_type == "mlx":
|
||||
from .mlx_backend import MLXTTSBackend
|
||||
backend = MLXTTSBackend()
|
||||
else:
|
||||
from .pytorch_backend import PyTorchTTSBackend
|
||||
backend = PyTorchTTSBackend()
|
||||
elif engine == "luxtts":
|
||||
from .luxtts_backend import LuxTTSBackend
|
||||
backend = LuxTTSBackend()
|
||||
else:
|
||||
from .pytorch_backend import PyTorchTTSBackend
|
||||
backend = PyTorchTTSBackend()
|
||||
elif engine == "luxtts":
|
||||
from .luxtts_backend import LuxTTSBackend
|
||||
backend = LuxTTSBackend()
|
||||
else:
|
||||
raise ValueError(f"Unknown TTS engine: {engine}. Supported: {list(TTS_ENGINES.keys())}")
|
||||
|
||||
_tts_backends[engine] = backend
|
||||
return backend
|
||||
raise ValueError(f"Unknown TTS engine: {engine}. Supported: {list(TTS_ENGINES.keys())}")
|
||||
|
||||
_tts_backends[engine] = backend
|
||||
return backend
|
||||
|
||||
|
||||
def get_stt_backend() -> STTBackend:
|
||||
|
||||
@@ -171,9 +171,10 @@ class LuxTTSBackend:
|
||||
"""
|
||||
await self.load_model()
|
||||
|
||||
if use_cache:
|
||||
# Include "luxtts" in the cache key so it doesn't collide with Qwen prompts
|
||||
cache_key = "luxtts_" + get_cache_key(audio_path, reference_text)
|
||||
# Compute cache key once for both lookup and storage
|
||||
cache_key = ("luxtts_" + get_cache_key(audio_path, reference_text)) if use_cache else None
|
||||
|
||||
if cache_key:
|
||||
cached = get_cached_voice_prompt(cache_key)
|
||||
if cached is not None and isinstance(cached, dict):
|
||||
return cached, True
|
||||
@@ -187,8 +188,7 @@ class LuxTTSBackend:
|
||||
|
||||
encoded = await asyncio.to_thread(_encode_sync)
|
||||
|
||||
if use_cache:
|
||||
cache_key = "luxtts_" + get_cache_key(audio_path, reference_text)
|
||||
if cache_key:
|
||||
cache_voice_prompt(cache_key, encoded)
|
||||
|
||||
return encoded, False
|
||||
@@ -206,7 +206,7 @@ class LuxTTSBackend:
|
||||
"""
|
||||
combined_audio = []
|
||||
for path in audio_paths:
|
||||
audio, sr = load_audio(path, sample_rate=24000)
|
||||
audio, _sr = load_audio(path, sample_rate=24000)
|
||||
audio = normalize_audio(audio)
|
||||
combined_audio.append(audio)
|
||||
|
||||
@@ -257,8 +257,8 @@ class LuxTTSBackend:
|
||||
return_smooth=False, # 48kHz output
|
||||
)
|
||||
|
||||
# LuxTTS returns a tensor, convert to numpy
|
||||
audio = wav.numpy().squeeze()
|
||||
# LuxTTS returns a tensor (may be on GPU/MPS), move to CPU first
|
||||
audio = wav.detach().cpu().numpy().squeeze()
|
||||
return audio, 48000
|
||||
|
||||
return await asyncio.to_thread(_generate_sync)
|
||||
|
||||
+17
-5
@@ -48,6 +48,18 @@ from .utils.tasks import get_task_manager
|
||||
from .utils.cache import clear_voice_prompt_cache
|
||||
from .platform_detect import get_backend_type
|
||||
|
||||
# Keep references to fire-and-forget background tasks to prevent GC
|
||||
_background_tasks: set = set()
|
||||
|
||||
|
||||
def _create_background_task(coro) -> asyncio.Task:
|
||||
"""Create a background task and prevent it from being garbage collected."""
|
||||
task = asyncio.create_task(coro)
|
||||
_background_tasks.add(task)
|
||||
task.add_done_callback(_background_tasks.discard)
|
||||
return task
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title="voicebox API",
|
||||
description="Production-quality Qwen3-TTS voice cloning API",
|
||||
@@ -622,7 +634,7 @@ async def generate_speech(
|
||||
task_manager.error_download(model_name, str(e))
|
||||
|
||||
task_manager.start_download(model_name)
|
||||
asyncio.create_task(download_model_background())
|
||||
_create_background_task(download_model_background())
|
||||
|
||||
raise HTTPException(
|
||||
status_code=202,
|
||||
@@ -646,7 +658,7 @@ async def generate_speech(
|
||||
task_manager.error_download(model_name, str(e))
|
||||
|
||||
task_manager.start_download(model_name)
|
||||
asyncio.create_task(download_luxtts_background())
|
||||
_create_background_task(download_luxtts_background())
|
||||
|
||||
raise HTTPException(
|
||||
status_code=202,
|
||||
@@ -986,7 +998,7 @@ async def transcribe_audio(
|
||||
get_task_manager().error_download(progress_model_name, str(e))
|
||||
|
||||
get_task_manager().start_download(progress_model_name)
|
||||
asyncio.create_task(download_whisper_background())
|
||||
_create_background_task(download_whisper_background())
|
||||
|
||||
# Return 202 Accepted
|
||||
raise HTTPException(
|
||||
@@ -1639,7 +1651,7 @@ async def trigger_model_download(request: models.ModelDownloadRequest):
|
||||
)
|
||||
|
||||
# Start download in background task (don't await)
|
||||
asyncio.create_task(download_in_background())
|
||||
_create_background_task(download_in_background())
|
||||
|
||||
# Return immediately - frontend should poll progress endpoint
|
||||
return {"message": f"Model {request.model_name} download started"}
|
||||
@@ -1889,7 +1901,7 @@ async def download_cuda_backend():
|
||||
import logging
|
||||
logging.getLogger(__name__).error(f"CUDA download failed: {e}")
|
||||
|
||||
asyncio.create_task(_download())
|
||||
_create_background_task(_download())
|
||||
return {"message": "CUDA backend download started", "progress_key": "cuda-backend"}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user