Add Mage-VL image and video support

README.md

@@ -29,7 +29,7 @@ multi-component export for pipelines.
 |---|---|
 | **Text Generation** | Llama 2/3/4, Mistral, Qwen 2/2.5/3/3.5/3.6, Phi-3/3.5, Gemma 1/2/3/4, Granite, GPT-2, OPT, OLMo, SmolLM3, and many more |
 | **Mixture of Experts** | PhiMoE, GPTOSS, Mixtral, OLMoE, DeepSeek-V2/V3, Qwen2-MoE, Qwen3-MoE, Qwen3-Next, GLM-4-MoE, Arctic, DBRX, Jamba |
-| **Multimodal** | Gemma 3/4, Phi-4MM (vision + audio + LoRA), LLaVA, InternVL2, Qwen2.5-VL, Qwen3-VL, Qwen3.5/3.6-VL, Pixtral |
+| **Multimodal** | Gemma 3/4, Phi-4MM (vision + audio + LoRA), LLaVA, InternVL2, Mage-VL (image + streaming video), Qwen2.5-VL, Qwen3-VL, Qwen3.5/3.6-VL, Pixtral |
 | **Encoder-only** | BERT, RoBERTa, ALBERT, DeBERTa, DistilBERT, ELECTRA, XLNet |
 | **Encoder-Decoder** | BART, T5/mT5, Marian, M2M-100, Pegasus, BigBird-Pegasus |
 | **Speech-to-Text** | Whisper, FastConformer-RNNT, FunASR, Qwen3-ASR, SenseVoice |

scripts/generate_golden.py

@@ -353,7 +353,7 @@ def _generate_seq2seq(case: TestCase, json_path: Path, device: str) -> None:
             input_ids=torch.from_numpy(input_ids).to(torch_device),
             decoder_input_ids=torch.from_numpy(decoder_start).to(torch_device),
         )
-    last_logits = outputs.logits[0, -1, :].cpu().numpy()
+    last_logits = outputs.logits[0, -1, :].float().cpu().numpy()
     golden = _extract_logits_golden(last_logits)
 
     # L5: greedy generation
@@ -413,8 +413,9 @@ def _generate_vision_language(case: TestCase, json_path: Path, device: str) -> N
         case.model_id, dtype=torch_dtype, device=device
     )
 
-    # Load images from testdata/
+    # Load real image/video media from testdata/.
     images = [Image.open(Path("testdata") / img_path) for img_path in case.images]
+    videos = [str(Path("testdata") / video_path) for video_path in case.videos]
 
     # Build a chat-formatted prompt when a usable template is available.
     # Phi-3 Vision exposes its template on the underlying tokenizer rather
@@ -425,6 +426,8 @@ def _generate_vision_language(case: TestCase, json_path: Path, device: str) -> N
         content: list[dict[str, str]] = []
         for img_path in case.images:
             content.append({"type": "image", "image": str(Path("testdata") / img_path)})
+        for video_path in case.videos:
+            content.append({"type": "video", "video": str(Path("testdata") / video_path)})
         content.append({"type": "text", "text": prompt_text})
         messages = [{"role": "user", "content": content}]
         try:
@@ -451,11 +454,37 @@ def _generate_vision_language(case: TestCase, json_path: Path, device: str) -> N
             prompt_text = processor.image_token * len(case.images) + prompt_text
 
     # Process multimodal inputs through the HF processor
-    processed = processor(
-        text=prompt_text,
-        images=images if images else None,
-        return_tensors="pt",
-    )
+    image_processor = getattr(processor, "image_processor", None)
+    video_processor = getattr(processor, "video_processor", None)
+    saved_image_max = getattr(image_processor, "max_pixels", None)
+    image_size = getattr(image_processor, "size", None)
+    saved_image_longest = getattr(image_size, "longest_edge", None)
+    saved_video_max = getattr(video_processor, "max_pixels", None)
+    try:
+        if case.media_max_pixels is not None:
+            if image_processor is not None:
+                image_processor.max_pixels = case.media_max_pixels
+                if image_size is not None and saved_image_longest is not None:
+                    image_size.longest_edge = case.media_max_pixels
+            if video_processor is not None:
+                video_processor.max_pixels = case.media_max_pixels
+        processor_kwargs: dict[str, object] = {
+            "text": prompt_text,
+            "return_tensors": "pt",
+        }
+        if images:
+            processor_kwargs["images"] = images
+        if videos:
+            processor_kwargs["videos"] = videos
+            processor_kwargs["num_frames"] = case.video_num_frames
+        processed = processor(**processor_kwargs)
+    finally:
+        if image_processor is not None and saved_image_max is not None:
+            image_processor.max_pixels = saved_image_max
+        if image_size is not None and saved_image_longest is not None:
+            image_size.longest_edge = saved_image_longest
+        if video_processor is not None and saved_video_max is not None:
+            video_processor.max_pixels = saved_video_max
 
     # Normalize the CLI/device-map selection to a concrete runtime device
     # before moving any tensors. `device="auto"` is handled by Transformers/
@@ -468,7 +497,7 @@ def _generate_vision_language(case: TestCase, json_path: Path, device: str) -> N
     with torch.no_grad():
         outputs = model(**processed)
 
-    last_logits = outputs.logits[0, -1, :].cpu().numpy()
+    last_logits = outputs.logits[0, -1, :].float().cpu().numpy()
     golden = _extract_logits_golden(last_logits)
     input_ids_np = processed["input_ids"].cpu().numpy()
 

src/mobius/__main__.py

@@ -399,6 +399,7 @@ def _save_package(
             hf_model_id=hf_model_id,
             ep=ep,
             local_config_dir=local_config_dir,
+            trust_remote_code=getattr(args, "trust_remote_code", False),
         )
         for name, path in artifacts.items():
             print(f"  {name}: {path}")

src/mobius/_configs/_base.py

@@ -494,8 +494,13 @@ class ArchitectureConfig(BaseModelConfig):
     # Vision shared fields (accessed as top-level config.X by tasks)
     mm_tokens_per_image: int | None = None
     image_token_id: int | None = None
+    video_token_id: int | None = None
+    vision_start_token_id: int | None = None
+    vision_end_token_id: int | None = None
     spatial_merge_size: int = 2
     temporal_patch_size: int = 2
+    frame_windows_size: int = 4
+    tokens_per_second: float = 1.0
     deepstack_visual_indexes: list[int] | None = None
     fullatt_block_indexes: list[int] | None = None
     window_size: int = 112
@@ -982,6 +987,8 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
 
         # Vision config (from multimodal models)
         options.update(_extract_vision_config(config, parent_config, model_type))
+        if getattr(parent_config, "model_type", None) == "mage_vl":
+            options["model_type"] = "mage_vl"
 
         # Audio config
         options.update(_extract_audio_config(config, parent_config, model_type))

src/mobius/_configs/_extractors.py

@@ -182,8 +182,13 @@ def extract_vision_config(config, parent_config, model_type: str) -> dict:
     for shared in (
         "mm_tokens_per_image",
         "image_token_id",
+        "video_token_id",
+        "vision_start_token_id",
+        "vision_end_token_id",
         "spatial_merge_size",
         "temporal_patch_size",
+        "frame_windows_size",
+        "tokens_per_second",
         "deepstack_visual_indexes",
         "fullatt_block_indexes",
         "window_size",

src/mobius/_configs/_sub_configs.py

@@ -47,6 +47,9 @@ class VisionConfig:
     norm_eps: float = 1e-6
     mm_tokens_per_image: int | None = None
     image_token_id: int | None = None
+    video_token_id: int | None = None
+    vision_start_token_id: int | None = None
+    vision_end_token_id: int | None = None
     # Pixtral / Mistral-3 vision fields
     model_type: str | None = None
     head_dim: int | None = None
@@ -56,6 +59,8 @@ class VisionConfig:
     in_channels: int = 3
     spatial_merge_size: int = 2
     temporal_patch_size: int = 2
+    frame_windows_size: int = 4
+    tokens_per_second: float = 1.0
     num_position_embeddings: int | None = None
     deepstack_visual_indexes: list[int] | None = None
     fullatt_block_indexes: list[int] | None = None

src/mobius/_configs/per_model/__init__.py

@@ -29,6 +29,7 @@
     _gemma4_unified_vision,
     _hunyuan_vl_mot_vision,
     _internvl_vision,
+    _mage_vl_vision,
     _phi4mm_audio,
     _phi4mm_vision,
     _phi_vision,

src/mobius/_configs/per_model/_mage_vl_vision.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from mobius._configs._extractors import register_vision_hook
+
+
+@register_vision_hook
+def _mage_vl_vision(config, parent_config, model_type: str, fields: dict):
+    """Extract the custom Mage-ViT configuration from a Mage-VL parent."""
+    del config, model_type
+    if parent_config is None or getattr(parent_config, "model_type", None) != "mage_vl":
+        return None
+
+    vision = getattr(parent_config, "vision_config", None)
+    if vision is None:
+        return None
+
+    fields.update(
+        model_type=getattr(vision, "model_type", "mage_vl_vision"),
+        hidden_size=getattr(vision, "hidden_size", 1024),
+        intermediate_size=getattr(vision, "intermediate_size", 4096),
+        num_hidden_layers=getattr(vision, "num_hidden_layers", 24),
+        num_attention_heads=getattr(vision, "num_attention_heads", 16),
+        image_size=getattr(vision, "image_size", 448),
+        patch_size=getattr(vision, "patch_size", 16),
+        in_channels=getattr(vision, "num_channels", 3),
+        out_hidden_size=getattr(vision, "out_hidden_size", 2560),
+        spatial_merge_size=getattr(vision, "spatial_merge_size", 2),
+        temporal_patch_size=getattr(vision, "temporal_patch_size", 1),
+        frame_windows_size=getattr(vision, "frame_windows_size", 4),
+        norm_eps=getattr(vision, "layer_norm_eps", 1e-6),
+        rope_theta=getattr(vision, "rope_theta", 10_000.0),
+        hidden_act=getattr(vision, "hidden_act", "gelu"),
+        image_token_id=getattr(parent_config, "image_token_id", None),
+        video_token_id=getattr(parent_config, "video_token_id", None),
+        vision_start_token_id=getattr(parent_config, "vision_start_token_id", None),
+        vision_end_token_id=getattr(parent_config, "vision_end_token_id", None),
+        tokens_per_second=getattr(parent_config, "tokens_per_second", 1.0),
+    )
+    return None

src/mobius/_registry.py

@@ -70,6 +70,7 @@
     LayerNormCausalLMModel,
     LLaDAModel,
     Llama4CausalLMModel,
+    MageVLForConditionalGeneration,
     MoECausalLMModel,
     NanoChatCausalLMModel,
     NemotronCausalLMModel,
@@ -592,6 +593,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "internvl": ModelRegistration(InternVL2Model, task="vision-language"),
     "internvl2": ModelRegistration(InternVL2Model, task="vision-language"),
     "internvl_chat": ModelRegistration(InternVL2Model, task="vision-language"),
+    "mage_vl": ModelRegistration(MageVLForConditionalGeneration, task="mage-vl"),
     "janus": ModelRegistration(LLaVAModel, task="vision-language"),
     "llava": ModelRegistration(LLaVAModel, task="vision-language"),
     "llava_next": ModelRegistration(LLaVAModel, task="vision-language"),
@@ -967,6 +969,7 @@ def _create_default_registry() -> ModelRegistry:
     "gemma4_unified": "google/gemma-4-12B",
     "gemma4_unified_text": "google/gemma-4-12B",
     "internvl2": "OpenGVLab/InternVL2-1B",
+    "mage_vl": "microsoft/Mage-VL",
     "phi4mm": "microsoft/Phi-4-multimodal-instruct",
     "phi4_multimodal": "microsoft/Phi-4-multimodal-instruct",
     "phi3_v": "microsoft/Phi-3.5-vision-instruct",

src/mobius/_testing/golden.py

@@ -88,6 +88,15 @@ class GoldenTestCase:
     images: list[str]
     """Image paths relative to ``testdata/`` (VL tasks)."""
 
+    videos: list[str]
+    """Video paths relative to ``testdata/`` (video-language tasks)."""
+
+    video_num_frames: int | None
+    """Optional deterministic number of frames sampled from each video."""
+
+    media_max_pixels: int | None
+    """Optional deterministic pixel budget applied to each image/video frame."""
+
     audio: list[str]
     """Audio paths relative to ``testdata/`` (speech tasks)."""
 
@@ -235,6 +244,9 @@ def load_test_case(yaml_path: Path) -> GoldenTestCase:
         level=data["level"],
         prompts=inputs.get("prompts", []) or [],
         images=inputs.get("images", []) or [],
+        videos=inputs.get("videos", []) or [],
+        video_num_frames=inputs.get("video_num_frames"),
+        media_max_pixels=inputs.get("media_max_pixels"),
         audio=inputs.get("audio", []) or [],
         decoder_prompt=inputs.get("decoder_prompt", "") or "",
         generation_params=generation,

src/mobius/_testing/golden_test.py

@@ -94,6 +94,9 @@ def test_load_minimal(self, tmp_path: Path):
         assert case.level == "L4"
         assert case.prompts == ["Hello world"]
         assert case.images == []
+        assert case.videos == []
+        assert case.video_num_frames is None
+        assert case.media_max_pixels is None
         assert case.audio == []
         assert case.decoder_prompt == ""
         assert case.generation_params == {}

src/mobius/_testing/torch_reference.py

@@ -5,14 +5,52 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
+from pathlib import Path
 
 import numpy as np
 import torch
 
 logger = logging.getLogger(__name__)
 
 
+@contextlib.contextmanager
+def _mage_vl_optional_streammind_import(model_id: str):
+    """Treat StreamMind's mamba-ssm dependency as optional for base Mage-VL.
+
+    The remote ``modeling_mage_vl.py`` imports ``streammind_gate`` only inside
+    StreamMind-specific methods, but Transformers recursively validates that
+    sibling module and otherwise requires mamba-ssm even for ordinary
+    image/video generation. mamba-ssm has no Windows wheel and is not used by
+    the base checkpoint path exercised here.
+    """
+    if model_id.lower() != "microsoft/mage-vl":
+        yield
+        return
+
+    import transformers.dynamic_module_utils as dynamic_module_utils
+
+    original_get_imports = dynamic_module_utils.get_imports
+
+    def _get_imports(filename):
+        imports = original_get_imports(filename)
+        if Path(filename).name == "streammind_gate.py":
+            return [name for name in imports if name != "mamba_ssm"]
+        return imports
+
+    dynamic_module_utils.get_imports = _get_imports
+    try:
+        yield
+    finally:
+        dynamic_module_utils.get_imports = original_get_imports
+
+
+def _load_mage_compatible(model_id: str, loader, *args, **kwargs):
+    with _mage_vl_optional_streammind_import(model_id):
+        return loader(*args, **kwargs)
+
+
 def _install_dynamic_cache_legacy_shims() -> None:
     """Restore ``DynamicCache`` methods removed in transformers 5.x.
 
@@ -229,8 +267,18 @@ def load_torch_multimodal_model(
     """
     import transformers
 
-    tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    processor = transformers.AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+    tokenizer = _load_mage_compatible(
+        model_id,
+        transformers.AutoTokenizer.from_pretrained,
+        model_id,
+        trust_remote_code=True,
+    )
+    processor = _load_mage_compatible(
+        model_id,
+        transformers.AutoProcessor.from_pretrained,
+        model_id,
+        trust_remote_code=True,
+    )
 
     # Shim: transformers 5.x removed DynamicCache.from_legacy_cache and
     # DynamicCache.get_usable_length, but some trust_remote_code models
@@ -241,7 +289,12 @@ def load_torch_multimodal_model(
     # Some models (e.g. Phi-3.5-vision-instruct) hardcode flash_attention_2
     # in their config.json, which causes an ImportError when flash_attn is
     # not installed.
-    config = transformers.AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+    config = _load_mage_compatible(
+        model_id,
+        transformers.AutoConfig.from_pretrained,
+        model_id,
+        trust_remote_code=True,
+    )
     config._attn_implementation = "eager"
 
     # Some trust_remote_code VLMs (e.g. Phi-3-Vision) are registered as
@@ -255,9 +308,21 @@ def load_torch_multimodal_model(
 
     def _load_from_pretrained(auto_cls):
         try:
-            return auto_cls.from_pretrained(model_id, dtype=dtype, **base_kwargs)
+            return _load_mage_compatible(
+                model_id,
+                auto_cls.from_pretrained,
+                model_id,
+                dtype=dtype,
+                **base_kwargs,
+            )
         except TypeError:
-            return auto_cls.from_pretrained(model_id, torch_dtype=dtype, **base_kwargs)
+            return _load_mage_compatible(
+                model_id,
+                auto_cls.from_pretrained,
+                model_id,
+                torch_dtype=dtype,
+                **base_kwargs,
+            )
 
     try:
         image_text_to_text_cls = transformers.AutoModelForImageTextToText

src/mobius/components/__init__.py

@@ -74,6 +74,7 @@
     "VisionEncoderLayer",
     "VisionModel",
     "apply_rms_norm",
+    "build_packed_token_offset",
     "create_attention_bias",
     "create_decoder_layer",
     "create_padding_mask",
@@ -120,6 +121,7 @@
     LayerNormNoBias,
     Linear,
     OffsetLayerNorm,
+    build_packed_token_offset,
     create_attention_bias,
     create_padding_mask,
     create_sliding_window_mask,