fix: MLX thread affinity crash + missing mlx-lm/mlx-audio in dev setup

backend/backends/base.py

@@ -171,17 +171,19 @@ 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
 
     if device == "cuda" and torch.cuda.is_available():
         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:

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)

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."""

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")
 

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.

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",

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