fix(lyrics): fallback to transformers/MPS when MLX fails

MLX load failure and the worker-thread stream crash both now fall
back to Qwen 2.5 7B via transformers on MPS. Fallback at load time
catches missing/incompatible MLX weights; fallback at generate time
catches the RuntimeError: There is no Stream(gpu, 0) bug. The global
singleton is swapped on first failure so subsequent calls skip MLX.

lyrics_lm.py

@@ -24,12 +24,15 @@ import re
 from dataclasses import dataclass
 from typing import Any
 
+import logging
+
 import ace_pipeline as ap
 
 _DEFAULT_MAC_ID = "mlx-community/Qwen2.5-7B-Instruct-4bit"
 _DEFAULT_CUDA_ID = "Qwen/Qwen2.5-7B-Instruct"
 
 _LM = None  # lazy module-level singleton
+_log = logging.getLogger("ams.lyrics")
 
 
 def build_system_prompt() -> str:
@@ -98,30 +101,36 @@ def _get_lm():
     return _LM
 
 
+def _load_hflm(device: str) -> "_HFLM":
+    """Load Qwen 2.5 7B via transformers on the given device (mps/cuda/cpu)."""
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    dtype = torch.bfloat16 if device in ("cuda", "mps") else torch.float32
+    tok = AutoTokenizer.from_pretrained(_DEFAULT_CUDA_ID)
+    model = AutoModelForCausalLM.from_pretrained(_DEFAULT_CUDA_ID, torch_dtype=dtype).to(device)
+    return _HFLM(model=model, tokenizer=tok)
+
+
 def _load_lm():
     """Construct the per-device LM wrapper.
 
-    On MPS we use ``mlx-lm`` which expects a model ID and returns
-    ``(model, tokenizer)``. On CUDA / CPU we use ``transformers`` with
-    ``apply_chat_template`` for the prompt.
+    On MPS, try mlx-lm first (4-bit, fast). If MLX fails to load, fall back
+    to transformers on MPS. On CUDA/CPU use transformers directly.
     """
     device = ap.detect_device()
     if device == "mps":
-        from mlx_lm import load  # type: ignore[import-not-found]
-
-        model, tokenizer = load(_DEFAULT_MAC_ID)
-        return _MLXLM(model=model, tokenizer=tokenizer)
+        try:
+            from mlx_lm import load  # type: ignore[import-not-found]
 
-    # CUDA / CPU fallback path. Use bfloat16 on CUDA, float32 on CPU.
-    import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+            model, tokenizer = load(_DEFAULT_MAC_ID)
+            return _MLXLM(model=model, tokenizer=tokenizer)
+        except Exception as exc:
+            _log.warning("MLX load failed (%s); falling back to transformers on MPS", exc)
+            return _load_hflm("mps")
 
-    tok = AutoTokenizer.from_pretrained(_DEFAULT_CUDA_ID)
-    dtype = torch.bfloat16 if device == "cuda" else torch.float32
-    model = AutoModelForCausalLM.from_pretrained(_DEFAULT_CUDA_ID, torch_dtype=dtype)
-    if device == "cuda":
-        model = model.to("cuda")
-    return _HFLM(model=model, tokenizer=tok)
+    # CUDA / CPU path.
+    return _load_hflm(device)
 
 
 @dataclass
@@ -136,36 +145,25 @@ class _MLXLM:
         import mlx_lm.generate as mlx_gen_mod  # type: ignore[import-not-found]
         from mlx_lm import generate  # type: ignore[import-not-found]
 
-        # Qwen's ChatML template — mlx-lm doesn't expose apply_chat_template
-        # the way HF does, so build the prompt manually here.
         prompt = (
             f"<|im_start|>system\n{system}<|im_end|>\n"
             f"<|im_start|>user\n{user}<|im_end|>\n"
             f"<|im_start|>assistant\n"
         )
-        # Gradio runs handlers in anyio worker threads. MLX maintains a
-        # *per-thread* default stream and a module-level ``generation_stream``
-        # that was created at mlx_lm import time on the MAIN thread. Both
-        # need to be valid in the *current* (worker) thread or
-        # ``wired_limit().__exit__`` crashes with "There is no Stream(gpu, 0)
-        # in current thread" when it calls ``mx.synchronize(generation_stream)``.
-        #
-        # Two-part fix:
-        #   1. ``mx.stream(mx.gpu)`` wrap installs the default GPU stream
-        #      for the current thread for the duration of the call.
-        #   2. Re-assign ``mlx_lm.generate.generation_stream`` to a stream
-        #      created in the *current* thread so ``mx.synchronize`` doesn't
-        #      reach across thread boundaries. The reassignment is safe
-        #      because Gradio's queue runs at default_concurrency_limit=1 —
-        #      no two lyrics drafts run concurrently.
-        with mx.stream(mx.gpu):
-            mlx_gen_mod.generation_stream = mx.new_stream(mx.default_device())
-            return generate(
-                self.model,
-                self.tokenizer,
-                prompt=prompt,
-                max_tokens=int(kw.get("max_new_tokens", 600)),
-            )
+        try:
+            with mx.stream(mx.gpu):
+                mlx_gen_mod.generation_stream = mx.new_stream(mx.default_device())
+                return generate(
+                    self.model,
+                    self.tokenizer,
+                    prompt=prompt,
+                    max_tokens=int(kw.get("max_new_tokens", 600)),
+                )
+        except RuntimeError as exc:
+            _log.warning("MLX generate failed (%s); switching to transformers on MPS", exc)
+            global _LM
+            _LM = _load_hflm("mps")
+            return _LM.generate(system, user, **kw)
 
 
 @dataclass