Labyricorn/voicebox
1
0
Fork 0
mirror of https://github.com/jamiepine/voicebox.git synced 2026-10-03 17:15:19 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix(mlx): fold Qwen3 LLM load+generate into one worker submission; bind model locally in generate closures

Browse Source 
Review follow-ups on the thread-affinity PR: MLXQwenLLMBackend.generate
awaited load_model as one worker submission and then submitted
_generate_sync as a second, leaving the same load/generate gap the TTS
and STT paths close (a concurrent load_model for another size or an
unload could swap or null out self.model/self.tokenizer in between).
Give the LLM backend the same _op_lock + single _reload_and_generate_sync
shape. The TTS/STT/LLM sync closures now bind the model to a local once,
so an inline unload_model() from the event loop mid-generation cannot
turn a later self.model read (the voice-clone fallback path in
particular) into an AttributeError.
This commit is contained in:
jamiepine
2026-10-04 00:01:47 +00:00
committed by capy-ai-staging[bot]
parent cf984885e9
commit dd8ab5cb20
2 changed files with 41 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/backends/qwen_llm_backend.py
+27 -10
View File
@@ -190,6 +190,9 @@ class MLXQwenLLMBackend:
self.tokenizer = None
self.model_size = model_size
self._current_model_size: Optional[str] = None
# Same role as MLXTTSBackend._op_lock: keeps two coroutines from
# racing their reload decisions around one load-then-generate.
self._op_lock = asyncio.Lock()
def is_loaded(self) -> bool:
return self.model is not None
@@ -217,10 +220,14 @@ class MLXQwenLLMBackend:
# one thread (see mlx_backend._run_on_mlx_thread / issue #699).
from .mlx_backend import _run_on_mlx_thread
if self.model is not None and self._current_model_size != model_size:
await _run_on_mlx_thread(self.unload_model)
async with self._op_lock:
await _run_on_mlx_thread(self._reload_sync, model_size)
await _run_on_mlx_thread(self._load_model_sync, model_size)
def _reload_sync(self, model_size: str) -> None:
"""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()
self._load_model_sync(model_size)
def _load_model_sync(self, model_size: str) -> None:
from mlx_lm import load as mlx_load
@@ -262,12 +269,21 @@ class MLXQwenLLMBackend:
model_size: Optional[str] = None,
examples: Optional[list[tuple[str, str]]] = None,
) -> str:
await self.load_model(model_size)
from .mlx_backend import _run_on_mlx_thread
return await _run_on_mlx_thread(
self._generate_sync, prompt, system, max_tokens, temperature, examples
)
resolved_size = model_size if model_size is not None else self.model_size
def _reload_and_generate_sync() -> str:
# One worker submission for load + generate, as in
# MLXTTSBackend._reload_and_generate_sync: no gap between the two
# where another request's load_model (different size) or an
# unload can swap or null out self.model / self.tokenizer.
if self.model is None or self._current_model_size != resolved_size:
self._reload_sync(resolved_size)
return self._generate_sync(prompt, system, max_tokens, temperature, examples)
async with self._op_lock:
return await _run_on_mlx_thread(_reload_and_generate_sync)
def _generate_sync(
self,
@@ -280,8 +296,9 @@ class MLXQwenLLMBackend:
from mlx_lm import generate as mlx_generate
from mlx_lm.sample_utils import make_sampler
model, tokenizer = self.model, self.tokenizer
messages = _build_messages(prompt, system, examples)
chat_prompt = self.tokenizer.apply_chat_template(
chat_prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
@@ -290,8 +307,8 @@ class MLXQwenLLMBackend:
sampler = make_sampler(temp=temperature, top_p=0.9) if temperature > 0 else None
text = mlx_generate(
self.model,
self.tokenizer,
model,
tokenizer,
prompt=chat_prompt,
max_tokens=max_tokens,
sampler=sampler,