fix

app.py

@@ -17,7 +17,10 @@ import hashlib
 from slugify import slugify
 import argparse 
 import importlib
+import importlib.machinery
+import importlib.util
 import sys
+import types
 from pathlib import Path
 import tempfile
 import zipfile
@@ -32,6 +35,7 @@ training_script_url = f"https://raw.githubusercontent.com/huggingface/diffusers/
 orchestrator_script_url = "https://huggingface.co/datasets/multimodalart/lora-ease-helper/raw/main/script.py"
 
 MAGE_VL_MODEL_ID = "microsoft/Mage-VL"
+MAGE_VL_REVISION = "d88b153285f1633a61b2f693c59c8576693af185"
 mage_vl_processor = None
 mage_vl_model = None
 caption_cache = {}
@@ -73,13 +77,17 @@ def get_captioner():
     """Load the embedded Mage-VL model once per Space process."""
     global mage_vl_processor, mage_vl_model
     if mage_vl_processor is None or mage_vl_model is None:
+        _install_image_only_mamba_shim()
         target_device = "cuda" if is_spaces or torch.cuda.is_available() else "cpu"
         dtype = torch.bfloat16 if target_device == "cuda" else torch.float32
         mage_vl_processor = AutoProcessor.from_pretrained(
-            MAGE_VL_MODEL_ID, trust_remote_code=True
+            MAGE_VL_MODEL_ID,
+            revision=MAGE_VL_REVISION,
+            trust_remote_code=True,
         )
         mage_vl_model = AutoModelForCausalLM.from_pretrained(
             MAGE_VL_MODEL_ID,
+            revision=MAGE_VL_REVISION,
             trust_remote_code=True,
             dtype=dtype,
             attn_implementation="sdpa",
@@ -87,6 +95,41 @@ def get_captioner():
     return mage_vl_processor, mage_vl_model
 
 
+def _install_image_only_mamba_shim():
+    """Satisfy Mage-VL's optional StreamMind import without compiling Mamba.
+
+    Transformers scans every relative remote-code file, including the lazily
+    loaded StreamMind gate. Still-image captioning never constructs that gate,
+    so requiring its CUDA extension would add a large, unnecessary build step.
+    """
+    if importlib.util.find_spec("mamba_ssm") is not None:
+        return
+
+    def module(name, package=False):
+        value = types.ModuleType(name)
+        value.__spec__ = importlib.machinery.ModuleSpec(
+            name, loader=None, is_package=package
+        )
+        if package:
+            value.__path__ = []
+        sys.modules[name] = value
+        return value
+
+    mamba_module = module("mamba_ssm", package=True)
+    models_module = module("mamba_ssm.models", package=True)
+    mixer_module = module("mamba_ssm.models.mixer_seq_simple")
+
+    def unavailable_create_block(*args, **kwargs):
+        raise RuntimeError(
+            "Mage-VL's optional StreamMind gate requires mamba-ssm; this Space "
+            "embeds Mage-VL for still-image captioning only."
+        )
+
+    mixer_module.create_block = unavailable_create_block
+    mamba_module.models = models_module
+    models_module.mixer_seq_simple = mixer_module
+
+
 def _mage_prompt(processor, instruction):
     messages = [
         {