From c82185f8bed597691581ee7033c2e4c2ba023df3 Mon Sep 17 00:00:00 2001 From: Alex Rieker Date: Sun, 12 Jul 2026 20:30:54 +0200 Subject: [PATCH 1/3] fix: MLX thread affinity crash + missing mlx-lm/mlx-audio in dev setup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two related bugs on Apple Silicon, both stemming from the same root cause (mlx-lm/mlx-audio never getting installed correctly by `just setup-python`), plus one cosmetic health-endpoint bug found along the way: 1. `backend/backends/mlx_backend.py`: pin every MLX model load/generate/ transcribe call to a single dedicated worker thread instead of asyncio.to_thread()'s default pool. MLX's Metal command encoder/stream is thread-local — if model load and generation land on different pool threads (which asyncio.to_thread() does not guarantee against), MLX raises "There is no Stream(gpu, N) in current thread." Fixes #699. 2. `justfile`: `just setup-python` installed requirements-mlx.txt (mlx + miniaudio) on Apple Silicon and stopped — it never ran the `pip install --no-deps mlx-lm==0.31.1` / `mlx-audio==0.4.1` steps that .github/workflows/release.yml already does for packaged builds. Without mlx-audio installed, get_backend_type() falls back to "pytorch" silently: no error, no log line calling this out, just a backend that's much slower than intended and never actually uses MLX. Fixes #823. Installing mlx-lm without --no-deps (the obvious thing to try once you notice #823) reintroduces a second problem: mlx-lm>=0.30.0 requires transformers>=5.0.0, which conflicts with the transformers<=4.57.6 cap in requirements.txt. transformers>=5.0 changed `check_model_inputs` from a decorator factory (`@check_model_inputs()`) to a plain decorator (`@check_model_inputs`, called as `check_model_inputs(func)`) — qwen_tts's CustomVoice tokenizer still uses the old calling convention, so loading Qwen CustomVoice raises `TypeError: check_model_inputs() missing 1 required positional argument: 'func'` once transformers gets silently upgraded. --no-deps avoids this entirely since mlx-lm's own transformers requirement is never consulted; documented this in requirements-mlx.txt so the next person doesn't hit the same trap. 3. `backend/routes/health.py`: `/health`'s `backend_variant` field never had an `mlx` branch, so it always reported "cpu" regardless of whether MLX was active or working — cosmetic, but it's what a couple of the linked issues (#706, #723) used to conclude MLX wasn't running on the GPU when `backend_type` said otherwise. Added the missing branch so `backend_variant` reports "mps" when the mlx backend is selected. All three verified on real Apple Silicon hardware (M4): repeated generations with no thread crash, `just setup-python` now installs a working MLX backend from a clean venv, Qwen CustomVoice downloads and generates successfully post-fix, and `/health` reports `backend_variant: "mps"` when MLX is active. Fixes #699, fixes #823 --- backend/backends/mlx_backend.py | 32 ++++++++++++++++------ backend/requirements-mlx.txt | 48 ++++++++++++++++++++++++++------- backend/routes/health.py | 2 ++ justfile | 8 ++++++ 4 files changed, 72 insertions(+), 18 deletions(-) diff --git a/backend/backends/mlx_backend.py b/backend/backends/mlx_backend.py index 9692e59b4..94ddb69dd 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 @@ -81,8 +97,8 @@ async def load_model_async(self, model_size: Optional[str] = None): if self.model is not None and self._current_model_size != model_size: 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 +274,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 +308,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 +379,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/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 From 7cfa6b86be89c780468ea47f668035d09fb2109b Mon Sep 17 00:00:00 2001 From: Alex Rieker Date: Sun, 12 Jul 2026 21:08:54 +0200 Subject: [PATCH 2/3] fix: enable MPS for the PyTorch Qwen3-TTS/CustomVoice engines on Apple Silicon MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit get_torch_device() already supports allow_mps, and other engines (luxtts_backend.py, qwen_llm_backend.py) already pass allow_mps=True — but PyTorchTTSBackend, PyTorchSTTBackend, and QwenCustomVoiceBackend never did, so any Apple Silicon install that ends up on the PyTorch backend (e.g. MLX unavailable, or before this PR's justfile fix) silently ran these three on plain CPU instead of falling back to PyTorch+MPS. This turned out to matter more than expected: benchmarking on M4 hardware, PyTorch+MPS generation (RTF ~1.7x) was faster than this same PR's MLX thread-fix (RTF ~3.4-4.3x) for the same qwen engine and model size, and faster than a CUDA RTX 2080 Ti (RTF ~2.3-3.1x) too. I'm not proposing changing the default backend selection here — one engine, one model size, one machine is not enough data for that, and MLX-Whisper is reportedly faster than PyTorch-CPU for STT — but MPS should clearly be an option rather than silently absent. Verified on M4: PyTorchTTSBackend()._get_device() and QwenCustomVoiceBackend()._get_device() both now return "mps" (were "cpu"). --- backend/backends/pytorch_backend.py | 4 ++-- backend/backends/qwen_custom_voice_backend.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) 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..6290688a9 100644 --- a/backend/backends/qwen_custom_voice_backend.py +++ b/backend/backends/qwen_custom_voice_backend.py @@ -65,7 +65,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 From e0553aac04b68b0732a593cbeaf4dee599dcb2ba Mon Sep 17 00:00:00 2001 From: Alex Rieker Date: Sun, 12 Jul 2026 22:23:12 +0200 Subject: [PATCH 3/3] =?UTF-8?q?fix:=20address=20CodeRabbit=20review=20find?= =?UTF-8?q?ings=20=E2=80=94=20MPS=20cache=20cleanup,=20MLX=20unload=20thre?= =?UTF-8?q?ad-safety?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two follow-ups from the automated review on this PR: 1. backend/backends/base.py: empty_device_cache() only handled "cuda" and "xpu", never "mps" — so once the previous commit enabled MPS for the PyTorch Qwen3-TTS/CustomVoice engines, switching model sizes on Apple Silicon would leak MPS-cached memory (torch.mps.empty_cache() never got called). Added the missing branch. 2. backend/backends/qwen_custom_voice_backend.py: unload_model() had its own inline `if torch.cuda.is_available(): torch.cuda.empty_cache()` instead of using the shared empty_device_cache() helper — meant it never freed MPS memory either, same bug as #1 but duplicated instead of shared. Switched to the shared helper (matches pytorch_backend.py's existing pattern). 3. backend/backends/mlx_backend.py: load_model_async()'s unload-on- model-size-change call went straight to self.unload_model() instead of routing through _run_on_mlx_thread like every other MLX call in this file (the whole point of the earlier commit in this PR). Fixed for consistency with the thread-pinning fix. Verified on M4: 0.6B -> generate still works post-changes (regression check), torch.mps.empty_cache() confirmed callable without error in this environment. --- backend/backends/base.py | 8 +++++--- backend/backends/mlx_backend.py | 6 ++++-- backend/backends/qwen_custom_voice_backend.py | 4 ++-- 3 files changed, 11 insertions(+), 7 deletions(-) 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 94ddb69dd..c867bdbed 100644 --- a/backend/backends/mlx_backend.py +++ b/backend/backends/mlx_backend.py @@ -93,9 +93,11 @@ 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 on the dedicated MLX thread await _run_on_mlx_thread(self._load_model_sync, model_size) diff --git a/backend/backends/qwen_custom_voice_backend.py b/backend/backends/qwen_custom_voice_backend.py index 6290688a9..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, ) @@ -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")