diff --git a/README.md b/README.md index f9f2320b6..bd5475220 100644 --- a/README.md +++ b/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 | @@ -38,6 +38,10 @@ multi-component export for pipelines. | **Diffusion** | Stable Diffusion (UNet + VAE + ControlNet), Flux, SD3, DiT, QwenImage, HunyuanDiT, CogVideoX | | **Adapters** | T2I-Adapter, IP-Adapter | +Mage-VL supports direct three-model ONNX export. ORT GenAI export is currently +rejected because the runtime cannot supply its required `patch_positions` input +or Mage-VL's 1D decoder positions. + Supports **290+ Transformers model types** and **10 Diffusers component types** across **40+ task types** and **100+ reusable components**. diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index 95dec62eb..25e497561 100644 --- a/scripts/generate_golden.py +++ b/scripts/generate_golden.py @@ -50,6 +50,41 @@ import numpy as np +_MISSING_ATTRIBUTE = object() + + +@contextlib.contextmanager +def _temporary_processor_max_pixels(processor: object, max_pixels: int | None): + """Temporarily override processor pixel limits and restore their exact state.""" + if max_pixels is None: + yield + return + + saved_attributes: list[tuple[object, str, object]] = [] + + def override(obj: object | None, name: str) -> None: + if obj is None: + return + saved_attributes.append((obj, name, getattr(obj, name, _MISSING_ATTRIBUTE))) + setattr(obj, name, max_pixels) + + image_processor = getattr(processor, "image_processor", None) + video_processor = getattr(processor, "video_processor", None) + override(image_processor, "max_pixels") + image_size = getattr(image_processor, "size", None) + if image_size is not None and hasattr(image_size, "longest_edge"): + override(image_size, "longest_edge") + override(video_processor, "max_pixels") + + try: + yield + finally: + for obj, name, previous in reversed(saved_attributes): + if previous is _MISSING_ATTRIBUTE: + delattr(obj, name) + else: + setattr(obj, name, previous) + def parse_args() -> argparse.Namespace: """Parse command-line arguments.""" @@ -353,7 +388,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 +448,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 +461,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 +489,17 @@ 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", - ) + with _temporary_processor_max_pixels(processor, 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) # Normalize the CLI/device-map selection to a concrete runtime device # before moving any tensors. `device="auto"` is handled by Transformers/ @@ -468,7 +512,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() diff --git a/scripts/generate_golden_test.py b/scripts/generate_golden_test.py index f6cf7ad87..2399894ba 100644 --- a/scripts/generate_golden_test.py +++ b/scripts/generate_golden_test.py @@ -66,6 +66,29 @@ def test_dtype_promotion(self): assert result["logits_summary"].dtype == np.float64 +def test_temporary_processor_max_pixels_restores_none_and_missing_attributes(): + class ProcessorPart: + pass + + image_processor = ProcessorPart() + image_processor.max_pixels = None + image_processor.size = ProcessorPart() + image_processor.size.longest_edge = None + video_processor = ProcessorPart() + processor = ProcessorPart() + processor.image_processor = image_processor + processor.video_processor = video_processor + + with generate_golden._temporary_processor_max_pixels(processor, 1234): + assert image_processor.max_pixels == 1234 + assert image_processor.size.longest_edge == 1234 + assert video_processor.max_pixels == 1234 + + assert image_processor.max_pixels is None + assert image_processor.size.longest_edge is None + assert not hasattr(video_processor, "max_pixels") + + class TestDryRun: """Tests for main() with --dry-run (no HF inference needed).""" diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 821490e8f..c4d0533a9 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -353,6 +353,17 @@ def _save_package( pkg, output_dir: str, args, optimize: str | None, component_filter: str | None ) -> None: """Save a ModelPackage to disk, applying optimizations and runtime configs.""" + runtime = getattr(args, "runtime", None) + if runtime == "ort-genai": + from mobius.integrations.ort_genai.auto_export import ( + _validate_ort_genai_compatibility, + ) + + try: + _validate_ort_genai_compatibility(pkg) + except ValueError as error: + raise SystemExit(f"Error: {error}") from error + components = (lambda name: name == component_filter) if component_filter else None for name, model in pkg.items(): if components is not None and not components(name): @@ -384,7 +395,6 @@ def _save_package( path = os.path.join(output_dir, "model.onnx") print(f"Saved {name} to {path}") - runtime = getattr(args, "runtime", None) if runtime == "ort-genai": from mobius.integrations.ort_genai import write_ort_genai_config @@ -399,6 +409,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}") diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 475b69b9a..e2b5c2e8d 100644 --- a/src/mobius/_configs/_base.py +++ b/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)) diff --git a/src/mobius/_configs/_extractors.py b/src/mobius/_configs/_extractors.py index ac4c6e41c..5f7e14027 100644 --- a/src/mobius/_configs/_extractors.py +++ b/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", diff --git a/src/mobius/_configs/_sub_configs.py b/src/mobius/_configs/_sub_configs.py index 79dfddc5f..aeca2590c 100644 --- a/src/mobius/_configs/_sub_configs.py +++ b/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 diff --git a/src/mobius/_configs/per_model/__init__.py b/src/mobius/_configs/per_model/__init__.py index 28e637637..c92dec766 100644 --- a/src/mobius/_configs/per_model/__init__.py +++ b/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, diff --git a/src/mobius/_configs/per_model/_mage_vl_vision.py b/src/mobius/_configs/per_model/_mage_vl_vision.py new file mode 100644 index 000000000..b2a974490 --- /dev/null +++ b/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 diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 1c611e97c..e59456e3f 100644 --- a/src/mobius/_registry.py +++ b/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", diff --git a/src/mobius/_testing/golden.py b/src/mobius/_testing/golden.py index 271e6afb8..37c4dda88 100644 --- a/src/mobius/_testing/golden.py +++ b/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, diff --git a/src/mobius/_testing/golden_test.py b/src/mobius/_testing/golden_test.py index 6dd26d385..b119125d2 100644 --- a/src/mobius/_testing/golden_test.py +++ b/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 == {} diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index 095e32110..7fb31b3b1 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -5,7 +5,9 @@ from __future__ import annotations +import contextlib import logging +from pathlib import Path import numpy as np import torch @@ -13,6 +15,42 @@ 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 diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 3baeb6e77..6610cd002 100644 --- a/src/mobius/components/__init__.py +++ b/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, diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 3b4786b5d..dd759d6e6 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -117,6 +117,7 @@ "qwen2_vl", "qwen2_5_vl", "qwen3_vl", + "mage_vl", "qwen3_vl_text", "qwen3_5", "qwen3_5_vl", @@ -420,6 +421,7 @@ def _write_vision_processor_config( output_dir: str, *, hf_model_id: str | None = None, + trust_remote_code: bool = False, ) -> str | None: """Write the vision processor config file for VLM models. @@ -443,6 +445,8 @@ def _write_vision_processor_config( - **Pixtral / Mistral3**: Writes ``processor_config.json`` with a 7-step pipeline (DecodeImage → ConvertRGB → Resize → Rescale → Normalize → Permute3D → PixtralImageSizes). + - **Mage-VL**: Writes ``image_processor.json`` with Qwen-style smart resize, + CLIP normalization, and packed patch extraction. - **Other VLMs**: Writes ``processor_config.json`` with a 5-step pipeline (DecodeImage → ConvertRGB → Resize → Rescale → Normalize). @@ -517,7 +521,10 @@ def _write_vision_processor_config( try: from transformers import AutoProcessor - hf_proc = AutoProcessor.from_pretrained(hf_model_id) + hf_proc = AutoProcessor.from_pretrained( + hf_model_id, + trust_remote_code=trust_remote_code, + ) ip = getattr(hf_proc, "image_processor", None) if ip is not None: image_mean = list(getattr(ip, "image_mean", image_mean)) @@ -607,7 +614,10 @@ def _write_vision_processor_config( try: from transformers import AutoProcessor - hf_proc = AutoProcessor.from_pretrained(hf_model_id) + hf_proc = AutoProcessor.from_pretrained( + hf_model_id, + trust_remote_code=trust_remote_code, + ) ip = getattr(hf_proc, "image_processor", None) if ip is not None: image_mean = list(getattr(ip, "image_mean", image_mean)) @@ -701,7 +711,10 @@ def _write_vision_processor_config( "transforms": transforms, } } - path = os.path.join(output_dir, "processor_config.json") + path = os.path.join( + output_dir, + "image_processor.json" if model_type == "mage_vl" else "processor_config.json", + ) with open(path, "w", encoding="utf-8") as f: json.dump(processor_config, f, indent=4) @@ -846,13 +859,12 @@ def _write_genai_config( if is_vlm: image_token_id = getattr(config, "image_token_id", None) if image_token_id is not None: + model_type = getattr(config, "model_type", "") vision_input_mapping = _introspect_inputs(pkg, "vision_encoder") embedding_input_mapping = _introspect_inputs(pkg, "embedding") - # spatial_merge_size and config_filename are config-level # properties that cannot be inferred from the graph. vision_kwargs: dict[str, Any] = {} - model_type = getattr(config, "model_type", "") if model_type in _GEMMA4_MODEL_TYPES: vision_cfg = getattr(config, "vision", None) vision_kwargs["spatial_merge_size"] = getattr( @@ -879,8 +891,12 @@ def _write_genai_config( ) if sms is not None: vision_kwargs["spatial_merge_size"] = sms - vision_kwargs["config_filename"] = "processor_config.json" - if model_type in {"qwen3_vl", "qwen3_vl_text"}: + vision_kwargs["config_filename"] = ( + "image_processor.json" + if model_type == "mage_vl" + else "processor_config.json" + ) + if model_type in {"mage_vl", "qwen3_vl", "qwen3_vl_text"}: patch_size = getattr(vision_cfg, "patch_size", None) window_size = getattr(vision_cfg, "window_size", None) if patch_size is not None: @@ -948,6 +964,17 @@ def _write_genai_config( return generator.write(output_dir) +def _validate_ort_genai_compatibility(pkg: ModelPackage) -> None: + """Reject packages whose required inputs cannot be supplied by ORT GenAI.""" + config = getattr(pkg, "config", None) + if getattr(config, "model_type", None) == "mage_vl": + raise ValueError( + "ORT GenAI does not support Mage-VL's required patch_positions vision " + "input or its 1D decoder position_ids contract. Export without " + "--runtime ort-genai to save the runnable direct three-model ONNX package." + ) + + def write_ort_genai_config( pkg: ModelPackage, directory: str, @@ -956,6 +983,7 @@ def write_ort_genai_config( ep: str = "cpu", context_length: int = 4096, local_config_dir: str | None = None, + trust_remote_code: bool = False, ) -> dict[str, str]: """Generate ORT-GenAI config artifacts for an already-built ModelPackage. @@ -986,6 +1014,8 @@ def write_ort_genai_config( this directory instead of downloaded from HuggingFace Hub. Typically set when the CLI ``--config`` flag points to a local directory rather than a HuggingFace model ID. + trust_remote_code: Allow custom HuggingFace configuration code when + resolving token IDs and model type. Returns: Dict mapping artifact name to file path, e.g.:: @@ -1007,6 +1037,7 @@ def write_ort_genai_config( "This is set automatically when building with mobius.build(). " "Diffusion models (which have no config) are not supported." ) + _validate_ort_genai_compatibility(pkg) os.makedirs(directory, exist_ok=True) @@ -1031,7 +1062,9 @@ def write_ort_genai_config( if hf_model_id is not None: import transformers - hf_config = transformers.AutoConfig.from_pretrained(hf_model_id) + hf_config = transformers.AutoConfig.from_pretrained( + hf_model_id, trust_remote_code=trust_remote_code + ) model_type = hf_config.model_type cfg_model_type = getattr(config, "model_type", None) # See _select_ort_model_type: decoder-only packages prefer the package's @@ -1160,7 +1193,12 @@ def write_ort_genai_config( result[tf] = os.path.join(directory, tf) # Write processor config for VLMs - processor_path = _write_vision_processor_config(config, directory, hf_model_id=hf_model_id) + processor_path = _write_vision_processor_config( + config, + directory, + hf_model_id=hf_model_id, + trust_remote_code=trust_remote_code, + ) if processor_path: result["processor_config"] = processor_path @@ -1187,6 +1225,7 @@ def export_package( ep: str = "cpu", context_length: int = 4096, local_config_dir: str | None = None, + trust_remote_code: bool = False, external_data: str = "onnx", progress_bar: bool = True, ) -> dict[str, str]: @@ -1224,6 +1263,8 @@ def export_package( larger. local_config_dir: Local model directory to copy tokenizer files from when ``hf_model_id`` is ``None``. + trust_remote_code: Allow custom HuggingFace configuration code when + resolving token IDs and model type. external_data: External-data format passed to :meth:`ModelPackage.save` (``"onnx"`` or ``"safetensors"``). progress_bar: Whether to show the save progress bar. @@ -1260,6 +1301,7 @@ def export_package( "Diffusion models (which have no config) are not supported — " "use ModelPackage.save() directly for those." ) + _validate_ort_genai_compatibility(pkg) os.makedirs(output_dir, exist_ok=True) @@ -1279,6 +1321,7 @@ def export_package( ep=ep, context_length=context_length, local_config_dir=local_config_dir, + trust_remote_code=trust_remote_code, ) # 3. Add ONNX paths to the manifest @@ -1385,6 +1428,7 @@ def auto_export( hf_model_id=model_id, ep=ep, context_length=context_length, + trust_remote_code=trust_remote_code, external_data=external_data, progress_bar=progress_bar, ) diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 174db33dd..740b0f1a8 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -26,6 +26,7 @@ _write_genai_config, _write_vision_processor_config, auto_export, + export_package, write_ort_genai_config, ) @@ -190,6 +191,77 @@ def test_writes_transform_pipeline(self, tmp_path): assert len(norm_attrs["mean"]) == 3 assert len(norm_attrs["std"]) == 3 + def test_mage_vl_writes_packed_patch_processor(self, tmp_path): + vision = types.SimpleNamespace( + image_size=448, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=1, + model_type="mage_vl_vision", + ) + config = types.SimpleNamespace( + vision=vision, + model_type="mage_vl", + spatial_merge_size=2, + temporal_patch_size=1, + ) + + path = _write_vision_processor_config(config, str(tmp_path)) + assert path is not None + assert path.endswith("image_processor.json") + with open(path) as f: + data = json.load(f) + + assert data["processor"]["name"] == "qwen2_5_image_processor" + transforms = data["processor"]["transforms"] + assert transforms[-1]["operation"] == { + "name": "patch_image", + "type": "PatchImage", + "attrs": { + "patch_size": 16, + "temporal_patch_size": 1, + "merge_size": 2, + }, + } + normalize = next( + transform["operation"] + for transform in transforms + if transform["operation"]["type"] == "Normalize" + ) + assert normalize["attrs"]["qwen2_5_vl"] == 1 + + def test_mage_vl_processor_propagates_trust_remote_code(self, tmp_path): + vision = types.SimpleNamespace( + image_size=448, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=1, + model_type="mage_vl_vision", + ) + config = types.SimpleNamespace( + vision=vision, + model_type="mage_vl", + spatial_merge_size=2, + temporal_patch_size=1, + ) + hf_processor = mock.MagicMock() + hf_processor.image_processor = None + with mock.patch( + "transformers.AutoProcessor.from_pretrained", + return_value=hf_processor, + ) as from_pretrained: + _write_vision_processor_config( + config, + str(tmp_path), + hf_model_id="microsoft/Mage-VL", + trust_remote_code=True, + ) + + from_pretrained.assert_called_once_with( + "microsoft/Mage-VL", + trust_remote_code=True, + ) + def test_gemma4_unified_skips_image_processor(self, tmp_path): """Encoder-free gemma4_unified has no native transform: no image_processor.json.""" vision = mock.MagicMock() @@ -764,6 +836,48 @@ class FakeConfig: assert model["vision"]["patch_size"] == 16 assert model["vision"]["window_size"] == 64 + def test_mage_vl_is_rejected_before_writing_runtime_artifacts(self, tmp_path): + import dataclasses + + from mobius._model_package import ModelPackage + from mobius.integrations.ort_genai.auto_export import write_ort_genai_config + + @dataclasses.dataclass + class FakeVision: + image_size: int = 448 + patch_size: int = 16 + spatial_merge_size: int = 2 + + @dataclasses.dataclass + class FakeConfig: + model_type: str = "mage_vl" + vocab_size: int = 151936 + hidden_size: int = 2560 + num_hidden_layers: int = 1 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + head_dim: int = 128 + image_token_id: int = 151655 + temporal_patch_size: int = 1 + vision: FakeVision = dataclasses.field(default_factory=FakeVision) + + pkg = ModelPackage( + { + "decoder": mock.MagicMock(), + "vision_encoder": mock.MagicMock(), + "embedding": mock.MagicMock(), + }, + config=FakeConfig(), + ) + + output_dir = tmp_path / "ort-genai" + with pytest.raises( + ValueError, + match=r"Mage-VL.*patch_positions.*1D decoder position_ids", + ): + write_ort_genai_config(pkg, str(output_dir)) + assert not output_dir.exists() + def test_processor_config_not_written_without_vision(self, tmp_path): """image_processor.json is NOT written when pkg.config has no vision attr.""" from mobius.integrations.ort_genai.auto_export import write_ort_genai_config @@ -1001,6 +1115,30 @@ def test_tokenizer_copied_when_model_id_provided(self, tmp_path): mock_copy.assert_called_once_with("fake/model", str(tmp_path)) assert "tokenizer.json" in result + def test_hf_config_propagates_trust_remote_code(self, tmp_path): + """Remote-code models can resolve their HuggingFace configuration.""" + from mobius.integrations.ort_genai.auto_export import write_ort_genai_config + + pkg = self._make_pkg() + with ( + mock.patch( + "mobius.integrations.ort_genai.auto_export._copy_tokenizer_files", + return_value=[], + ), + mock.patch("transformers.AutoConfig.from_pretrained") as mock_hf, + ): + mock_hf.return_value = mock.MagicMock( + model_type="mage_vl", bos_token_id=1, eos_token_id=2, pad_token_id=0 + ) + write_ort_genai_config( + pkg, + str(tmp_path), + hf_model_id="microsoft/Mage-VL", + trust_remote_code=True, + ) + + mock_hf.assert_called_once_with("microsoft/Mage-VL", trust_remote_code=True) + def test_ep_default_normalizes_to_cpu(self, tmp_path): """ep='default' is normalized to cpu (provider_options=[]).""" from mobius.integrations.ort_genai.auto_export import write_ort_genai_config @@ -1220,6 +1358,18 @@ def fake_save(self, directory, **kwargs): # ONNX path is in the manifest (single-component package) assert result["model"] == os.path.join(str(tmp_path), "model.onnx") + def test_mage_vl_is_rejected_before_saving_onnx(self, tmp_path): + pkg = self._make_pkg() + pkg.config.model_type = "mage_vl" + + with ( + mock.patch.object(pkg, "save") as save, + pytest.raises(ValueError, match=r"Mage-VL.*patch_positions"), + ): + export_package(pkg, str(tmp_path)) + + save.assert_not_called() + def test_propagates_save_kwargs(self, tmp_path, monkeypatch): """external_data and progress_bar are forwarded to pkg.save.""" from mobius.integrations.ort_genai.auto_export import export_package @@ -1803,6 +1953,18 @@ def fake_export_package(pkg, output_dir, **kwargs): assert captured["execution_provider"] == "default" assert captured["text_only"] is False + def test_auto_export_rejects_mage_vl_before_saving(self, tmp_path): + pkg = _make_fake_llm_pkg("mage_vl") + + with ( + mock.patch("mobius._builder.build", return_value=pkg), + mock.patch.object(pkg, "save") as save, + pytest.raises(ValueError, match=r"Mage-VL.*patch_positions"), + ): + auto_export("microsoft/Mage-VL", str(tmp_path)) + + save.assert_not_called() + def test_auto_export_produces_genai_config(self, tmp_path): """Mock build() to return a tiny package, verify genai_config.""" import onnx_ir as ir diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 9ca9e2ce3..a5ed6af97 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -69,6 +69,7 @@ "IPAdapterModel", "InternLM2CausalLMModel", "InternVL2Model", + "MageVLForConditionalGeneration", "JambaCausalLMModel", "JetMoeCausalLMModel", "Llama4CausalLMModel", @@ -215,6 +216,7 @@ from mobius.models.llama4 import Llama4CausalLMModel from mobius.models.llava import LLaVAModel from mobius.models.longcat_flash import LongcatFlashCausalLMModel +from mobius.models.mage_vl import MageVLForConditionalGeneration from mobius.models.mamba import Mamba2CausalLMModel, MambaCausalLMModel from mobius.models.mimi import MimiModel, mimi_default_config from mobius.models.minimax import MiniMaxCausalLMModel diff --git a/src/mobius/models/mage_vl.py b/src/mobius/models/mage_vl.py new file mode 100644 index 000000000..b76f900e5 --- /dev/null +++ b/src/mobius/models/mage_vl.py @@ -0,0 +1,604 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Mage-VL streaming image/video-language model with a custom Mage-ViT encoder. + +Replicates Hugging Face ``MageVLForConditionalGeneration`` as a standardized +decoder, vision-encoder, and embedding package. The vision tower consumes packed +Qwen2-VL patches, applies 3D 4:6:6 RoPE at explicit sampled-frame positions, and +limits bidirectional attention to independent four-frame windows. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import onnx_ir as ir +import torch +from onnxscript import OpBuilder, nn + +from mobius._build_context import ep_capabilities, get_build_dtype +from mobius._configs import ArchitectureConfig +from mobius.components import ( + Conv2dNoBias, + Embedding, + LayerNorm, + Linear, + build_packed_token_offset, +) +from mobius.models.base import TextModel + +if TYPE_CHECKING: + from onnx_ir import Value + + +class _GELU(nn.Module): + def forward(self, op: OpBuilder, hidden_states: Value): + return op.Gelu(hidden_states) + + +class MageVLVisionRotaryEmbedding(nn.Module): + """Construct Mage-VL's 4:6:6 temporal/height/width rotary frequencies.""" + + def __init__(self, head_dim: int, rope_theta: float): + super().__init__() + if head_dim % 32 != 0: + raise ValueError( + f"Mage-VL vision head_dim must be divisible by 32, got {head_dim}" + ) + half_dim = head_dim // 2 + unit = half_dim // 16 + self.t_size = 4 * unit + self.h_size = 6 * unit + self.w_size = 6 * unit + self.rope_theta = rope_theta + + def _axis_freqs(self, op: OpBuilder, positions: Value, size: int): + indices = op.Range( + op.Constant(value_int=0), + op.Constant(value_int=size), + op.Constant(value_int=1), + ) + indices = op.Cast(indices, to=ir.DataType.FLOAT) + exponents = op.Div(indices, float(size)) + inv_freq = op.Reciprocal( + op.Pow(op.Constant(value_float=float(self.rope_theta)), exponents) + ) + positions = op.Cast(positions, to=ir.DataType.FLOAT) + return op.Mul(op.Unsqueeze(positions, [-1]), op.Unsqueeze(inv_freq, [0])) + + def forward(self, op: OpBuilder, patch_positions: Value): + # Explicit temporal positions preserve the original sampled video frame indices. + t_pos = op.Gather(patch_positions, op.Constant(value_int=0), axis=1) + h_pos = op.Gather(patch_positions, op.Constant(value_int=1), axis=1) + w_pos = op.Gather(patch_positions, op.Constant(value_int=2), axis=1) + half_freqs = op.Concat( + self._axis_freqs(op, t_pos, self.t_size), + self._axis_freqs(op, h_pos, self.h_size), + self._axis_freqs(op, w_pos, self.w_size), + axis=-1, + ) + # The reference concatenates D/2 frequencies with themselves before + # applying its interleaved rotate-half implementation. + freqs = op.Concat(half_freqs, half_freqs, axis=-1) # (patches, head_dim) + freqs = op.Unsqueeze(freqs, [0]) # (1, patches, head_dim) + return op.Cos(freqs), op.Sin(freqs) + + +class MageVLVisionAttention(nn.Module): + """Fused-QKV bidirectional attention matching MageVLVisionAttention.""" + + def __init__(self, hidden_size: int, num_heads: int): + super().__init__() + self.num_heads = num_heads + self.scale = (hidden_size // num_heads) ** -0.5 + self.qkv = Linear(hidden_size, 3 * hidden_size) + self.proj = Linear(hidden_size, hidden_size) + + def _apply_rope(self, op: OpBuilder, x: Value, cos: Value, sin: Value): + # Mage-VL's reference duplicates the complete D/2 frequency vector, + # then rotates adjacent even/odd channels. This differs from the ONNX + # RotaryEmbedding interleaved contract, which shares one frequency per + # pair, so reproduce the reference explicitly in FP32. + x_shape = op.Shape(x) + x_heads = op.Reshape( + x, + op.Concat( + op.Shape(x, start=0, end=2), + op.Constant(value_ints=[self.num_heads, -1]), + axis=0, + ), + ) + x_float = op.Cast(x_heads, to=ir.DataType.FLOAT) + even = op.Slice(x_float, starts=[0], ends=[9223372036854775807], axes=[3], steps=[2]) + odd = op.Slice(x_float, starts=[1], ends=[9223372036854775807], axes=[3], steps=[2]) + rotated = op.Reshape( + op.Concat(op.Unsqueeze(op.Neg(odd), [-1]), op.Unsqueeze(even, [-1]), axis=-1), + op.Shape(x_float), + ) + cos = op.Unsqueeze(cos, [2]) + sin = op.Unsqueeze(sin, [2]) + embedded = op.Add(op.Mul(x_float, cos), op.Mul(rotated, sin)) + return op.Reshape(op.CastLike(embedded, x), x_shape) + + def forward( + self, + op: OpBuilder, + hidden_states: Value, + attention_mask: Value | None, + cu_seqlens: Value, + cos: Value, + sin: Value, + ): + q, k, v = op.Split( + self.qkv(op, hidden_states), + num_outputs=3, + axis=-1, + _outputs=3, + ) + q = self._apply_rope(op, q, cos, sin) + k = self._apply_rope(op, k, cos, sin) + if ep_capabilities().supports_packed_multi_head_attention: + query = op.Squeeze(q, [0]) + key = op.Squeeze(k, [0]) + value = op.Squeeze(v, [0]) + token_offset = build_packed_token_offset(op, cu_seqlens) + if get_build_dtype() == ir.DataType.BFLOAT16: + query_mha = op.Cast(query, to=ir.DataType.FLOAT16) + key_mha = op.Cast(key, to=ir.DataType.FLOAT16) + value_mha = op.Cast(value, to=ir.DataType.FLOAT16) + else: + query_mha, key_mha, value_mha = query, key, value + hidden_states = op.PackedMultiHeadAttention( + query_mha, + key_mha, + value_mha, + None, + token_offset, + op.Cast(cu_seqlens, to=ir.DataType.INT32), + num_heads=self.num_heads, + scale=self.scale, + _domain="com.microsoft", + _outputs=1, + ) + hidden_states = op.Unsqueeze(op.CastLike(hidden_states, query), [0]) + else: + hidden_states = op.Attention( + q, + k, + v, + attention_mask, + q_num_heads=self.num_heads, + kv_num_heads=self.num_heads, + scale=self.scale, + is_causal=0, + _outputs=1, + ) + return self.proj(op, hidden_states) + + +class MageVLVisionMLP(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.fc1 = Linear(hidden_size, intermediate_size) + self.fc2 = Linear(intermediate_size, hidden_size) + + def forward(self, op: OpBuilder, hidden_states: Value): + return self.fc2(op, op.Gelu(self.fc1(op, hidden_states))) + + +class MageVLVisionEncoderLayer(nn.Module): + """Pre-norm Mage-ViT block with fused-QKV attention and GELU MLP.""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + num_heads: int, + norm_eps: float, + ): + super().__init__() + self.self_attn = MageVLVisionAttention(hidden_size, num_heads) + self.layer_norm1 = LayerNorm(hidden_size, eps=norm_eps) + self.mlp = MageVLVisionMLP(hidden_size, intermediate_size) + self.layer_norm2 = LayerNorm(hidden_size, eps=norm_eps) + + def forward( + self, + op: OpBuilder, + hidden_states: Value, + attention_mask: Value | None, + cu_seqlens: Value, + cos: Value, + sin: Value, + ): + residual = hidden_states + hidden_states = self.self_attn( + op, + self.layer_norm1(op, hidden_states), + attention_mask, + cu_seqlens, + cos, + sin, + ) + hidden_states = op.Add(residual, hidden_states) + residual = hidden_states + hidden_states = self.mlp(op, self.layer_norm2(op, hidden_states)) + return op.Add(residual, hidden_states) + + +class MageVLVisionEncoder(nn.Module): + def __init__( + self, + num_layers: int, + hidden_size: int, + intermediate_size: int, + num_heads: int, + norm_eps: float, + ): + super().__init__() + self.layers = nn.ModuleList( + [ + MageVLVisionEncoderLayer( + hidden_size, + intermediate_size, + num_heads, + norm_eps, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + op: OpBuilder, + hidden_states: Value, + attention_mask: Value | None, + cu_seqlens: Value, + cos: Value, + sin: Value, + ): + for layer in self.layers: + hidden_states = layer(op, hidden_states, attention_mask, cu_seqlens, cos, sin) + return hidden_states + + +class MageVLVisionEmbeddings(nn.Module): + """Bias-free Conv2d projection for pre-extracted Qwen2-VL patches.""" + + def __init__(self, in_channels: int, hidden_size: int, patch_size: int): + super().__init__() + self.in_channels = in_channels + self.hidden_size = hidden_size + self.patch_size = patch_size + self.patch_embedding = Conv2dNoBias( + in_channels, + hidden_size, + kernel_size=patch_size, + stride=patch_size, + ) + + def forward(self, op: OpBuilder, hidden_states: Value): + patches = op.Reshape( + hidden_states, + op.Constant(value_ints=[-1, self.in_channels, self.patch_size, self.patch_size]), + ) + hidden_states = self.patch_embedding(op, patches) # (patches, hidden, 1, 1) + return op.Reshape( + hidden_states, + op.Constant(value_ints=[1, -1, self.hidden_size]), + ) # (1, patches, hidden) + + +class MageVLVisionPatchMerger(nn.Module): + """Merge each contiguous 2x2 packed patch group into one text-width token.""" + + def __init__( + self, + output_size: int, + context_size: int, + spatial_merge_size: int, + norm_eps: float, + ): + super().__init__() + merged_size = context_size * spatial_merge_size**2 + self.ln_q = LayerNorm(context_size, eps=norm_eps) + self.mlp = nn.Sequential( + Linear(merged_size, merged_size), + _GELU(), + Linear(merged_size, output_size), + ) + self.merged_size = merged_size + + def forward(self, op: OpBuilder, hidden_states: Value): + hidden_states = self.ln_q(op, hidden_states) + hidden_states = op.Reshape( + hidden_states, + op.Constant(value_ints=[-1, self.merged_size]), + ) + return self.mlp(op, hidden_states) + + +class MageVLVisionPretrainedModel(nn.Module): + """Custom Mage-ViT backbone used by microsoft/Mage-VL.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + vc = config.vision + if vc is None or vc.hidden_size is None: + raise ValueError("Mage-VL requires a complete vision configuration") + if vc.num_hidden_layers is None or vc.num_attention_heads is None: + raise ValueError("Mage-VL vision layer and head counts are required") + hidden_size = vc.hidden_size + num_heads = vc.num_attention_heads + norm_eps = vc.norm_eps + self.frame_windows_size = vc.frame_windows_size + self.embeddings = MageVLVisionEmbeddings( + vc.in_channels, + hidden_size, + vc.patch_size or 16, + ) + self.layernorm_pre = LayerNorm(hidden_size, eps=norm_eps) + self.encoder = MageVLVisionEncoder( + vc.num_hidden_layers, + hidden_size, + vc.intermediate_size or 4 * hidden_size, + num_heads, + norm_eps, + ) + self.video_rope = MageVLVisionRotaryEmbedding( + hidden_size // num_heads, + vc.rope_theta or 10_000.0, + ) + self.merger = MageVLVisionPatchMerger( + vc.out_hidden_size or config.hidden_size, + hidden_size, + vc.spatial_merge_size, + norm_eps, + ) + + def _attention_metadata(self, op: OpBuilder, grid_thw: Value, hidden_states: Value): + # Map each flattened patch to its visual sample and four-frame chunk. + total = op.Shape(hidden_states, start=1, end=2) + patch_ids = op.Range( + op.Constant(value_int=0), + op.Squeeze(total, [0]), + op.Constant(value_int=1), + ) + sample_lengths = op.ReduceProd(grid_thw, axes=[1], keepdims=0) + sample_ends = op.CumSum(sample_lengths, op.Constant(value_int=0)) + sample_ids = op.ReduceSum( + op.Cast( + op.GreaterOrEqual( + op.Unsqueeze(patch_ids, [1]), + op.Unsqueeze(sample_ends, [0]), + ), + to=ir.DataType.INT64, + ), + axes=[1], + keepdims=0, + ) + sample_starts = op.Concat( + op.Constant(value_ints=[0]), + op.Slice(sample_ends, starts=[0], ends=[-1], axes=[0]), + axis=0, + ) + local_ids = op.Sub(patch_ids, op.Gather(sample_starts, sample_ids)) + spatial_sizes = op.Mul( + op.Gather(grid_thw, op.Constant(value_int=1), axis=1), + op.Gather(grid_thw, op.Constant(value_int=2), axis=1), + ) + window_sizes = op.Mul( + op.Gather(spatial_sizes, sample_ids), + op.Constant(value_int=self.frame_windows_size), + ) + window_ids = op.Div(local_ids, window_sizes) + same_as_previous = op.And( + op.Equal( + op.Slice(sample_ids, starts=[1], ends=[9223372036854775807], axes=[0]), + op.Slice(sample_ids, starts=[0], ends=[-1], axes=[0]), + ), + op.Equal( + op.Slice(window_ids, starts=[1], ends=[9223372036854775807], axes=[0]), + op.Slice(window_ids, starts=[0], ends=[-1], axes=[0]), + ), + ) + segment_starts = op.Concat( + op.Cast(op.Constant(value_ints=[1]), to=ir.DataType.BOOL), + op.Not(same_as_previous), + axis=0, + ) + cu_seqlens = op.Concat( + op.Compress(patch_ids, segment_starts), + total, + axis=0, + ) + if ep_capabilities().supports_packed_multi_head_attention: + return None, cu_seqlens + + segment_ids = op.CumSum( + op.Cast(segment_starts, to=ir.DataType.INT64), + op.Constant(value_int=0), + ) + attention_mask = op.Equal( + op.Unsqueeze(segment_ids, [1]), + op.Unsqueeze(segment_ids, [0]), + ) + return op.Unsqueeze(attention_mask, [0, 1]), cu_seqlens + + def forward( + self, + op: OpBuilder, + hidden_state: Value, + grid_thw: Value, + patch_positions: Value, + ): + hidden_states = self.embeddings(op, hidden_state) + cos, sin = self.video_rope(op, patch_positions) + attention_mask, cu_seqlens = self._attention_metadata(op, grid_thw, hidden_states) + hidden_states = self.layernorm_pre(op, hidden_states) + hidden_states = self.encoder( + op, + hidden_states, + attention_mask, + cu_seqlens, + cos, + sin, + ) + return self.merger(op, hidden_states) + + +class _MageVLModelBody(nn.Module): + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.language_model = TextModel(config) + + def forward( + self, + op: OpBuilder, + inputs_embeds: Value, + attention_mask: Value, + position_ids: Value, + past_key_values: list | None, + ): + return self.language_model( + op, + input_ids=None, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + ) + + +class MageVLDecoderModel(nn.Module): + """Qwen3 decoder with the checkpoint's ``model.language_model`` hierarchy.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.model = _MageVLModelBody(config) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward( + self, + op: OpBuilder, + inputs_embeds: Value, + attention_mask: Value, + position_ids: Value, + past_key_values: list | None = None, + ): + hidden_states, present_key_values = self.model( + op, + inputs_embeds, + attention_mask, + position_ids, + past_key_values, + ) + return self.lm_head(op, hidden_states), present_key_values + + +class _MageVLVisualBody(nn.Module): + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.visual = MageVLVisionPretrainedModel(config) + + +class MageVLVisionEncoderModel(nn.Module): + """Standalone Mage-ViT encoder for image and sampled-video patches.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.model = _MageVLVisualBody(config) + + def forward( + self, + op: OpBuilder, + pixel_values: Value, + image_grid_thw: Value, + patch_positions: Value, + ): + return self.model.visual(op, pixel_values, image_grid_thw, patch_positions) + + +class MageVLEmbeddingModel(nn.Module): + """Scatter packed image/video features at Mage-VL visual placeholder tokens.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + config.pad_token_id, + ) + if config.image_token_id is None: + raise ValueError("Mage-VL requires image_token_id") + self.image_token_id = config.image_token_id + self.video_token_id = config.video_token_id + + def forward(self, op: OpBuilder, input_ids: Value, image_features: Value): + text_embeddings = self.embed_tokens(op, input_ids) + visual_mask = op.Equal(input_ids, op.Constant(value_int=self.image_token_id)) + if self.video_token_id is not None: + visual_mask = op.Or( + visual_mask, + op.Equal(input_ids, op.Constant(value_int=self.video_token_id)), + ) + flat_visual_mask = op.Reshape( + visual_mask, + op.Constant(value_ints=[-1]), + ) + flat_indices = op.CumSum( + op.Cast(flat_visual_mask, to=ir.DataType.INT64), + op.Constant(value_int=0), + ) + indices = op.Reshape(flat_indices, op.Shape(input_ids)) + flat_text = op.Reshape( + text_embeddings, + op.Constant(value_ints=[-1, self.embed_tokens.weight.shape[1]]), + ) + zero_feature = op.Mul( + op.Slice(flat_text, starts=[0], ends=[1], axes=[0]), + 0.0, + ) + padded_features = op.Concat(zero_feature, image_features, axis=0) + visual_embeddings = op.Gather(padded_features, indices, axis=0) + return op.Where( + op.Unsqueeze(visual_mask, [-1]), + visual_embeddings, + text_embeddings, + ) + + +class MageVLForConditionalGeneration(nn.Module): + """Mage-VL streaming image/video-language model with a Qwen3 decoder.""" + + default_task: str = "mage-vl" + category: str = "Multimodal" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.config = config + self.decoder = MageVLDecoderModel(config) + self.vision_encoder = MageVLVisionEncoderModel(config) + self.embedding = MageVLEmbeddingModel(config) + + def forward(self, op: OpBuilder, **kwargs): + raise NotImplementedError("MageVLTask exports each pipeline component separately") + + def preprocess_weights( + self, + state_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Route every Hugging Face checkpoint tensor to its package component.""" + renamed: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if key.startswith("model.language_model."): + if key == "model.language_model.embed_tokens.weight": + renamed["embedding.embed_tokens.weight"] = value + else: + renamed[f"decoder.{key}"] = value + elif key.startswith("model.visual."): + renamed[f"vision_encoder.{key[len('model.') :]}"] = value + elif key.startswith("lm_head."): + renamed[f"decoder.{key}"] = value + return renamed diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index d7f947ffc..6a8bd328b 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -46,6 +46,7 @@ "ImageClassificationTask", "ModelTask", "MllamaVisionLanguageTask", + "MageVLTask", "MaskedDiffusionTask", "MoshiDepformerTask", "MoshiTemporalTask", @@ -124,6 +125,7 @@ from mobius.tasks._vision_language_3model import ( Cosmos3EdgeVLTask, HybridQwenVLTask, + MageVLTask, MllamaVisionLanguageTask, PixtralVLTask, QwenVLTask, @@ -163,6 +165,7 @@ "cosmos3-edge-vl": Cosmos3EdgeVLTask, "pixtral-vl": PixtralVLTask, "mllama-vision-language": MllamaVisionLanguageTask, + "mage-vl": MageVLTask, "qwen-vl": QwenVLTask, "hybrid-qwen-vl": HybridQwenVLTask, "qwen3-vl-vision-language": Qwen3VLVisionLanguageTask, diff --git a/src/mobius/tasks/_vision_language_3model.py b/src/mobius/tasks/_vision_language_3model.py index 42d804ed9..1fd5d0035 100644 --- a/src/mobius/tasks/_vision_language_3model.py +++ b/src/mobius/tasks/_vision_language_3model.py @@ -219,6 +219,46 @@ def _build_vision( return _make_model(graph) +class MageVLTask(VisionLanguageTask): + """Mage-VL split with packed patches and explicit sampled-frame positions.""" + + def _build_vision( + self, + vision: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + total_patches = ir.SymbolicDim("total_patches") + num_visuals = ir.SymbolicDim("num_visuals") + patch_size = (config.vision.patch_size if config.vision else None) or 16 + in_channels = config.vision.in_channels if config.vision else 3 + pixel_dim = in_channels * patch_size * patch_size + + graph, builder = _make_graph(name="vision_encoder") + pixel_values = builder.input( + "pixel_values", + dtype=config.dtype, + shape=[total_patches, pixel_dim], + ) + image_grid_thw = builder.input( + "image_grid_thw", + dtype=ir.DataType.INT64, + shape=[num_visuals, 3], + ) + patch_positions = builder.input( + "patch_positions", + dtype=ir.DataType.INT64, + shape=[total_patches, 3], + ) + image_features = vision( + builder.op, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + patch_positions=patch_positions, + ) + builder.add_output(image_features, "image_features") + return _make_model(graph) + + class HybridQwenVLTask(QwenVLTask): """Qwen VL 3-model split with hybrid KV + DeltaNet cache. diff --git a/testdata/cases/schema.json b/testdata/cases/schema.json index 418278824..99fc99b7f 100644 --- a/testdata/cases/schema.json +++ b/testdata/cases/schema.json @@ -104,6 +104,25 @@ "minItems": 0, "description": "Image file paths relative to testdata/ (vision and vision-language tasks)." }, + "videos": { + "type": "array", + "items": { + "type": "string", + "minLength": 1 + }, + "minItems": 0, + "description": "Video file paths relative to testdata/ (video-language tasks)." + }, + "video_num_frames": { + "type": "integer", + "minimum": 1, + "description": "Optional deterministic frame count for video-language preprocessing." + }, + "media_max_pixels": { + "type": "integer", + "minimum": 1, + "description": "Optional deterministic per-image/frame pixel budget for multimodal golden tests." + }, "audio": { "type": "array", "items": { diff --git a/testdata/cases/vision-language/mage-vl.yaml b/testdata/cases/vision-language/mage-vl.yaml new file mode 100644 index 000000000..5dbee839b --- /dev/null +++ b/testdata/cases/vision-language/mage-vl.yaml @@ -0,0 +1,25 @@ +model_id: "microsoft/Mage-VL" +model_type: "mage_vl" +revision: "d88b153285f1633a61b2f693c59c8576693af185" +task_type: "image-text-to-text" +dtype: "bfloat16" +trust_remote_code: true + +inputs: + prompts: + - "Describe the image and then summarize the main action in the video." + images: + - "mage-vl-dog.jpg" + videos: + - "mage-vl-soccer-broadcast.mp4" + video_num_frames: 5 + media_max_pixels: 200704 + +level: "L4+L5" + +generation: + max_new_tokens: 24 + do_sample: false + +ci_skip_reason: "Mage-VL is a 4.7B remote-code model; real-media golden validation requires a large GPU/CPU-offload runner." +notes: "Microsoft Mage-VL with its upstream Apache-2.0 dog image and soccer video; exercises mixed image/video ordering and explicit sampled-frame patch positions." diff --git a/testdata/golden/vision-language/mage-vl.json b/testdata/golden/vision-language/mage-vl.json new file mode 100644 index 000000000..77769fe71 --- /dev/null +++ b/testdata/golden/vision-language/mage-vl.json @@ -0,0 +1,1184 @@ +{ + "top1_id": 785, + "top2_id": 32, + "top10_ids": [ + 785, + 32, + 1782, + 27, + 64, + 34024, + 66755, + 4064, + 52670, + 26972 + ], + "top10_logits": [ + "0x1.b400000000000p+3", + "0x1.9400000000000p+3", + "0x1.9200000000000p+3", + "0x1.9000000000000p+3", + "0x1.6e00000000000p+3", + "0x1.6400000000000p+3", + "0x1.6400000000000p+3", + "0x1.5c00000000000p+3", + "0x1.5200000000000p+3", + "0x1.4a00000000000p+3" + ], + "logits_summary": [ + "0x1.b400000000000p+3", + "-0x1.d400000000000p+3", + "-0x1.5f2fa87e4eef6p+1", + "0x1.4888c01384aafp+1" + ], + "input_ids": [ + 151644, + 8948, + 198, + 2610, + 525, + 264, + 10950, + 17847, + 13, + 151645, + 198, + 151644, + 872, + 198, + 151652, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151653, + 27, + 15, + 13, + 15, + 6486, + 29, + 151652, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151653, + 27, + 22, + 13, + 20, + 6486, + 29, + 151652, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151653, + 27, + 16, + 20, + 13, + 15, + 6486, + 29, + 151652, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151653, + 27, + 17, + 17, + 13, + 20, + 6486, + 29, + 151652, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151653, + 27, + 18, + 15, + 13, + 15, + 6486, + 29, + 151652, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151655, + 151653, + 74785, + 279, + 2168, + 323, + 1221, + 62079, + 279, + 1887, + 1917, + 304, + 279, + 2766, + 13, + 151645, + 198, + 151644, + 77091, + 198 + ] +} diff --git a/testdata/golden/vision-language/mage-vl_generation.json b/testdata/golden/vision-language/mage-vl_generation.json new file mode 100644 index 000000000..91211afe6 --- /dev/null +++ b/testdata/golden/vision-language/mage-vl_generation.json @@ -0,0 +1,31 @@ +{ + "model_id": "microsoft/Mage-VL", + "prompt": "Describe the image and then summarize the main action in the video.", + "generated_tokens": [ + 785, + 2168, + 61891, + 264, + 1874, + 315, + 3040, + 2953, + 11259, + 2163, + 264, + 1965, + 304, + 264, + 23889, + 11, + 1817, + 9963, + 264, + 13753, + 42395, + 448, + 279, + 18096 + ], + "generated_text": "The image depicts a group of four men standing around a table in a stadium, each holding a yellow microphone with the BBC" +} diff --git a/testdata/mage-vl-dog.jpg b/testdata/mage-vl-dog.jpg new file mode 100644 index 000000000..8fc3a3cba Binary files /dev/null and b/testdata/mage-vl-dog.jpg differ diff --git a/testdata/mage-vl-soccer-broadcast.mp4 b/testdata/mage-vl-soccer-broadcast.mp4 new file mode 100644 index 000000000..590a8ccc6 Binary files /dev/null and b/testdata/mage-vl-soccer-broadcast.mp4 differ diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 19ea16689..3e8ef5ae3 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -2166,6 +2166,37 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: ("internvl_chat", {"vision": _TINY_VISION, "image_token_id": 32000}, True), ("internvl2", {"vision": _TINY_VISION, "image_token_id": 32000}, False), ("internvl", {"vision": _TINY_VISION, "image_token_id": 32000}, False), + ( + "mage_vl", + { + "attn_qk_norm": True, + "image_token_id": 100, + "video_token_id": 101, + "vision_start_token_id": 102, + "vision_end_token_id": 103, + "vision": VisionConfig( + model_type="mage_vl_vision", + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=2, + image_size=8, + patch_size=4, + norm_eps=1e-6, + in_channels=3, + out_hidden_size=TINY_HIDDEN, + spatial_merge_size=2, + temporal_patch_size=1, + frame_windows_size=4, + rope_theta=10_000.0, + hidden_act="gelu", + ), + "spatial_merge_size": 2, + "temporal_patch_size": 1, + "frame_windows_size": 4, + }, + True, + ), # --- Gemma3 multimodal (requires rope_local_base_freq, layer_types) --- ( "gemma3", diff --git a/tests/cli_test.py b/tests/cli_test.py index 7a3e7cf6e..363ee1b65 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -453,6 +453,59 @@ def test_runtime_ort_genai_calls_write_ort_genai_config(self): mock_export.assert_called_once() call_kwargs = mock_export.call_args assert call_kwargs.kwargs.get("hf_model_id") == "Qwen/Qwen2.5-0.5B" + assert call_kwargs.kwargs.get("trust_remote_code") is False + + def test_runtime_ort_genai_propagates_trust_remote_code(self): + """--trust-remote-code also applies to runtime config generation.""" + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "mobius.integrations.ort_genai.write_ort_genai_config", + return_value={}, + ) as mock_export, + ): + main( + [ + "build", + "--model", + "Qwen/Qwen2.5-0.5B", + tmpdir, + "--no-weights", + "--trust-remote-code", + "--runtime", + "ort-genai", + ] + ) + + assert mock_export.call_args.kwargs["trust_remote_code"] is True + + def test_runtime_ort_genai_rejects_mage_vl_before_saving(self): + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch("mobius._model_package.ModelPackage.save") as save, + mock.patch( + "mobius.integrations.ort_genai.write_ort_genai_config" + ) as config_writer, + pytest.raises( + SystemExit, + match=r"Mage-VL.*patch_positions.*1D decoder position_ids", + ), + ): + main( + [ + "build", + "--model", + "microsoft/Mage-VL", + tmpdir, + "--no-weights", + "--trust-remote-code", + "--runtime", + "ort-genai", + ] + ) + + save.assert_not_called() + config_writer.assert_not_called() def test_runtime_onnx_genai_uses_native_vlm_emitter(self): pkg = mock.MagicMock() diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index 37ea7ca2b..97299f91c 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -24,6 +24,7 @@ from __future__ import annotations +import contextlib import dataclasses import functools import os @@ -64,6 +65,40 @@ def _load_phi3_vision_projector_weights_cached(model_id: str): # Root of test data (images, audio, etc.) _TESTDATA_DIR = Path(__file__).resolve().parent.parent / "testdata" +_MISSING_ATTRIBUTE = object() + + +@contextlib.contextmanager +def _temporary_processor_max_pixels(processor: object, max_pixels: int | None): + """Temporarily override processor pixel limits and restore their exact state.""" + if max_pixels is None: + yield + return + + saved_attributes: list[tuple[object, str, object]] = [] + + def override(obj: object | None, name: str) -> None: + if obj is None: + return + saved_attributes.append((obj, name, getattr(obj, name, _MISSING_ATTRIBUTE))) + setattr(obj, name, max_pixels) + + image_processor = getattr(processor, "image_processor", None) + video_processor = getattr(processor, "video_processor", None) + override(image_processor, "max_pixels") + image_size = getattr(image_processor, "size", None) + if image_size is not None and hasattr(image_size, "longest_edge"): + override(image_size, "longest_edge") + override(video_processor, "max_pixels") + + try: + yield + finally: + for obj, name, previous in reversed(saved_attributes): + if previous is _MISSING_ATTRIBUTE: + delattr(obj, name) + else: + setattr(obj, name, previous) def _get_test_build_ep() -> str: @@ -226,6 +261,29 @@ def test_phi3_vision_projector_weights_are_cached(): _load_phi3_vision_projector_weights_cached.cache_clear() +def test_temporary_processor_max_pixels_restores_none_and_missing_attributes(): + class ProcessorPart: + pass + + image_processor = ProcessorPart() + image_processor.max_pixels = None + image_processor.size = ProcessorPart() + image_processor.size.longest_edge = None + video_processor = ProcessorPart() + processor = ProcessorPart() + processor.image_processor = image_processor + processor.video_processor = video_processor + + with _temporary_processor_max_pixels(processor, 1234): + assert image_processor.max_pixels == 1234 + assert image_processor.size.longest_edge == 1234 + assert video_processor.max_pixels == 1234 + + assert image_processor.max_pixels is None + assert image_processor.size.longest_edge is None + assert not hasattr(video_processor, "max_pixels") + + def _make_empty_kv_cache( session: OnnxModelSession, config: object, @@ -834,13 +892,25 @@ def _run_vl_vision_to_image_features( for name in vis_session.input_names: if name in processed: val = processed[name] - vis_feeds[name] = val if isinstance(val, np.ndarray) else np.array(val) + array = val if isinstance(val, np.ndarray) else np.array(val) + target_dtype = vis_session.get_input_dtype(name) + vis_feeds[name] = ( + array.astype(target_dtype) + if target_dtype is not None and array.dtype != target_dtype + else array + ) else: # Handle HF↔ONNX name mismatches (e.g. HF "image_position_ids" # vs ONNX "pixel_position_ids"). for hf_key, val in processed.items(): if hf_key.replace("image_", "pixel_") == name: - vis_feeds[name] = val if isinstance(val, np.ndarray) else np.array(val) + array = val if isinstance(val, np.ndarray) else np.array(val) + target_dtype = vis_session.get_input_dtype(name) + vis_feeds[name] = ( + array.astype(target_dtype) + if target_dtype is not None and array.dtype != target_dtype + else array + ) break vis_out = vis_session.run(vis_feeds) return vis_out, vis_out[next(iter(vis_out))] @@ -848,6 +918,46 @@ def _run_vl_vision_to_image_features( vis_session.close() +def _prepare_vl_inputs( + case: GoldenTestCase, + processor: object, +) -> dict[str, np.ndarray]: + """Apply an HF multimodal processor to the case's ordered image/video media.""" + from PIL import Image + + images = [Image.open(_TESTDATA_DIR / path) for path in case.images] + videos = [str(_TESTDATA_DIR / path) for path in case.videos] + prompt_text = case.prompts[0] + if videos: + content: list[dict[str, str]] = [ + *[{"type": "image", "image": str(_TESTDATA_DIR / path)} for path in case.images], + *[{"type": "video", "video": str(_TESTDATA_DIR / path)} for path in case.videos], + {"type": "text", "text": prompt_text}, + ] + prompt_text = processor.apply_chat_template( + [{"role": "user", "content": content}], + tokenize=False, + add_generation_prompt=True, + ) + else: + prompt_text = _build_mm_prompt(processor, prompt_text, case.images, "image") + + kwargs: dict[str, object] = { + "text": prompt_text, + "images": images or None, + "return_tensors": "pt", + } + if videos: + kwargs["videos"] = videos + kwargs["num_frames"] = case.video_num_frames + with _temporary_processor_max_pixels(processor, case.media_max_pixels): + processed_pt = processor(**kwargs) + return { + key: value.numpy() if hasattr(value, "numpy") else np.array(value) + for key, value in processed_pt.items() + } + + def _run_vision_language_prefill( pkg: ModelPackage, case: GoldenTestCase, @@ -864,22 +974,12 @@ def _run_vision_language_prefill( golden generation. """ import transformers - from PIL import Image - # --- Step 0: Preprocess image with HF processor --- + # --- Step 0: Preprocess image/video media with the HF processor --- processor = transformers.AutoProcessor.from_pretrained( case.model_id, trust_remote_code=case.trust_remote_code ) - image = Image.open(_TESTDATA_DIR / case.images[0]) - - # Build the prompt (chat template when available, else manual placeholder). - prompt_text = _build_mm_prompt(processor, case.prompts[0], case.images, "image") - - # Use PyTorch tensors then convert — some processors don't support np - processed_pt = processor(text=prompt_text, images=[image], return_tensors="pt") - processed: dict[str, np.ndarray] = { - k: v.numpy() if hasattr(v, "numpy") else np.array(v) for k, v in processed_pt.items() - } + processed = _prepare_vl_inputs(case, processor) # --- Step 1: Run vision encoder (+ host-side projector where required) --- vis_out, image_features = _run_vl_vision_to_image_features(pkg, case, processed) @@ -1029,21 +1129,13 @@ def _run_vl_generation( Returns newly generated token IDs (prompt excluded). """ import transformers - from PIL import Image # --- Step 0: prepare multimodal inputs --- processor = transformers.AutoProcessor.from_pretrained( case.model_id, trust_remote_code=case.trust_remote_code ) - image = Image.open(_TESTDATA_DIR / case.images[0]) - - prompt_text = _build_mm_prompt(processor, case.prompts[0], case.images, "image") suppress_ids = _load_suppress_token_ids(case.model_id, case.trust_remote_code) - - processed_pt = processor(text=prompt_text, images=[image], return_tensors="pt") - processed: dict[str, np.ndarray] = { - k: v.numpy() if hasattr(v, "numpy") else np.array(v) for k, v in processed_pt.items() - } + processed = _prepare_vl_inputs(case, processor) # --- Step 1: vision encoder (+ host-side projector where required) --- _vis_out, image_features = _run_vl_vision_to_image_features(pkg, case, processed) diff --git a/tests/mage_vl_parity_test.py b/tests/mage_vl_parity_test.py new file mode 100644 index 000000000..981bf9913 --- /dev/null +++ b/tests/mage_vl_parity_test.py @@ -0,0 +1,240 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Synthetic end-to-end parity for the Mage-VL video vision path.""" + +from __future__ import annotations + +import dataclasses + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import torch +from _test_configs import VL_CONFIGS, _base_config +from torch.nn import functional + +from mobius import build_from_module +from mobius._registry import registry +from mobius._testing import count_op_type +from mobius._weight_loading import apply_weights +from mobius.tasks import get_task + + +def _linear(x: torch.Tensor, state: dict[str, torch.Tensor], prefix: str) -> torch.Tensor: + return functional.linear(x, state[f"{prefix}.weight"], state.get(f"{prefix}.bias")) + + +def _layer_norm( + x: torch.Tensor, + state: dict[str, torch.Tensor], + prefix: str, +) -> torch.Tensor: + return functional.layer_norm( + x, + (x.shape[-1],), + state[f"{prefix}.weight"], + state[f"{prefix}.bias"], + 1e-6, + ) + + +def _reference_vision( + pixel_values: torch.Tensor, + grid_thw: torch.Tensor, + patch_positions: torch.Tensor, + state: dict[str, torch.Tensor], +) -> torch.Tensor: + prefix = "vision_encoder.visual" + patch_weight = state[f"{prefix}.embeddings.patch_embedding.weight"] + hidden = functional.conv2d(pixel_values.reshape(-1, 3, 4, 4), patch_weight, stride=4) + hidden = hidden.reshape(1, -1, 64) + + # Mage-VL uses independent frequency scales for its 4:6:6 T/H/W split. + axis_freqs = [] + for axis, size in enumerate((4, 6, 6)): + inv_freq = 1.0 / (10_000.0 ** (torch.arange(size, dtype=torch.float32) / size)) + axis_freqs.append(patch_positions[:, axis].float().unsqueeze(1) * inv_freq) + half_freqs = torch.cat(axis_freqs, dim=-1) + freqs = torch.cat((half_freqs, half_freqs), dim=-1).unsqueeze(0).unsqueeze(2) + cos, sin = freqs.cos(), freqs.sin() + + segment_ids = [] + for sample_id, (t, h, w) in enumerate(grid_thw.tolist()): + patches_per_window = 4 * h * w + for local_id in range(t * h * w): + segment_ids.append((sample_id, local_id // patches_per_window)) + mask = torch.tensor( + [[left == right for right in segment_ids] for left in segment_ids], + dtype=torch.bool, + ) + + hidden = _layer_norm(hidden, state, f"{prefix}.layernorm_pre") + for layer_idx in range(2): + layer = f"{prefix}.encoder.layers.{layer_idx}" + residual = hidden + normed = _layer_norm(hidden, state, f"{layer}.layer_norm1") + qkv = _linear(normed, state, f"{layer}.self_attn.qkv") + q, k, v = qkv.chunk(3, dim=-1) + + def _rope(x: torch.Tensor, target_dtype: torch.dtype = qkv.dtype) -> torch.Tensor: + x = x.reshape(1, x.shape[1], 2, 32).float() + even, odd = x[..., ::2], x[..., 1::2] + rotated = torch.stack((-odd, even), dim=-1).flatten(-2) + return (x * cos + rotated * sin).to(target_dtype) + + q = _rope(q).transpose(1, 2) + k = _rope(k).transpose(1, 2) + v = v.reshape(1, v.shape[1], 2, 32).transpose(1, 2) + scores = torch.matmul(q, k.transpose(-1, -2)) * (32**-0.5) + scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min) + attended = torch.matmul(torch.softmax(scores, dim=-1), v) + attended = attended.transpose(1, 2).reshape(1, -1, 64) + hidden = residual + _linear(attended, state, f"{layer}.self_attn.proj") + + residual = hidden + hidden = _layer_norm(hidden, state, f"{layer}.layer_norm2") + hidden = _linear(hidden, state, f"{layer}.mlp.fc1") + hidden = functional.gelu(hidden) + hidden = residual + _linear(hidden, state, f"{layer}.mlp.fc2") + + hidden = _layer_norm(hidden, state, f"{prefix}.merger.ln_q").reshape(-1, 256) + hidden = functional.gelu(_linear(hidden, state, f"{prefix}.merger.mlp.0")) + return _linear(hidden, state, f"{prefix}.merger.mlp.2") + + +@pytest.mark.parametrize( + ("dtype", "torch_dtype", "atol"), + [ + (ir.DataType.FLOAT, torch.float32, 1e-4), + (ir.DataType.FLOAT16, torch.float16, 1e-2), + ], +) +def test_mage_vl_synthetic_video_parity(tmp_path, dtype, torch_dtype, atol): + """A nonzero five-frame clip matches PyTorch across the four-frame boundary.""" + overrides = next(overrides for mt, overrides, _ in VL_CONFIGS if mt == "mage_vl") + config = dataclasses.replace(_base_config(**overrides), dtype=dtype) + package = build_from_module( + registry.get("mage_vl")(config), + config, + task="mage-vl", + ) + vision = package["vision_encoder"] + + generator = torch.Generator().manual_seed(0) + state = { + name: (torch.randn(tuple(value.shape), generator=generator) * 0.02).to(torch_dtype) + for name, value in vision.graph.initializers.items() + if value.const_value is None + } + pixel_values = torch.randn((20, 48), generator=generator).to(torch_dtype) + grid_thw = torch.tensor([[5, 2, 2]], dtype=torch.int64) + patch_positions = torch.tensor( + [ + (frame, height, width) + for frame in (0, 3, 7, 12, 18) + for height in range(2) + for width in range(2) + ], + dtype=torch.int64, + ) + expected = _reference_vision(pixel_values, grid_thw, patch_positions, state) + + apply_weights(vision, state) + model_path = tmp_path / "mage_vl_vision.onnx" + ir.save(vision, model_path) + session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + (actual,) = session.run( + None, + { + "pixel_values": pixel_values.numpy(), + "image_grid_thw": grid_thw.numpy(), + "patch_positions": patch_positions.numpy(), + }, + ) + np.testing.assert_allclose(actual, expected.detach().numpy(), atol=atol, rtol=atol) + + +def test_mage_vl_decode_embedding_accepts_no_new_media(tmp_path): + """Single-token decode can run with an empty packed feature tensor.""" + overrides = next(overrides for mt, overrides, _ in VL_CONFIGS if mt == "mage_vl") + config = _base_config(**overrides) + embedding = get_task("mage-vl").build(registry.get("mage_vl")(config), config)["embedding"] + state = { + name: torch.randn(tuple(value.shape), generator=torch.Generator().manual_seed(1)) + for name, value in embedding.graph.initializers.items() + if value.const_value is None + } + apply_weights(embedding, state) + model_path = tmp_path / "mage_vl_embedding.onnx" + ir.save(embedding, model_path) + session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + (actual,) = session.run( + None, + { + "input_ids": np.array([[7]], dtype=np.int64), + "image_features": np.empty((0, config.hidden_size), dtype=np.float32), + }, + ) + np.testing.assert_allclose( + actual, + state["embedding.embed_tokens.weight"][[7]].numpy().reshape(1, 1, -1), + ) + + +def test_mage_vl_cuda_uses_packed_attention(): + """CUDA builds preserve four-frame segments without a quadratic dense mask.""" + overrides = next(overrides for mt, overrides, _ in VL_CONFIGS if mt == "mage_vl") + config = _base_config(**overrides) + package = build_from_module( + registry.get("mage_vl")(config), + config, + task="mage-vl", + execution_provider="cuda", + ) + vision = package["vision_encoder"] + + assert count_op_type(vision.graph, "PackedMultiHeadAttention") == 2 + assert count_op_type(vision.graph, "Attention") == 0 + + +def test_mage_vl_embedding_uses_global_media_order_across_batch(tmp_path): + """Packed visual features continue across prompt rows instead of restarting.""" + overrides = next(overrides for mt, overrides, _ in VL_CONFIGS if mt == "mage_vl") + config = _base_config(**overrides) + embedding = get_task("mage-vl").build(registry.get("mage_vl")(config), config)["embedding"] + state = { + name: torch.randn(tuple(value.shape), generator=torch.Generator().manual_seed(2)) + for name, value in embedding.graph.initializers.items() + if value.const_value is None + } + apply_weights(embedding, state) + model_path = tmp_path / "mage_vl_batched_embedding.onnx" + ir.save(embedding, model_path) + session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + input_ids = np.array( + [ + [config.image_token_id, 7], + [8, config.video_token_id], + ], + dtype=np.int64, + ) + image_features = np.stack( + [ + np.full(config.hidden_size, 11.0, dtype=np.float32), + np.full(config.hidden_size, 22.0, dtype=np.float32), + ] + ) + (actual,) = session.run( + None, + { + "input_ids": input_ids, + "image_features": image_features, + }, + ) + + np.testing.assert_allclose(actual[0, 0], image_features[0]) + np.testing.assert_allclose(actual[1, 1], image_features[1]) + np.testing.assert_allclose(actual[0, 1], state["embedding.embed_tokens.weight"][7].numpy()) + np.testing.assert_allclose(actual[1, 0], state["embedding.embed_tokens.weight"][8].numpy()) diff --git a/tests/weight_alignment_test.py b/tests/weight_alignment_test.py index 39ae516a2..2913162f6 100644 --- a/tests/weight_alignment_test.py +++ b/tests/weight_alignment_test.py @@ -33,6 +33,7 @@ ENCODER_CONFIGS, SEQ2SEQ_CONFIGS, VISION_CONFIGS, + VL_CONFIGS, _base_config, ) @@ -239,3 +240,31 @@ class TestDetectionWeightAlignment: def test_identity_state_dict_roundtrip(self, model_type: str, config_overrides: dict): _assert_identity_roundtrip(model_type, config_overrides) + + +def test_mage_vl_huggingface_checkpoint_alignment(): + """Every Mage-VL package parameter is populated from its exact HF key.""" + config_overrides = next(overrides for mt, overrides, _ in VL_CONFIGS if mt == "mage_vl") + config = _base_config(**config_overrides) + module = registry.get("mage_vl")(config) + pkg = get_task("mage-vl").build(module, config) + parameter_names = _collect_parameter_names(pkg) + + hf_state: dict[str, torch.Tensor] = {} + for name in parameter_names: + if name == "embedding.embed_tokens.weight": + continue + if name.startswith("decoder.model.language_model."): + hf_name = name[len("decoder.") :] + elif name.startswith("decoder.lm_head."): + hf_name = name[len("decoder.") :] + elif name.startswith("vision_encoder.visual."): + hf_name = f"model.{name[len('vision_encoder.') :]}" + else: + raise AssertionError(f"Unrecognized Mage-VL parameter name: {name}") + hf_state[hf_name] = torch.ones(1) + + # The shared source embedding must populate both decoder and embedding models. + hf_state["model.language_model.embed_tokens.weight"] = torch.ones(1) + aligned = module.preprocess_weights(hf_state) + assert set(aligned) == parameter_names