mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-10-03 17:15:19 -07:00
Ron David Ben Ishayandcapy-ai-staging[bot]
fa8db820d7
Two real follow-on issues found by CodeRabbit on PR #989: 1. MLXTTSBackend.load_model_async called self.unload_model() directly on the event-loop thread before dispatching _load_model_sync to the MLX worker thread - model teardown ran on the wrong OS thread, same class of bug the PR itself fixes. Combined unload+load into one _reload_sync callable submitted as a single MLX-thread operation. 2. Public unload_model() (called synchronously from the /models/unload routes) also ran on the caller thread. Now submits to the MLX executor and blocks on the result, so teardown always happens on the worker thread regardless of caller. Same fix applied to MLXSTTBackend. 3. Both backends cache self.model on the instance, but generate()/ transcribe() only locked their own internal load step - a concurrent request for a different model_size could swap self.model between one request's load and its inference (real race: routes/ generations.py calls load_engine_model() and generate_chunked() as separate awaited steps with a gap between them). Added a per-backend asyncio.Lock held across the full load+inference sequence in both generate() and transcribe(). Note: this closes the race for callers using the backend's own public methods consistently; the wider route-level orchestration race (load_engine_model + generate_chunked as two separate calls) is a follow-up outside this file's scope. Verified: same-size and cross-size-switch generations both complete clean after the patch; /models/{name}/unload route returns 200 without deadlocking the (still-responsive) server.
420 lines
16 KiB
Python
420 lines
16 KiB
Python
"""
|
|
MLX backend implementation for TTS and STT using mlx-audio.
|
|
"""
|
|
|
|
from typing import Optional, List, Tuple
|
|
import asyncio
|
|
import logging
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# MLX's Metal backend keeps a per-thread stream registry. Loading a model on
|
|
# one worker thread (via asyncio.to_thread, which round-robins across the
|
|
# default executor's pool) and then generating on a different worker thread
|
|
# raises "There is no Stream(gpu, N) in current thread." All MLX calls in
|
|
# this module must therefore run on the SAME OS thread for the process
|
|
# lifetime — route them through this single-worker executor instead of
|
|
# asyncio.to_thread's shared multi-worker pool.
|
|
_mlx_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mlx-worker")
|
|
|
|
|
|
def _run_on_mlx_thread(func, *args):
|
|
loop = asyncio.get_running_loop()
|
|
return loop.run_in_executor(_mlx_executor, func, *args)
|
|
|
|
# PATCH: Import and apply offline patch BEFORE any huggingface_hub usage
|
|
# This prevents mlx_audio from making network requests when models are cached
|
|
from ..utils.hf_offline_patch import patch_huggingface_hub_offline, ensure_original_qwen_config_cached
|
|
|
|
patch_huggingface_hub_offline()
|
|
ensure_original_qwen_config_cached()
|
|
|
|
from . import TTSBackend, STTBackend, LANGUAGE_CODE_TO_NAME, WHISPER_HF_REPOS
|
|
from .base import is_model_cached, combine_voice_prompts as _combine_voice_prompts, model_load_progress
|
|
from ..utils.cache import get_cache_key, get_cached_voice_prompt, cache_voice_prompt
|
|
|
|
|
|
class MLXTTSBackend:
|
|
"""MLX-based TTS backend using mlx-audio."""
|
|
|
|
def __init__(self, model_size: str = "1.7B"):
|
|
self.model = None
|
|
self.model_size = model_size
|
|
self._current_model_size = None
|
|
# Guards the whole load-then-use sequence in generate()/create_voice_prompt()
|
|
# so a concurrent request for a different model_size can't swap self.model
|
|
# out from under an in-flight request between its load and its inference.
|
|
self._op_lock = asyncio.Lock()
|
|
|
|
def is_loaded(self) -> bool:
|
|
"""Check if model is loaded."""
|
|
return self.model is not None
|
|
|
|
def _get_model_path(self, model_size: str) -> str:
|
|
"""
|
|
Get the MLX model path.
|
|
|
|
Args:
|
|
model_size: Model size (1.7B or 0.6B)
|
|
|
|
Returns:
|
|
HuggingFace Hub model ID for MLX
|
|
"""
|
|
mlx_model_map = {
|
|
"1.7B": "mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16",
|
|
"0.6B": "mlx-community/Qwen3-TTS-12Hz-0.6B-Base-bf16",
|
|
}
|
|
|
|
if model_size not in mlx_model_map:
|
|
raise ValueError(f"Unknown model size: {model_size}")
|
|
|
|
hf_model_id = mlx_model_map[model_size]
|
|
logger.info("Will download MLX model from HuggingFace Hub: %s", hf_model_id)
|
|
|
|
return hf_model_id
|
|
|
|
def _is_model_cached(self, model_size: str) -> bool:
|
|
return is_model_cached(
|
|
self._get_model_path(model_size),
|
|
weight_extensions=(".safetensors", ".bin", ".npz"),
|
|
)
|
|
|
|
async def load_model_async(self, model_size: Optional[str] = None):
|
|
"""
|
|
Lazy load the MLX TTS model.
|
|
|
|
Args:
|
|
model_size: Model size to load (1.7B or 0.6B)
|
|
"""
|
|
if model_size is None:
|
|
model_size = self.model_size
|
|
|
|
# If already loaded with correct size, return
|
|
if self.model is not None and self._current_model_size == model_size:
|
|
return
|
|
|
|
# Unload (if needed) and load as ONE callable on the MLX worker thread.
|
|
# Doing this as two separate _run_on_mlx_thread calls would run the
|
|
# unload on whichever thread issues the second call — usually still
|
|
# correct, but a caller-side await gap between them would let another
|
|
# coroutine slip a conflicting load in between. One callable removes
|
|
# the gap.
|
|
await _run_on_mlx_thread(self._reload_sync, model_size)
|
|
|
|
# Alias for compatibility
|
|
load_model = load_model_async
|
|
|
|
def _reload_sync(self, model_size: str):
|
|
"""Unload a mismatched model and load the requested one, in one MLX-thread op."""
|
|
if self.model is not None and self._current_model_size != model_size:
|
|
self._unload_model_sync()
|
|
self._load_model_sync(model_size)
|
|
|
|
def _load_model_sync(self, model_size: str):
|
|
"""Synchronous model loading."""
|
|
model_path = self._get_model_path(model_size)
|
|
model_name = f"qwen-tts-{model_size}"
|
|
is_cached = self._is_model_cached(model_size)
|
|
|
|
with model_load_progress(model_name, is_cached):
|
|
from mlx_audio.tts import load
|
|
|
|
logger.info("Loading MLX TTS model %s...", model_size)
|
|
|
|
self.model = load(model_path)
|
|
|
|
self._current_model_size = model_size
|
|
self.model_size = model_size
|
|
logger.info("MLX TTS model %s loaded successfully", model_size)
|
|
|
|
def unload_model(self):
|
|
"""Unload the model to free memory.
|
|
|
|
Safe to call from any thread (e.g. the FastAPI event loop, from the
|
|
/models/unload routes): the actual teardown is submitted to the MLX
|
|
worker thread and awaited synchronously here, so Metal resources are
|
|
always released on the same OS thread that created them. Do not call
|
|
this from within a callable already running ON the MLX worker thread
|
|
(e.g. from _reload_sync) — use _unload_model_sync directly there, or
|
|
this would deadlock the single-worker executor waiting on itself.
|
|
"""
|
|
if self.model is not None:
|
|
_mlx_executor.submit(self._unload_model_sync).result()
|
|
|
|
def _unload_model_sync(self):
|
|
if self.model is not None:
|
|
del self.model
|
|
self.model = None
|
|
self._current_model_size = None
|
|
logger.info("MLX TTS model unloaded")
|
|
|
|
async def create_voice_prompt(
|
|
self,
|
|
audio_path: str,
|
|
reference_text: str,
|
|
use_cache: bool = True,
|
|
) -> Tuple[dict, bool]:
|
|
"""
|
|
Create voice prompt from reference audio.
|
|
|
|
MLX backend stores voice prompt as a dict with audio path and text.
|
|
The actual voice prompt processing happens during generation.
|
|
|
|
Args:
|
|
audio_path: Path to reference audio file
|
|
reference_text: Transcript of reference audio
|
|
use_cache: Whether to use cached prompt if available
|
|
|
|
Returns:
|
|
Tuple of (voice_prompt_dict, was_cached)
|
|
"""
|
|
async with self._op_lock:
|
|
await self.load_model_async(None)
|
|
|
|
# Check cache if enabled
|
|
if use_cache:
|
|
cache_key = get_cache_key(audio_path, reference_text)
|
|
cached_prompt = get_cached_voice_prompt(cache_key)
|
|
if cached_prompt is not None:
|
|
# Return cached prompt (should be dict format)
|
|
if isinstance(cached_prompt, dict):
|
|
# Validate that the cached audio file still exists
|
|
cached_audio_path = cached_prompt.get("ref_audio") or cached_prompt.get("ref_audio_path")
|
|
if cached_audio_path and Path(cached_audio_path).exists():
|
|
return cached_prompt, True
|
|
else:
|
|
# Cached file no longer exists, invalidate cache
|
|
logger.warning("Cached audio file not found: %s, regenerating prompt", cached_audio_path)
|
|
|
|
# MLX voice prompt format - store audio path and text
|
|
# The model will process this during generation
|
|
voice_prompt_items = {
|
|
"ref_audio": str(audio_path),
|
|
"ref_text": reference_text,
|
|
}
|
|
|
|
# Cache if enabled
|
|
if use_cache:
|
|
cache_key = get_cache_key(audio_path, reference_text)
|
|
cache_voice_prompt(cache_key, voice_prompt_items)
|
|
|
|
return voice_prompt_items, False
|
|
|
|
async def combine_voice_prompts(self, audio_paths, reference_texts):
|
|
return await _combine_voice_prompts(audio_paths, reference_texts)
|
|
|
|
async def generate(
|
|
self,
|
|
text: str,
|
|
voice_prompt: dict,
|
|
language: str = "en",
|
|
seed: Optional[int] = None,
|
|
instruct: Optional[str] = None,
|
|
) -> Tuple[np.ndarray, int]:
|
|
"""
|
|
Generate audio from text using voice prompt.
|
|
|
|
Args:
|
|
text: Text to synthesize
|
|
voice_prompt: Voice prompt dictionary with ref_audio and ref_text
|
|
language: Language code (en or zh) - may not be fully supported by MLX
|
|
seed: Random seed for reproducibility
|
|
instruct: Natural language instruction (may not be supported by MLX)
|
|
|
|
Returns:
|
|
Tuple of (audio_array, sample_rate)
|
|
"""
|
|
logger.info("Generating audio for text: %s", text)
|
|
|
|
def _generate_sync():
|
|
"""Run synchronous generation in thread pool."""
|
|
# MLX generate() returns a generator yielding GenerationResult objects
|
|
audio_chunks = []
|
|
sample_rate = 24000
|
|
lang = LANGUAGE_CODE_TO_NAME.get(language, "auto")
|
|
|
|
# Set seed if provided (MLX uses numpy random)
|
|
if seed is not None:
|
|
import mlx.core as mx
|
|
|
|
np.random.seed(seed)
|
|
mx.random.seed(seed)
|
|
|
|
# Extract voice prompt info
|
|
ref_audio = voice_prompt.get("ref_audio") or voice_prompt.get("ref_audio_path")
|
|
ref_text = voice_prompt.get("ref_text", "")
|
|
|
|
# Validate that the audio file exists
|
|
if ref_audio and not Path(ref_audio).exists():
|
|
logger.warning("Audio file not found: %s", ref_audio)
|
|
logger.warning("This may be due to a cached voice prompt referencing a deleted temp file.")
|
|
logger.warning("Regenerating without voice prompt.")
|
|
ref_audio = None
|
|
|
|
# Inference runs with the process's default HF_HUB_OFFLINE
|
|
# state. Forcing offline here (previously used to avoid lazy
|
|
# mlx_audio lookups hanging when the network drops mid-inference,
|
|
# issue #462) regressed online users because libraries make
|
|
# legitimate metadata calls during generation.
|
|
try:
|
|
if ref_audio:
|
|
# Check if generate accepts ref_audio parameter
|
|
import inspect
|
|
|
|
sig = inspect.signature(self.model.generate)
|
|
if "ref_audio" in sig.parameters:
|
|
# Generate with voice cloning
|
|
for result in self.model.generate(text, ref_audio=ref_audio, ref_text=ref_text, lang_code=lang):
|
|
audio_chunks.append(np.array(result.audio))
|
|
sample_rate = result.sample_rate
|
|
else:
|
|
# Fallback: generate without voice cloning
|
|
for result in self.model.generate(text, lang_code=lang):
|
|
audio_chunks.append(np.array(result.audio))
|
|
sample_rate = result.sample_rate
|
|
else:
|
|
# No voice prompt, generate normally
|
|
for result in self.model.generate(text, lang_code=lang):
|
|
audio_chunks.append(np.array(result.audio))
|
|
sample_rate = result.sample_rate
|
|
except Exception as e:
|
|
# If voice cloning fails, try without it
|
|
logger.warning("Voice cloning failed, generating without voice prompt: %s", e)
|
|
for result in self.model.generate(text, lang_code=lang):
|
|
audio_chunks.append(np.array(result.audio))
|
|
sample_rate = result.sample_rate
|
|
|
|
# Concatenate all chunks
|
|
if audio_chunks:
|
|
audio = np.concatenate([np.asarray(chunk, dtype=np.float32) for chunk in audio_chunks])
|
|
else:
|
|
# Fallback: empty audio
|
|
audio = np.array([], dtype=np.float32)
|
|
|
|
return audio, sample_rate
|
|
|
|
# Hold the op lock across load + inference so a concurrent request
|
|
# for a different model_size can't swap self.model in between (the
|
|
# model is read inside _generate_sync via closure, after this point).
|
|
async with self._op_lock:
|
|
await self.load_model_async(None)
|
|
audio, sample_rate = await _run_on_mlx_thread(_generate_sync)
|
|
|
|
return audio, sample_rate
|
|
|
|
|
|
class MLXSTTBackend:
|
|
"""MLX-based STT backend using mlx-audio Whisper."""
|
|
|
|
def __init__(self, model_size: str = "base"):
|
|
self.model = None
|
|
self.model_size = model_size
|
|
# See MLXTTSBackend._op_lock — same reason.
|
|
self._op_lock = asyncio.Lock()
|
|
|
|
def is_loaded(self) -> bool:
|
|
"""Check if model is loaded."""
|
|
return self.model is not None
|
|
|
|
def _is_model_cached(self, model_size: str) -> bool:
|
|
hf_repo = WHISPER_HF_REPOS.get(model_size, f"openai/whisper-{model_size}")
|
|
return is_model_cached(hf_repo, weight_extensions=(".safetensors", ".bin", ".npz"))
|
|
|
|
async def load_model_async(self, model_size: Optional[str] = None):
|
|
"""
|
|
Lazy load the MLX Whisper model.
|
|
|
|
Args:
|
|
model_size: Model size (tiny, base, small, medium, large)
|
|
"""
|
|
if model_size is None:
|
|
model_size = self.model_size
|
|
|
|
if self.model is not None and self.model_size == model_size:
|
|
return
|
|
|
|
# Run blocking load in thread pool
|
|
await _run_on_mlx_thread(self._load_model_sync, model_size)
|
|
|
|
# Alias for compatibility
|
|
load_model = load_model_async
|
|
|
|
def _load_model_sync(self, model_size: str):
|
|
"""Synchronous model loading."""
|
|
progress_model_name = f"whisper-{model_size}"
|
|
is_cached = self._is_model_cached(model_size)
|
|
|
|
with model_load_progress(progress_model_name, is_cached):
|
|
from mlx_audio.stt import load
|
|
|
|
model_name = WHISPER_HF_REPOS.get(model_size, f"openai/whisper-{model_size}")
|
|
logger.info("Loading MLX Whisper model %s...", model_size)
|
|
|
|
self.model = load(model_name)
|
|
|
|
self.model_size = model_size
|
|
logger.info("MLX Whisper model %s loaded successfully", model_size)
|
|
|
|
def unload_model(self):
|
|
"""Unload the model to free memory.
|
|
|
|
Safe to call from any thread — see MLXTTSBackend.unload_model for why.
|
|
"""
|
|
if self.model is not None:
|
|
_mlx_executor.submit(self._unload_model_sync).result()
|
|
|
|
def _unload_model_sync(self):
|
|
if self.model is not None:
|
|
del self.model
|
|
self.model = None
|
|
logger.info("MLX Whisper model unloaded")
|
|
|
|
async def transcribe(
|
|
self,
|
|
audio_path: str,
|
|
language: Optional[str] = None,
|
|
model_size: Optional[str] = None,
|
|
) -> str:
|
|
"""
|
|
Transcribe audio to text.
|
|
|
|
Args:
|
|
audio_path: Path to audio file
|
|
language: Optional language hint
|
|
model_size: Optional model size override
|
|
|
|
Returns:
|
|
Transcribed text
|
|
"""
|
|
def _transcribe_sync():
|
|
"""Run synchronous transcription in thread pool."""
|
|
# MLX Whisper transcription using generate method
|
|
# The generate method accepts audio path directly
|
|
decode_options = {}
|
|
if language:
|
|
decode_options["language"] = language
|
|
|
|
# Inference runs with the process's default HF_HUB_OFFLINE
|
|
# state — see the comment in MLXTTSBackend.generate for the
|
|
# regression this revert fixes (issue #462).
|
|
result = self.model.generate(str(audio_path), **decode_options)
|
|
|
|
# Extract text from result
|
|
if isinstance(result, str):
|
|
return result.strip()
|
|
elif isinstance(result, dict):
|
|
return result.get("text", "").strip()
|
|
elif hasattr(result, "text"):
|
|
return result.text.strip()
|
|
else:
|
|
return str(result).strip()
|
|
|
|
# Hold the op lock across load + inference so a concurrent request
|
|
# for a different model_size can't swap self.model in between.
|
|
async with self._op_lock:
|
|
await self.load_model_async(model_size)
|
|
return await _run_on_mlx_thread(_transcribe_sync)
|