diff --git a/backend/backends/base.py b/backend/backends/base.py index 70ec11ef3..5d5343562 100644 --- a/backend/backends/base.py +++ b/backend/backends/base.py @@ -171,10 +171,10 @@ def check_cuda_compatibility() -> tuple[bool, str | None]: def empty_device_cache(device: str) -> None: """ - Free cached memory on the given device (CUDA or XPU). + Free cached memory on the given device (CUDA, XPU, or MPS). - Backends should call this after unloading models so VRAM is returned - to the OS. + Backends should call this after unloading models so VRAM/unified + memory is returned to the OS. """ import torch @@ -182,6 +182,8 @@ def empty_device_cache(device: str) -> None: torch.cuda.empty_cache() elif device == "xpu" and hasattr(torch, "xpu"): torch.xpu.empty_cache() + elif device == "mps" and hasattr(torch, "mps"): + torch.mps.empty_cache() def manual_seed(seed: int, device: str) -> None: diff --git a/backend/backends/mlx_backend.py b/backend/backends/mlx_backend.py index 9692e59b4..c867bdbed 100644 --- a/backend/backends/mlx_backend.py +++ b/backend/backends/mlx_backend.py @@ -5,11 +5,27 @@ from typing import Optional, List, Tuple import asyncio import logging +from concurrent.futures import ThreadPoolExecutor import numpy as np from pathlib import Path logger = logging.getLogger(__name__) +# MLX's Metal command encoder/stream is thread-local: it's bound to +# whichever thread first touches the GPU device. asyncio.to_thread() uses +# the loop's default executor, which can (and does) hand different calls +# to different worker threads — model load on one thread, then generate() +# on another raises "There is no Stream(gpu, N) in current thread." Pin +# every MLX call to the same single dedicated worker instead. +# See https://github.com/jamiepine/voicebox/issues/699 +_mlx_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mlx-worker") + + +def _run_on_mlx_thread(func, *args): + """Run ``func(*args)`` on the single dedicated MLX worker thread.""" + 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 @@ -77,12 +93,14 @@ async def load_model_async(self, model_size: Optional[str] = None): if self.model is not None and self._current_model_size == model_size: return - # Unload existing model if different size requested + # Unload existing model if different size requested — routed + # through the same dedicated thread as load/generate so it stays + # serialized with any in-flight MLX call (see _run_on_mlx_thread). if self.model is not None and self._current_model_size != model_size: - self.unload_model() + await _run_on_mlx_thread(self.unload_model) - # Run blocking load in thread pool - await asyncio.to_thread(self._load_model_sync, model_size) + # Run blocking load on the dedicated MLX thread + await _run_on_mlx_thread(self._load_model_sync, model_size) # Alias for compatibility load_model = load_model_async @@ -258,8 +276,8 @@ def _generate_sync(): return audio, sample_rate - # Run blocking inference in thread pool - audio, sample_rate = await asyncio.to_thread(_generate_sync) + # Run blocking inference on the dedicated MLX thread + audio, sample_rate = await _run_on_mlx_thread(_generate_sync) return audio, sample_rate @@ -292,8 +310,8 @@ async def load_model_async(self, model_size: Optional[str] = None): if self.model is not None and self.model_size == model_size: return - # Run blocking load in thread pool - await asyncio.to_thread(self._load_model_sync, model_size) + # Run blocking load on the dedicated MLX thread + await _run_on_mlx_thread(self._load_model_sync, model_size) # Alias for compatibility load_model = load_model_async @@ -363,5 +381,5 @@ def _transcribe_sync(): else: return str(result).strip() - # Run blocking transcription in thread pool - return await asyncio.to_thread(_transcribe_sync) + # Run blocking transcription on the dedicated MLX thread + return await _run_on_mlx_thread(_transcribe_sync) diff --git a/backend/backends/pytorch_backend.py b/backend/backends/pytorch_backend.py index f8ae79b86..99fd70a51 100644 --- a/backend/backends/pytorch_backend.py +++ b/backend/backends/pytorch_backend.py @@ -34,7 +34,7 @@ def __init__(self, model_size: str = "1.7B"): def _get_device(self) -> str: """Get the best available device.""" - return get_torch_device(allow_xpu=True, allow_directml=True) + return get_torch_device(allow_xpu=True, allow_directml=True, allow_mps=True) def is_loaded(self) -> bool: """Check if model is loaded.""" @@ -256,7 +256,7 @@ def __init__(self, model_size: str = "base"): def _get_device(self) -> str: """Get the best available device.""" - return get_torch_device(allow_xpu=True, allow_directml=True) + return get_torch_device(allow_xpu=True, allow_directml=True, allow_mps=True) def is_loaded(self) -> bool: """Check if model is loaded.""" diff --git a/backend/backends/qwen_custom_voice_backend.py b/backend/backends/qwen_custom_voice_backend.py index 74f739bba..3eddedced 100644 --- a/backend/backends/qwen_custom_voice_backend.py +++ b/backend/backends/qwen_custom_voice_backend.py @@ -25,6 +25,7 @@ from .base import ( is_model_cached, get_torch_device, + empty_device_cache, combine_voice_prompts as _combine_voice_prompts, model_load_progress, ) @@ -65,7 +66,7 @@ def __init__(self, model_size: str = "1.7B"): self._current_model_size: Optional[str] = None def _get_device(self) -> str: - return get_torch_device(allow_xpu=True, allow_directml=True) + return get_torch_device(allow_xpu=True, allow_directml=True, allow_mps=True) def is_loaded(self) -> bool: return self.model is not None @@ -127,8 +128,7 @@ def unload_model(self) -> None: self.model = None self._current_model_size = None - if torch.cuda.is_available(): - torch.cuda.empty_cache() + empty_device_cache(self.device) logger.info("Qwen CustomVoice unloaded") diff --git a/backend/requirements-mlx.txt b/backend/requirements-mlx.txt index b1a82d5d9..cf6dc692c 100644 --- a/backend/requirements-mlx.txt +++ b/backend/requirements-mlx.txt @@ -9,14 +9,42 @@ mlx>=0.30.0 # M1 installs with `ModuleNotFoundError: miniaudio` (issue #505). miniaudio>=1.59 -# NOTE: mlx-audio is intentionally not listed here. From 0.3.1 onward it -# declares `transformers==5.0.0rc3` / `>=5.0.0`, which conflicts with the -# `transformers<=4.57.6` cap in requirements.txt and breaks CI's clean -# resolver. The mlx-audio API surface we use (mlx_audio.tts.load, -# mlx_audio.stt.load) works fine on transformers 4.57.x in practice. +# sounddevice is a runtime dep of mlx-audio (device enumeration for its +# playback helpers). mlx-audio 0.4.1 pins sounddevice==0.5.3 exactly; not +# listing it here left it unresolved on a from-scratch install, which +# doesn't crash immediately but leaves pip's resolver warning about a +# missing dependency and mlx-audio's playback helper unusable. +sounddevice==0.5.3 + +# NOTE: mlx-lm and mlx-audio are intentionally not listed here — both are +# installed separately with --no-deps (see `just setup-python` and +# .github/workflows/release.yml), NOT via this requirements file. From +# 0.30.0 onward mlx-lm declares `transformers>=5.0.0` (mlx-audio>=0.3.1 +# does the same), which conflicts with the `transformers<=4.57.6` cap in +# requirements.txt. A plain `pip install -r` here would let pip resolve +# mlx-lm freely and silently upgrade transformers past that cap, which +# breaks the qwen_tts (PyTorch) engines at *runtime*, not install time: +# qwen_tts's `@check_model_inputs()` call site +# (qwen_tts/core/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py) uses +# the pre-5.0 signature (a decorator factory called with no arguments). +# transformers>=5.0 changed `check_model_inputs` to a plain decorator +# (`check_model_inputs(func)`, no parens at the call site), so loading +# Qwen CustomVoice through qwen_tts on the newer transformers raises +# `TypeError: check_model_inputs() missing 1 required positional +# argument: 'func'`. # -# Install it via `pip install --no-deps mlx-audio==0.4.1` after this file -# (see .github/workflows/release.yml). Most other mlx-audio runtime deps -# (huggingface_hub, librosa, mlx-lm, numba, numpy, protobuf, pyloudnorm, -# sounddevice, tqdm) are already in requirements.txt or pulled in by -# other engines. +# Install mlx-lm and mlx-audio via: +# pip install --no-deps mlx-lm==0.31.1 +# pip install --no-deps mlx-audio==0.4.1 +# --no-deps means both packages' own `transformers>=5.0.0` requirement is +# never consulted, so transformers stays at whatever requirements.txt +# installed. The mlx_audio.tts/mlx_audio.stt API surface we use works +# fine against transformers 4.57.x in practice (verified in dev and CI). +# `just setup-python` previously installed this file's contents (mlx + +# miniaudio + sounddevice) and stopped — mlx_lm and mlx_audio were never +# installed at all, so the backend silently fell back to the PyTorch +# backend on Apple Silicon instead of MLX (issue #823), and once mlx_lm +# was installed by hand without --no-deps (the obvious fix to try), it +# reintroduced the transformers conflict above (issue #699 is the +# MLX-generation-side symptom of the same root cause). See the two +# `pip install --no-deps` lines added to `just setup-python` below. diff --git a/backend/routes/health.py b/backend/routes/health.py index 1568455dc..482dd2a1e 100644 --- a/backend/routes/health.py +++ b/backend/routes/health.py @@ -175,6 +175,8 @@ async def health(): default_variant = "cuda" elif has_xpu: default_variant = "xpu" + elif backend_type == "mlx": + default_variant = "mps" return models.HealthResponse( status="healthy", diff --git a/justfile b/justfile index 156806d88..95e737849 100644 --- a/justfile +++ b/justfile @@ -72,6 +72,14 @@ setup-python: if [ "$(uname -m)" = "arm64" ] && [ "$(uname)" = "Darwin" ]; then echo "Detected Apple Silicon — installing MLX dependencies..." {{ pip }} install -r {{ backend_dir }}/requirements-mlx.txt + # mlx-lm and mlx-audio are installed --no-deps and separately from + # requirements-mlx.txt on purpose — see the NOTE at the bottom of + # that file. Without this step the backend silently falls back to + # the PyTorch backend on Apple Silicon (issue #823), and installing + # mlx-lm without --no-deps upgrades transformers past the + # requirements.txt cap and breaks qwen_tts at runtime (issue #699). + {{ pip }} install --no-deps mlx-lm==0.31.1 + {{ pip }} install --no-deps mlx-audio==0.4.1 fi {{ pip }} install git+https://github.com/QwenLM/Qwen3-TTS.git {{ pip }} install pyinstaller ruff pytest pytest-asyncio -q