diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 5e66c71e147b..aaced448e647 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -1437,6 +1437,8 @@ title: mllama - local: model_doc/mm-grounding-dino title: MM Grounding DINO + - local: model_doc/molmo2 + title: Molmo2 - local: model_doc/musicflamingo title: MusicFlamingo - local: model_doc/neomme diff --git a/docs/source/en/model_doc/molmo2.md b/docs/source/en/model_doc/molmo2.md new file mode 100644 index 000000000000..2ffc2fc58cd2 --- /dev/null +++ b/docs/source/en/model_doc/molmo2.md @@ -0,0 +1,202 @@ + +*This model was published in HF papers on 2026-01-15 and contributed to Hugging Face Transformers on 2026-10-08.* + +
+
+FlashAttention +SDPA
+
+ +# Molmo2 + +[Molmo2](https://huggingface.co/papers/2601.10611) is a family of open-weight vision-language models by AllenAI that are state-of-the-art among open-source models, with exceptional capabilities in point-driven grounding for single image, multi-image, and video tasks. The architecture combines a Vision Transformer (ViT) for image processing with an adapter layer connecting vision and text modalities, and a text decoder based on transformer architecture with rotary position embeddings. + +The abstract from the paper is the following: + +*Today's strongest video-language models (VLMs) remain proprietary. The strongest open-weight models either rely on synthetic data from proprietary VLMs, effectively distilling from them, or do not disclose their training data or recipe. As a result, the open-source community lacks the foundations needed to improve on the state-of-the-art video (and image) language models. Crucially, many downstream applications require more than just high-level video understanding; they require grounding -- either by pointing or by tracking in pixels. Even proprietary models lack this capability. We present Molmo2, a new family of VLMs that are state-of-the-art among open-source models and demonstrate exceptional new capabilities in point-driven grounding in single image, multi-image, and video tasks. Our key contribution is a collection of 7 new video datasets and 2 multi-image datasets, including a dataset of highly detailed video captions for pre-training, a free-form video Q&A dataset for fine-tuning, a new object tracking dataset with complex queries, and an innovative new video pointing dataset, all collected without the use of closed VLMs. We also present a training recipe for this data utilizing an efficient packing and message-tree encoding scheme, and show bi-directional attention on vision tokens and a novel token-weight strategy improves performance. Our best-in-class 8B model outperforms others in the class of open weight and data models on short videos, counting, and captioning, and is competitive on long-videos. On video-grounding Molmo2 significantly outperforms existing open-weight models like Qwen3-VL (35.5 vs 29.6 accuracy on video counting) and surpasses proprietary models like Gemini 3 Pro on some tasks (38.4 vs 20.0 F1 on video pointing and 56.2 vs 41.1 J&F on video tracking).* + +You can find all the original Molmo2 checkpoints under the [Molmo2](https://huggingface.co/collections/allenai/molmo2-67d6b5b0e138c5d621de1e5d) collection. + +## Usage example + +### Image-text-to-text generation + +The example below demonstrates how to generate text based on an image with [`Pipeline`] or the [`AutoModel`] class. + + + + +```python +from transformers import pipeline + +pipeline = pipeline( + task="image-text-to-text", + model="allenai/Molmo2-8B", + device_map="auto", +) + +messages = [ + { + "role": "user", + "content": [ + {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"}, + {"type": "text", "text": "Describe this image."}, + ], + } +] +pipeline(text=messages, max_new_tokens=128, return_full_text=False) +``` + + + + +```python +from transformers import AutoModelForImageTextToText, AutoProcessor + +processor = AutoProcessor.from_pretrained("allenai/Molmo2-8B") +model = AutoModelForImageTextToText.from_pretrained( + "allenai/Molmo2-8B", + device_map="auto", +) + +messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image."}, + {"type": "image", "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"}, + ], + } +] + +inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, +) + +generated_ids = model.generate(**inputs, max_new_tokens=128) +generated_text = processor.decode( + generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True +) +print(generated_text) +# In this captivating image, a snow leopard is captured ... +``` + + + + +### Video pointing + +Molmo2 can also generate grounded points for video inputs. The points are returned in the generated text as a +`` payload, where each coordinate group contains a timestamp followed by point ids and normalized +image coordinates. + +```python +import torch + +from transformers import AutoModelForImageTextToText, AutoProcessor + + +model_id = "allenai/Molmo2-8B" +video_url = "https://storage.googleapis.com/oe-training-public/demo_videos/many_penguins.mp4" + +processor = AutoProcessor.from_pretrained(model_id) +model = AutoModelForImageTextToText.from_pretrained(model_id, device_map="auto") + +messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Point to the penguins."}, + {"type": "video", "video": video_url}, + ], + } +] +inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, +) + +generated_ids = model.generate(**inputs, max_new_tokens=128, do_sample=False) + +input_len = inputs["input_ids"].shape[1] +generated_text = processor.decode(generated_ids[0][input_len:], skip_special_tokens=True) +print(generated_text) +``` + +## Molmo2Config + +[[autodoc]] Molmo2Config + +## Molmo2VisionConfig + +[[autodoc]] Molmo2VisionConfig + +## Molmo2AdapterConfig + +[[autodoc]] Molmo2AdapterConfig + +## Molmo2TextConfig + +[[autodoc]] Molmo2TextConfig + +## Molmo2Processor + +[[autodoc]] Molmo2Processor + - __call__ + +## Molmo2ImageProcessor + +[[autodoc]] Molmo2ImageProcessor + - __call__ + - preprocess + +## Molmo2VideoProcessor + +[[autodoc]] Molmo2VideoProcessor + - __call__ + +## Molmo2Model + +[[autodoc]] Molmo2Model + - forward + +## Molmo2TextModel + +[[autodoc]] Molmo2TextModel + - forward + +## Molmo2Adapter + +[[autodoc]] Molmo2Adapter + - forward + +## Molmo2VisionModel + +[[autodoc]] Molmo2VisionModel + - forward + +## Molmo2ForConditionalGeneration + +[[autodoc]] Molmo2ForConditionalGeneration + - forward diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index b16d58e71f32..5e3cf0be76b1 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -27,6 +27,7 @@ MergeModulelist, PermuteForRope, PrefixChange, + Split, Transpose, WeightConverter, WeightRenaming, @@ -804,6 +805,70 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(source_patterns=r"^vision_model", target_patterns="model.vision_model"), WeightRenaming(source_patterns=r"^multi_modal_projector", target_patterns="model.multi_modal_projector"), ], + "molmo2": [ + WeightRenaming( + source_patterns=r"(? ConversionOps: return Concatenate(self.dim, self.num_shards_attribute) +class Split(ConversionOps): + """Split a tensor along `dim` into one piece per target, sized like that target's model parameter. + WARNING this op is not TP-safe, a sharded fused tensor needs each target's local slice, not the full parameter sizes.""" + + def __init__(self, dim: int = 0): + self.dim = dim + + @torch.no_grad + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + full_layer_name: str, + model: PreTrainedModel, + **kwargs, + ) -> dict[str, torch.Tensor]: + if len(input_dict) > 1: + raise ValueError("Undefined Operation encountered") + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + sizes = [ + model.get_parameter(full_layer_name.replace(target_patterns[0], target)).shape[self.dim] + for target in target_patterns + ] + chunks = torch.split(tensor, sizes, dim=self.dim) + return {target: chunk.contiguous() for target, chunk in zip(target_patterns, chunks)} + + @property + def reverse_op(self) -> ConversionOps: + return Concatenate(self.dim) + + class Concatenate(ConversionOps): """Concatenate tensors along `dim`. Additionally, if concatenating an aribitrary number of tensors, `num_shards_attribute` is a config field to read to know how many tensors to recreate when using the opposite Ops.""" diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 3111734c9822..c43d61bbd8e1 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -166,7 +166,11 @@ "image_sizes_videos", "pixel_attention_mask", "pixel_values_images", - "num_local_patches", + "image_token_pooling", + "image_grids", + "image_num_crops", + "video_token_pooling", + "video_grids", ) diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py index d1868dc08275..30a822ae08f6 100644 --- a/src/transformers/models/__init__.py +++ b/src/transformers/models/__init__.py @@ -315,6 +315,7 @@ from .modernbert import * from .modernbert_decoder import * from .modernvbert import * + from .molmo2 import * from .moonshine import * from .moonshine_streaming import * from .moshi import * diff --git a/src/transformers/models/auto/auto_mappings.py b/src/transformers/models/auto/auto_mappings.py index e48fff4de60a..a712081b85be 100644 --- a/src/transformers/models/auto/auto_mappings.py +++ b/src/transformers/models/auto/auto_mappings.py @@ -430,6 +430,10 @@ ("modernbert", "ModernBertConfig"), ("modernbert-decoder", "ModernBertDecoderConfig"), ("modernvbert", "ModernVBertConfig"), + ("molmo2", "Molmo2Config"), + ("molmo2_adapter", "Molmo2AdapterConfig"), + ("molmo2_text", "Molmo2TextConfig"), + ("molmo2_vision", "Molmo2VisionConfig"), ("moonshine", "MoonshineConfig"), ("moonshine_streaming", "MoonshineStreamingConfig"), ("moonshine_streaming_encoder", "MoonshineStreamingEncoderConfig"), @@ -922,6 +926,9 @@ ("mllama_vision_model", "mllama"), ("mm-grounding-dino", "mm_grounding_dino"), ("modernbert-decoder", "modernbert_decoder"), + ("molmo2_adapter", "molmo2"), + ("molmo2_text", "molmo2"), + ("molmo2_vision", "molmo2"), ("moonshine_streaming_encoder", "moonshine_streaming"), ("moshi_depth", "moshi"), ("muse_glimmer_text", "muse_glimmer"), @@ -1073,6 +1080,7 @@ ("llava_onevision", {"torchvision": "LlavaOnevisionVideoProcessor"}), ("minicpmv4_6", {"torchvision": "MiniCPMV4_6VideoProcessor"}), ("minimax_m3_vl", {"torchvision": "MiniMaxM3VLVideoProcessor"}), + ("molmo2", {"torchvision": "Molmo2VideoProcessor"}), ("muse_glimmer", {"torchvision": "MuseGlimmerVideoProcessor"}), ("nemotron_h_omni", {"torchvision": "NemotronH_Omni_Reasoning_V3VideoProcessor"}), ("pe_video", {"torchvision": "PeVideoVideoProcessor"}), @@ -1207,6 +1215,7 @@ ("minicpmv4_7", "MiniCPMV4_7Processor"), ("minimax_m3_vl", "MiniMaxM3VLProcessor"), ("mllama", "MllamaProcessor"), + ("molmo2", "Molmo2Processor"), ("moonshine_streaming", "MoonshineStreamingProcessor"), ("muse_glimmer", "MuseGlimmerProcessor"), ("musicflamingo", "MusicFlamingoProcessor"), @@ -1350,6 +1359,7 @@ ("mobilenet_v1", {"pil": "MobileNetV1ImageProcessorPil", "torchvision": "MobileNetV1ImageProcessor"}), ("mobilenet_v2", {"pil": "MobileNetV2ImageProcessorPil", "torchvision": "MobileNetV2ImageProcessor"}), ("mobilevit", {"pil": "MobileViTImageProcessorPil", "torchvision": "MobileViTImageProcessor"}), + ("molmo2", {"torchvision": "Molmo2ImageProcessor"}), ("muse_glimmer", {"torchvision": "MuseGlimmerImageProcessor"}), ("nemotron_h_omni", {"torchvision": "NemotronH_Omni_Reasoning_V3ImageProcessor"}), ("neomme", {"torchvision": "NeoMMEImageProcessor"}), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 9a0eeef5b43a..cf8d7e722ed9 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -360,6 +360,8 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("modernbert", "ModernBertModel"), ("modernbert-decoder", "ModernBertDecoderModel"), ("modernvbert", "ModernVBertModel"), + ("molmo2", "Molmo2Model"), + ("molmo2_text", "Molmo2TextModel"), ("moonshine", "MoonshineModel"), ("moonshine_streaming", "MoonshineStreamingModel"), ("moshi", "MoshiModel"), @@ -1159,6 +1161,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("mistral3", "Mistral3ForConditionalGeneration"), ("mistral4", "Mistral4ForCausalLM"), ("mllama", "MllamaForConditionalGeneration"), + ("molmo2", "Molmo2ForConditionalGeneration"), ("muse_glimmer", "MuseGlimmerForConditionalGeneration"), ("nemotron_h_omni", "NemotronH_Omni_Reasoning_V3"), ("ovis2", "Ovis2ForConditionalGeneration"), @@ -2005,6 +2008,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("longformer", "LongformerModel"), ("mllama", "MllamaTextModel"), ("mobilebert", "MobileBertModel"), + ("molmo2_text", "Molmo2TextModel"), ("mt5", "MT5EncoderModel"), ("nystromformer", "NystromformerModel"), ("reformer", "ReformerModel"), diff --git a/src/transformers/models/molmo2/__init__.py b/src/transformers/models/molmo2/__init__.py new file mode 100644 index 000000000000..fd9267f0d5a2 --- /dev/null +++ b/src/transformers/models/molmo2/__init__.py @@ -0,0 +1,36 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from the modules in `src/transformers/models/molmo2`. +# Do NOT edit this file manually as any edits will be overwritten by auto-generation of the file. +# A module is picked up once it exports a name through `__all__`. +# Regenerate the file with: `python utils/check_inits.py --fix_and_overwrite` +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_molmo2 import * + from .image_processing_molmo2 import * + from .modeling_molmo2 import * + from .processing_molmo2 import * + from .video_processing_molmo2 import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/src/transformers/models/molmo2/configuration_molmo2.py b/src/transformers/models/molmo2/configuration_molmo2.py new file mode 100644 index 000000000000..b9ee2cd05ec1 --- /dev/null +++ b/src/transformers/models/molmo2/configuration_molmo2.py @@ -0,0 +1,240 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/molmo2/modular_molmo2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_molmo2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig, SubConfigSpec +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="allenai/Molmo2-8B") +@strict +class Molmo2VisionConfig(PreTrainedConfig): + r""" + image_size (`list[int]`, *optional*): + Input image size as (height, width), `[378, 378]` when not provided. + num_position_embeddings (`int`, *optional*): + Number of positional embeddings, one per image patch. Derived from `image_size` and `patch_size` when not + provided. + """ + + model_type = "molmo2_vision" + base_config_key = "vision_config" + + hidden_size: int = 1152 + intermediate_size: int = 4304 + num_hidden_layers: int = 25 + num_attention_heads: int = 16 + num_key_value_heads: int = 16 + head_dim: int = 72 + hidden_act: str = "gelu_pytorch_tanh" + layer_norm_eps: float = 1e-6 + image_size: list[int] | None = None + patch_size: int = 14 + num_position_embeddings: int | None = None + num_channels: int = 3 + attention_dropout: float = 0.0 + attention_bias: bool = True + initializer_range: float = 0.02 + + def __post_init__(self, **kwargs): + # Keys of the released checkpoints' `config.json`. + kwargs.pop("model_type", None) + self.image_size = kwargs.pop("image_default_input_size", self.image_size) + self.patch_size = kwargs.pop("image_patch_size", self.patch_size) + self.num_position_embeddings = kwargs.pop("image_num_pos", self.num_position_embeddings) + if self.image_size is None: + self.image_size = [378, 378] + if self.num_position_embeddings is None: + self.num_position_embeddings = (self.image_size[0] // self.patch_size) * ( + self.image_size[1] // self.patch_size + ) + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="allenai/Molmo2-8B") +@strict +class Molmo2AdapterConfig(PreTrainedConfig): + r""" + vision_feature_layer (`list[int]`, *optional*): + Indices of the ViT layers whose outputs are concatenated and pooled, `[24, 18]` when not provided. Negative + indices count from the last layer. + text_hidden_size (`int`, *optional*, defaults to 3584): + Hidden size of the text model (used for projection). + """ + + model_type = "molmo2_adapter" + base_config_key = "adapter_config" + + vision_feature_layer: list[int] | None = None + hidden_size: int = 1152 + num_attention_heads: int = 16 + num_key_value_heads: int = 16 + head_dim: int = 72 + attention_dropout: float = 0.0 + attention_bias: bool = True + hidden_act: str = "silu" + intermediate_size: int = 18944 + text_hidden_size: int = 3584 + mlp_bias: bool = False + initializer_range: float = 0.02 + + def __post_init__(self, **kwargs): + # Keys of the released checkpoints' `config.json`. + kwargs.pop("model_type", None) + self.vision_feature_layer = kwargs.pop("vit_layers", self.vision_feature_layer) + if self.vision_feature_layer is None: + self.vision_feature_layer = [24, 18] + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="allenai/Molmo2-8B") +@strict +class Molmo2TextConfig(PreTrainedConfig): + r""" + additional_vocab_size (`int`, *optional*, defaults to 128): + Number of additional vocabulary tokens beyond the base vocabulary. + qk_norm_type (`str`, *optional*, defaults to `"qwen3"`): + Query/key normalization layout used by the checkpoint. `"qwen3"` normalizes per head; `"olmo"` normalizes the + full projected query/key tensors. + rope_scaling_layers (`list[int]`, *optional*): + Indices of the layers that apply the scaled RoPE described by `rope_parameters`. The remaining layers use an + unscaled RoPE with the same theta. All layers are scaled when not provided. + norm_after (`bool`, *optional*, defaults to `False`): + Whether to apply layer normalization after the attention/FFN blocks instead of before. + """ + + model_type = "molmo2_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"qkv_bias": "attention_bias", "layer_norm_eps": "rms_norm_eps"} + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + hidden_size: int = 4096 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + head_dim: int = 128 + vocab_size: int = 151936 + additional_vocab_size: int = 128 + attention_bias: bool = False + qk_norm_type: str = "qwen3" + num_hidden_layers: int = 36 + intermediate_size: int = 12288 + hidden_act: str = "silu" + attention_dropout: float = 0.0 + max_position_embeddings: int = 36864 + rope_parameters: RopeParameters | dict | None = None + rope_scaling_layers: list[int] | None = None + rms_norm_eps: float = 1e-6 + norm_after: bool = False + initializer_range: float = 0.02 + use_cache: bool = True + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + if self.rope_scaling_layers is None: + self.rope_scaling_layers = list(range(self.num_hidden_layers)) + super().__post_init__(**kwargs) + + def validate_architecture(self): + super().validate_architecture() + if self.qk_norm_type not in ("qwen3", "olmo"): + raise ValueError(f"Unsupported `qk_norm_type`: {self.qk_norm_type}") + if self.additional_vocab_size < 0: + raise ValueError(f"`additional_vocab_size` must be >= 0, got {self.additional_vocab_size}") + + +@auto_docstring(checkpoint="allenai/Molmo2-8B") +@strict +class Molmo2Config(PreTrainedConfig): + r""" + vision_config (`Molmo2VisionConfig`, *optional*): + Configuration for the vision transformer backbone. + adapter_config (`Molmo2AdapterConfig`, *optional*): + Configuration for the vision-to-language adapter. + image_start_token_id (`int`, *optional*, defaults to 151936): + Token ID marking the start of an image region. + low_res_image_start_token_id (`int`, *optional*, defaults to 151940): + Token ID marking the start of a low-resolution image crop. + image_end_token_id (`int`, *optional*, defaults to 151937): + Token ID marking the end of an image region. + image_patch_id (`int`, *optional*, defaults to 151938): + Token ID for image patches. + frame_start_token_id (`int`, *optional*, defaults to 151943): + Token ID marking the start of a video frame. + frame_end_token_id (`int`, *optional*, defaults to 151944): + Token ID marking the end of a video frame. + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether the model's input and output word embeddings should be tied. + """ + + model_type = "molmo2" + attribute_map = {"image_token_id": "image_patch_id", "video_token_id": "image_patch_id"} + sub_configs_defaults = { + "text_config": SubConfigSpec(config_class=Molmo2TextConfig), + "vision_config": SubConfigSpec(config_class=Molmo2VisionConfig), + "adapter_config": SubConfigSpec(config_class=Molmo2AdapterConfig), + } + + vision_config: dict | PreTrainedConfig | None = None + adapter_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + image_start_token_id: int = 151936 + low_res_image_start_token_id: int = 151940 + image_end_token_id: int = 151937 + image_patch_id: int = 151938 + frame_start_token_id: int = 151943 + frame_end_token_id: int = 151944 + initializer_range: float = 0.02 + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + # Checkpoints serialized before the rename store the vision sub-config under `vit_config`. + legacy_vision_config = kwargs.pop("vit_config", None) + if self.vision_config is None and legacy_vision_config is not None: + self.vision_config = legacy_vision_config + + super().__post_init__(**kwargs) + + # The hub `config.json` declares a 27-layer ViT, but the released weights stop at layer 24, the deepest one + # `vit_layers` reads. Without this trim, layers 25-26 are built with random weights. Do not remove. + if legacy_vision_config is not None: + num_vit_layers = self.vision_config.num_hidden_layers + self.adapter_config.vision_feature_layer = [ + layer % num_vit_layers for layer in self.adapter_config.vision_feature_layer + ] + self.vision_config.num_hidden_layers = max(self.adapter_config.vision_feature_layer) + 1 + + def validate_architecture(self): + super().validate_architecture() + num_vit_layers = self.vision_config.num_hidden_layers + if not all(-num_vit_layers <= layer < num_vit_layers for layer in self.adapter_config.vision_feature_layer): + raise ValueError( + f"`adapter_config.vision_feature_layer` {self.adapter_config.vision_feature_layer} must index the " + f"{num_vit_layers} layers of `vision_config`." + ) + + +__all__ = ["Molmo2AdapterConfig", "Molmo2Config", "Molmo2TextConfig", "Molmo2VisionConfig"] diff --git a/src/transformers/models/molmo2/image_processing_molmo2.py b/src/transformers/models/molmo2/image_processing_molmo2.py new file mode 100644 index 000000000000..bbdd0699c9db --- /dev/null +++ b/src/transformers/models/molmo2/image_processing_molmo2.py @@ -0,0 +1,453 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/molmo2/modular_molmo2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_molmo2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import itertools +import math + +import torch +from torch.nn import functional as F + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, ImageInput, PILImageResampling, SizeDict +from ...processing_utils import Unpack +from ...utils import TensorType, auto_docstring +from .processing_molmo2 import Molmo2ImagesKwargs + + +def select_tiling(height: int, width: int, patch_size: int, max_num_crops: int) -> tuple[int, int]: + """Same as `get_optimal_tiled_canvas` in cohere2_vision, with the candidates in (height, width) order and ties + broken on height, as in the original Molmo2 processor.""" + tilings = [ + (tile_height, tile_width) + for tile_height, tile_width in itertools.product(range(1, max_num_crops + 1), repeat=2) + if tile_height * tile_width <= max_num_crops + ] + tilings.sort(key=lambda x: (x[0] * x[1], x[0])) + + candidate_resolutions = torch.tensor(tilings, dtype=torch.int32) * patch_size + original_size = torch.tensor([height, width], dtype=torch.float32) + + required_scales = candidate_resolutions.to(torch.float32) / original_size + required_scale = required_scales.amin(dim=-1, keepdim=True) + + if torch.all(required_scale < 1): + return tilings[int(required_scale.argmax())] + + required_scale = torch.where(required_scale < 1.0, 10e9, required_scale) + return tilings[int(required_scale.argmin())] + + +def compute_crop_geometry( + height: int, + width: int, + max_crops: int, + overlap_margins: list[int], + crop_size: int, + patch_size: int, +) -> tuple[int, int, int, int]: + """Tiling `(tiling_h, tiling_w)` and canvas pixel size `(canvas_height, canvas_width)` for one image.""" + left_margin, right_margin = overlap_margins + crop_patches = crop_size // patch_size + window_patches = crop_patches - (left_margin + right_margin) + window_size = window_patches * patch_size + margin_size = (left_margin + right_margin) * patch_size + tiling_h, tiling_w = select_tiling(height - margin_size, width - margin_size, window_size, max_crops) + canvas_height = tiling_h * window_size + margin_size + canvas_width = tiling_w * window_size + margin_size + return tiling_h, tiling_w, canvas_height, canvas_width + + +def batch_pixels_to_patches(images: torch.Tensor, patch_size: int) -> torch.Tensor: + """Reshape images of [num_images, h, w, channels] -> [num_images, n_patches, pixels_per_patch]""" + num_crops, height, width, channels = images.shape + num_patches_height = height // patch_size + num_patches_width = width // patch_size + images = images.reshape(num_crops, num_patches_height, patch_size, num_patches_width, patch_size, channels) + images = images.transpose(2, 3) + images = images.reshape(num_crops, num_patches_height * num_patches_width, patch_size * patch_size * channels) + return images + + +def arange_for_pooling( + index_grid: torch.Tensor, + pool_h: int, + pool_w: int, +) -> tuple[torch.Tensor, int, int]: + """Group a [h, w] patch-index grid into pooling windows: `[num_rows * num_cols, pool_h * pool_w]`, padded + with -1. Also returns the pooled grid's `(num_rows, num_cols)`.""" + height_padding = pool_h * ((index_grid.shape[0] + pool_h - 1) // pool_h) - index_grid.shape[0] + width_padding = pool_w * ((index_grid.shape[1] + pool_w - 1) // pool_w) - index_grid.shape[1] + index_grid = F.pad( + index_grid, + (width_padding // 2, (width_padding + 1) // 2, height_padding // 2, (height_padding + 1) // 2), + mode="constant", + value=-1, + ) + num_rows, num_cols = index_grid.shape[0] // pool_h, index_grid.shape[1] // pool_w + pooling_indices = ( + index_grid.reshape(num_rows, pool_h, num_cols, pool_w) + .permute(0, 2, 1, 3) + .reshape(num_rows * num_cols, pool_h * pool_w) + ) + return pooling_indices, num_rows, num_cols + + +def resize_and_normalize_image( + backend: "TorchvisionBackend", + image: torch.Tensor, + output_size: list[int], + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], +) -> torch.Tensor: + # Under `torch.compile`, torchvision's uint8 bilinear resize can produce values slightly outside 0-255, + # which overflow and wrap around (-1 becomes 255, so a near-black pixel becomes white). Resizing in + # float, then rounding and clamping back, avoids that and gives the exact same bytes in eager mode. + is_integer_input = not image.dtype.is_floating_point + resized = backend.resize( + image.float() if is_integer_input else image, + size=SizeDict(height=output_size[0], width=output_size[1]), + resample=resample, + antialias=False, + ) + if is_integer_input: + resized = resized.round().clamp(0, 255).to(image.dtype) + return backend.rescale_and_normalize(resized, do_rescale, rescale_factor, do_normalize, image_mean, image_std) + + +def build_resized_image( + backend: "TorchvisionBackend", + images: torch.Tensor, + crop_size: int, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + image_patch_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # `images`: a batch of same-shape images `[N, C, H, W]`; resize the whole batch at once. + resized = resize_and_normalize_image( + backend, + images, + [crop_size, crop_size], + resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + ) + # [N, C, S, S] -> [N, 1, S, S, C]: one global (low-res) view per image. + resized = resized.permute(0, 2, 3, 1).unsqueeze(1) + # The per-patch index grid depends only on the (shared) shape, so it is built once. + patches_per_side = crop_size // image_patch_size + resized_index_grid = torch.arange(patches_per_side**2, dtype=torch.int32, device=images.device).reshape( + patches_per_side, patches_per_side + ) + return resized, resized_index_grid + + +@auto_docstring +class Molmo2ImageProcessor(TorchvisionBackend): + valid_kwargs = Molmo2ImagesKwargs + model_input_names = ["pixel_values", "image_token_pooling", "image_grids", "image_num_crops"] + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 378, "width": 378} + do_resize = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + max_crops = 8 + overlap_margins = [4, 4] + patch_size = 14 + pooling_size = [2, 2] + + def __init__(self, **kwargs: Unpack[Molmo2ImagesKwargs]): + super().__init__(**kwargs) + + def _build_overlapping_crops( + self, + images: torch.Tensor, + max_crops: int, + overlap_margins: list[int], + crop_size: int, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + image_patch_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Tile a batch of same-shape images `[N, C, H, W]` into overlapping square crops. The tiling and the per-patch index grid depend only on the (shared) shape, so they are computed once and the resize/unfold are batched over N.""" + left_margin, right_margin = overlap_margins + crop_patches = crop_size // image_patch_size + window_patches = crop_patches - (left_margin + right_margin) + window_size = window_patches * image_patch_size + + num_images, num_channels, original_height, original_width = images.shape + tiling_h, tiling_w, canvas_height, canvas_width = compute_crop_geometry( + original_height, + original_width, + max_crops=max_crops, + overlap_margins=overlap_margins, + crop_size=crop_size, + patch_size=image_patch_size, + ) + + canvas = resize_and_normalize_image( + self, + images, + [canvas_height, canvas_width], + resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + ) + + # [N, C, canvas_height, canvas_width] -> unfold spatial dims -> [N, C, tiling_h, tiling_w, crop, crop] + crops = canvas.unfold(2, crop_size, window_size).unfold(3, crop_size, window_size) + crops = ( + crops.permute(0, 2, 3, 4, 5, 1) + .reshape(num_images, tiling_h * tiling_w, crop_size, crop_size, num_channels) + .contiguous() + ) + + patch_index_grid = torch.arange( + tiling_h * tiling_w * crop_patches * crop_patches, dtype=torch.int32, device=images.device + ).reshape(tiling_h, tiling_w, crop_patches, crop_patches) + if left_margin: + patch_index_grid[1:, :, :left_margin, :] = -1 + patch_index_grid[:, 1:, :, :left_margin] = -1 + if right_margin: + patch_index_grid[:-1, :, -right_margin:, :] = -1 + patch_index_grid[:, :-1, :, -right_margin:] = -1 + + patch_index_grid = patch_index_grid.permute(0, 2, 1, 3).reshape(-1) + patch_index_grid = patch_index_grid[patch_index_grid >= 0].reshape( + canvas_height // image_patch_size, canvas_width // image_patch_size + ) + + return crops, patch_index_grid + + def _image_batch_to_patches_and_grids( + self, + images: torch.Tensor, + max_crops: int, + overlap_margins: list[int], + crop_size: int, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + image_patch_size: int, + image_pooling_w: int, + image_pooling_h: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Process a batch of same-shape images `[N, C, H, W]`. The grid and pooling indices are + shape-dependent (shared across the batch) and returned expanded to `[N, ...]` so the caller + can reorder them per-image alongside the batched patch tensor.""" + patches_per_crop = (crop_size // image_patch_size) ** 2 + + crops, patch_index_grid = self._build_overlapping_crops( + images, + max_crops=max_crops, + overlap_margins=overlap_margins, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=image_patch_size, + ) + pooling_indices, num_patch_rows, num_patch_cols = arange_for_pooling( + patch_index_grid, image_pooling_h, image_pooling_w + ) + + resized, resized_index_grid = build_resized_image( + self, + images, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=image_patch_size, + ) + # [N, 1, S, S, C] + [N, ncrops, S, S, C] -> [N, 1 + ncrops, S, S, C] + crops = torch.cat([resized, crops], dim=1) + + resized_pooling_indices, resized_h, resized_w = arange_for_pooling( + resized_index_grid, image_pooling_h, image_pooling_w + ) + + pooling_indices = torch.where(pooling_indices >= 0, pooling_indices + patches_per_crop, -1) + pooling_indices = torch.cat([resized_pooling_indices, pooling_indices]) + image_grid = torch.tensor( + [[resized_h, resized_w, num_patch_rows, num_patch_cols]], dtype=torch.int64, device=images.device + ) + + # [N, total_crops, S, S, C] -> patches [N, total_crops, n_patch, pixels_per_patch] + num_images, total_crops = crops.shape[0], crops.shape[1] + pixels_per_patch = image_patch_size * image_patch_size * crops.shape[-1] + patches = batch_pixels_to_patches( + crops.reshape(num_images * total_crops, *crops.shape[2:]), image_patch_size + ).reshape(num_images, total_crops, -1, pixels_per_patch) + + # The grid and pooling indices are shared by all images of this shape; expand them to one row + # per image so the batch can later be restored to the original image order. + return image_grid.expand(num_images, -1), patches, pooling_indices.unsqueeze(0).expand(num_images, -1, -1) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + **kwargs: Unpack[Molmo2ImagesKwargs], + ) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + do_convert_rgb: bool, + max_crops: int, + overlap_margins: list[int], + patch_size: int, + pooling_size: list[int], + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + if size.height != size.width: + raise ValueError(f"Molmo2 only supports a square `size`, got height={size.height}, width={size.width}.") + crop_size = size.height + image_pooling_h, image_pooling_w = pooling_size + + # Group images by shape and process each unique shape as a single batch (the tiling and all + # patch/pooling indices are shape-dependent only), then restore the original order. + device = images[0].device + grouped_images, grouped_index = group_images_by_shape(images, disable_grouping=disable_grouping) + + grids_grouped: dict = {} + patches_grouped: dict = {} + pooled_grouped: dict = {} + for shape, stacked_images in grouped_images.items(): + image_grid, patches, pooled_indices = self._image_batch_to_patches_and_grids( + stacked_images, + max_crops=max_crops, + overlap_margins=overlap_margins, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=patch_size, + image_pooling_w=image_pooling_w, + image_pooling_h=image_pooling_h, + ) + grids_grouped[shape] = image_grid + patches_grouped[shape] = patches + pooled_grouped[shape] = pooled_indices + + grids = reorder_images(grids_grouped, grouped_index) + patches = reorder_images(patches_grouped, grouped_index) + pooled = reorder_images(pooled_grouped, grouped_index) + + all_crops: list[torch.Tensor] = [] + all_pooled: list[torch.Tensor] = [] + patch_offset = 0 + for crops, pooled_indices in zip(patches, pooled): + all_pooled.append(torch.where(pooled_indices >= 0, pooled_indices + patch_offset, pooled_indices)) + all_crops.append(crops) + patch_offset += crops.shape[0] * crops.shape[1] + + data = { + "pixel_values": torch.cat(all_crops, dim=0), + "image_token_pooling": torch.cat(all_pooled, dim=0), + "image_grids": torch.stack(grids, dim=0), + "image_num_crops": torch.tensor([crops.shape[0] for crops in patches], dtype=torch.int64, device=device), + } + return BatchFeature(data=data, tensor_type=return_tensors) + + def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None) -> int: + image_grid, _ = self.get_image_grid_and_crops(height, width, images_kwargs) + resized_h, resized_w, high_res_rows, high_res_cols = image_grid + return resized_h * resized_w + high_res_rows * high_res_cols + + def get_image_grid_and_crops(self, height: int, width: int, images_kwargs=None) -> tuple[list[int], int]: + """Return the `[low_res_h, low_res_w, high_res_h, high_res_w]` pooled grid and the crop count for one image.""" + if images_kwargs is None: + images_kwargs = {} + max_crops = images_kwargs.get("max_crops", self.max_crops) + overlap_margins = images_kwargs.get("overlap_margins", self.overlap_margins) + patch_size = images_kwargs.get("patch_size", self.patch_size) + pooling_size = images_kwargs.get("pooling_size", self.pooling_size) + size = images_kwargs.get("size", self.size) + + base_h = size["height"] if isinstance(size, dict) else size.height + base_w = size["width"] if isinstance(size, dict) else size.width + pooling_h, pooling_w = pooling_size + + tiling_h, tiling_w, canvas_height, canvas_width = compute_crop_geometry( + height, + width, + max_crops=max_crops, + overlap_margins=overlap_margins, + crop_size=base_h, + patch_size=patch_size, + ) + num_patch_rows_high = math.ceil((canvas_height // patch_size) / pooling_h) + num_patch_cols_high = math.ceil((canvas_width // patch_size) / pooling_w) + + resized_h = math.ceil((base_h // patch_size) / pooling_h) + resized_w = math.ceil((base_w // patch_size) / pooling_w) + + return [resized_h, resized_w, num_patch_rows_high, num_patch_cols_high], tiling_h * tiling_w + 1 + + +__all__ = ["Molmo2ImageProcessor"] diff --git a/src/transformers/models/molmo2/modeling_molmo2.py b/src/transformers/models/molmo2/modeling_molmo2.py new file mode 100644 index 000000000000..c0e2a783e5e9 --- /dev/null +++ b/src/transformers/models/molmo2/modeling_molmo2.py @@ -0,0 +1,1068 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/molmo2/modular_molmo2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_molmo2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...backbone_utils import filter_output_hidden_states +from ...cache_utils import Cache, DynamicCache +from ...configuration_utils import PreTrainedConfig +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...masking_utils import create_bidirectional_mask, create_causal_mask, create_masks_for_generate +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_molmo2 import Molmo2AdapterConfig, Molmo2Config, Molmo2TextConfig, Molmo2VisionConfig + + +@auto_docstring( + custom_intro=""" + Base class for Molmo2 causal language model (or autoregressive) outputs. + """ +) +@dataclass +class Molmo2CausalLMOutputWithPast(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + image_hidden_states: torch.FloatTensor | None = None + + +@auto_docstring( + custom_intro=""" + Base class for Molmo2 outputs, with hidden states and attentions. + """ +) +@dataclass +class Molmo2ModelOutputWithPast(BaseModelOutputWithPast): + r""" + past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + image_hidden_states (`torch.FloatTensor`, *optional*): + A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. + image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. + """ + + image_hidden_states: torch.FloatTensor | None = None + + +class Molmo2RotaryEmbedding(nn.Module): + """RoPE from the flat `rope_parameters`; `scaled=False` keeps the theta but skips the scaling, for the + layers outside `rope_scaling_layers`.""" + + def __init__(self, config: Molmo2TextConfig, scaled: bool = True): + super().__init__() + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + + # `rope_type = "default"` also keeps the generic `_init_weights` rotary re-init on the unscaled table. + self.rope_type = self.config.rope_parameters["rope_type"] if scaled else "default" + rope_init_fn: Callable = ( + self.compute_default_rope_parameters + if self.rope_type == "default" + else ROPE_INIT_FUNCTIONS[self.rope_type] + ) + inv_freq, self.attention_scaling = rope_init_fn(self.config) + + self.inv_freq = nn.Buffer(inv_freq, persistent=False) + self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False) + + @staticmethod + def compute_default_rope_parameters(config: Molmo2Config, **kwargs) -> tuple[torch.Tensor, float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + # Compute the inverse frequencies + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + freqs = position_ids[..., None].float() * self.inv_freq.to(device=x.device, dtype=torch.float) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +@use_kernel_forward_from_hub("RMSNorm") +class Molmo2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + Molmo2RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +@use_kernel_forward_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernelized_func(apply_rotary_pos_emb) +class Molmo2Attention(nn.Module): + """Olmo2 attention whose q/k RMSNorm runs either over the full projection (`olmo`) or per head (`qwen3`).""" + + def __init__(self, config: Molmo2TextConfig, layer_idx: int) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + self.q_norm = Molmo2RMSNorm( + self.head_dim if config.qk_norm_type == "qwen3" else config.num_attention_heads * self.head_dim, + eps=config.rms_norm_eps, + ) + self.k_norm = Molmo2RMSNorm( + self.head_dim if config.qk_norm_type == "qwen3" else config.num_key_value_heads * self.head_dim, + eps=config.rms_norm_eps, + ) + self.qk_norm_type = config.qk_norm_type + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + if self.qk_norm_type == "olmo": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.view(hidden_shape) + key_states = key_states.view(hidden_shape) + + if self.qk_norm_type == "qwen3": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Molmo2MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class Molmo2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Molmo2TextConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Molmo2Attention(config=config, layer_idx=layer_idx) + + self.mlp = Molmo2MLP(config) + self.input_layernorm = Molmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Molmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm_after = config.norm_after + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + # `norm_after=True` normalizes each sublayer's output (post-norm) instead of its input (pre-norm) + residual = hidden_states + if not self.norm_after: + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + if self.norm_after: + hidden_states = self.input_layernorm(hidden_states) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + if not self.norm_after: + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + if self.norm_after: + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class Molmo2PreTrainedModel(PreTrainedModel): + config: Molmo2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Molmo2DecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + + _can_compile_fullgraph = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": Molmo2DecoderLayer, + "attentions": Molmo2Attention, + } + input_modalities = ("image", "video", "text") + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, Molmo2VisionModel): + init.normal_(module.positional_embedding, mean=0.0, std=self.config.initializer_range) + + +class Molmo2VisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +# Modular automatically inherits RoPE, hence no inheritance for now +class Molmo2VisionAttention(nn.Module): + def __init__(self, config: Molmo2VisionConfig | Molmo2AdapterConfig, hidden_size: int | None = None): + super().__init__() + self.config = config + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = False + self.q_proj = nn.Linear( + hidden_size or config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + hidden_size or config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + hidden_size or config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + encoder_hidden_states = hidden_states if encoder_hidden_states is None else encoder_hidden_states + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + key_value_shape = (*encoder_hidden_states.shape[:-1], -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(encoder_hidden_states).view(key_value_shape).transpose(1, 2) + value_states = self.v_proj(encoder_hidden_states).view(key_value_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + is_causal=self.is_causal, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +class Molmo2VisionEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Molmo2VisionConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.self_attn = Molmo2VisionAttention(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = Molmo2VisionMLP(config) + + @auto_docstring + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.FloatTensor: + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +@auto_docstring +class Molmo2VisionModel(Molmo2PreTrainedModel): + config_class = Molmo2VisionConfig + main_input_name = "pixel_values" + input_modalities = ("image",) + _no_split_modules = ["Molmo2VisionEncoderLayer"] + _can_record_outputs = { + "hidden_states": Molmo2VisionEncoderLayer, + "attentions": Molmo2VisionAttention, + } + + def __init__(self, config: Molmo2VisionConfig): + super().__init__(config) + self.positional_embedding = nn.Parameter(torch.zeros(config.num_position_embeddings, config.hidden_size)) + self.patch_embedding = nn.Linear( + config.patch_size * config.patch_size * config.num_channels, config.hidden_size + ) + # trf-ignore: TRF034 (false positive: Siglip2EncoderLayer subclasses GradientCheckpointingLayer) + self.layers = nn.ModuleList([Molmo2VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.post_init() + + @merge_with_config_defaults + @capture_outputs(tie_last_hidden_states=False) + @auto_docstring + def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutputWithPooling: + hidden_states = self.patch_embedding(pixel_values.to(dtype=self.dtype)) + # patch count == num_position_embeddings, locked by config; only retraining with a different grid breaks this. + hidden_states = hidden_states + self.positional_embedding[None, :, :].to(hidden_states.dtype) + + for encoder_layer in self.layers: + hidden_states = encoder_layer(hidden_states, None, **kwargs) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=hidden_states.mean(dim=1), + ) + + +class Molmo2ImageProjectorMLP(nn.Module): + def __init__(self, config: Molmo2AdapterConfig): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(config.intermediate_size, config.text_hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +@auto_docstring( + custom_intro=""" + The Molmo2 vision adapter: pools ViT patch features with a cross-attention layer and projects them into the + language model's embedding space. + """ +) +class Molmo2Adapter(Molmo2PreTrainedModel): + config_class = Molmo2AdapterConfig + input_modalities = ("image",) + _no_split_modules = ["Molmo2VisionAttention"] + _can_record_outputs = {"attentions": Molmo2VisionAttention} + + def __init__(self, config: Molmo2AdapterConfig): + super().__init__(config) + pooling_input_dim = config.hidden_size * len(config.vision_feature_layer) + self.image_pooling_2d = Molmo2VisionAttention(config, hidden_size=pooling_input_dim) + self.image_projector = Molmo2ImageProjectorMLP(config) + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + image_features: torch.Tensor, + pooled_patches_idx: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + r""" + image_features (`torch.Tensor` of shape `(num_crops, num_patches, hidden_size * len(vision_feature_layer))`): + Concatenated intermediate ViT features of every crop. + pooled_patches_idx (`torch.Tensor` of shape `(num_tokens, pool_h * pool_w)`): + Indices into the flattened patch sequence pooled by each output token; `-1` marks padding slots. + """ + flat_features = image_features.reshape(-1, image_features.shape[-1]) + + valid_mask = pooled_patches_idx >= 0 + valid_token_mask = torch.any(valid_mask, -1) + + patches_to_pool = flat_features[torch.clip(pooled_patches_idx, 0)] + patches_to_pool = patches_to_pool * valid_mask.to(patches_to_pool.dtype)[..., None] + + num_valid_patches = valid_mask.float().sum(-1) + num_valid_patches = torch.where(num_valid_patches == 0, 1, num_valid_patches) + query = patches_to_pool.sum(-2, keepdim=True) / num_valid_patches[:, None, None].to(patches_to_pool.dtype) + + attention_mask = create_bidirectional_mask( + config=self.config, inputs_embeds=query, attention_mask=valid_mask, encoder_hidden_states=patches_to_pool + ) + pooled_features, _ = self.image_pooling_2d(query, patches_to_pool, attention_mask=attention_mask, **kwargs) + pooled_features = pooled_features.squeeze(1) + pooled_features = self.image_projector(pooled_features) + return BaseModelOutput(last_hidden_state=pooled_features[valid_token_mask]) + + +@auto_docstring +class Molmo2TextModel(Molmo2PreTrainedModel): + config: Molmo2TextConfig + + def __init__(self, config: Molmo2TextConfig): + super().__init__(config) + self.vocab_size = config.vocab_size + # The checkpoint's extra-vocabulary table is concatenated onto the base one at load time + # (see `conversion_mapping.py`), so the embedding covers `vocab_size + additional_vocab_size`. + self.embed_tokens = nn.Embedding(config.vocab_size + config.additional_vocab_size, config.hidden_size) + # trf-ignore: TRF034 (false positive: LlamaDecoderLayer subclasses GradientCheckpointingLayer) + self.layers = nn.ModuleList( + [Molmo2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Molmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Molmo2RotaryEmbedding(config) + self.rotary_emb_unscaled = Molmo2RotaryEmbedding(config, scaled=False) + self.rope_types = [ + "scaled" if layer_idx in config.rope_scaling_layers else "unscaled" + for layer_idx in range(config.num_hidden_layers) + ] + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = { + "scaled": self.rotary_emb(hidden_states, position_ids=position_ids), + "unscaled": self.rotary_emb_unscaled(hidden_states, position_ids=position_ids), + } + + for layer_idx, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings[self.rope_types[layer_idx]], + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring( + custom_intro=""" + The Molmo2 model which consists of a vision backbone, a pooling adapter and a language model, without a language + modeling head. + """ +) +class Molmo2Model(Molmo2PreTrainedModel): + config: Molmo2Config + + def __init__(self, config: Molmo2Config): + super().__init__(config) + self.vision_tower = Molmo2VisionModel(config.vision_config) + self.multi_modal_projector = Molmo2Adapter(config.adapter_config) + self.language_model = Molmo2TextModel(config.text_config) + self.post_init() + + @can_return_tuple + @filter_output_hidden_states + @auto_docstring( + custom_intro="Obtains pooled image features from the vision tower and the adapter: the `pooler_output` is a " + "tuple with one `(num_image_tokens, hidden_size)` tensor per image." + ) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_token_pooling: torch.Tensor, + image_grids: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + image_shape = None + if pixel_values.dim() == 4: + image_shape = pixel_values.shape[:3] + pixel_values = pixel_values.reshape(-1, pixel_values.shape[-2], pixel_values.shape[-1]) + + image_outputs: BaseModelOutputWithPooling = self.vision_tower(pixel_values, **kwargs) + image_features = torch.cat( + [ + image_outputs.hidden_states[layer + 1 if layer >= 0 else layer] + for layer in self.config.adapter_config.vision_feature_layer + ], + dim=-1, + ) + + if image_shape is not None: + image_features = image_features.reshape(*image_shape, -1) + pooled_features = self.multi_modal_projector(image_features, image_token_pooling).last_hidden_state + split_sizes = (image_grids[:, 0] * image_grids[:, 1] + image_grids[:, 2] * image_grids[:, 3]).tolist() + image_outputs.pooler_output = torch.split(pooled_features, split_sizes) + return image_outputs + + def get_placeholder_mask( + self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor + ): + """ + Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is + equal to the length of multimodal features. If the lengths are different, an error is raised. + """ + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_image_mask = special_image_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + + n_image_tokens = special_image_mask.sum() + n_image_features = image_features.shape[0] * image_features.shape[1] + special_image_mask = special_image_mask.unsqueeze(-1).to(inputs_embeds.device) + torch_compilable_check( + n_image_tokens * inputs_embeds.shape[-1] == image_features.numel(), + f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}", + ) + return special_image_mask + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + image_token_pooling: torch.Tensor | None = None, + image_grids: torch.Tensor | None = None, + image_num_crops: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_token_pooling: torch.Tensor | None = None, + video_grids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + past_key_values: Cache | None = None, + mm_token_type_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Molmo2ModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if pixel_values is not None and pixel_values_videos is not None: + raise ValueError("pixel_values and pixel_values_videos are provided at the same time") + if (pixel_values is not None or pixel_values_videos is not None) and mm_encoder_outputs is not None: + raise ValueError( + "You cannot specify both pixel_values/pixel_values_videos and mm_encoder_outputs at the same time" + ) + if self.training and mm_token_type_ids is None: + raise ValueError("`mm_token_type_ids` is required as a model input when training") + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + mm_encoder_outputs = mm_encoder_outputs if mm_encoder_outputs is not None else {} + if mm_encoder_outputs.get("image") is None and pixel_values is not None: + mm_encoder_outputs["image"] = self.get_image_features(pixel_values, image_token_pooling, image_grids) + if mm_encoder_outputs.get("video") is None and pixel_values_videos is not None: + mm_encoder_outputs["video"] = self.get_video_features( + pixel_values_videos, video_token_pooling, video_grids + ) + + image_features: torch.FloatTensor | None = None + if mm_encoder_outputs.get("image") is not None: + image_features = mm_encoder_outputs["image"].pooler_output + elif mm_encoder_outputs.get("video") is not None: + image_features = mm_encoder_outputs["video"].pooler_output + + if image_features is not None: + image_features = torch.cat(image_features, dim=0) + # `get_placeholder_mask` returns a [batch, seq, 1] mask; Molmo2 *adds* the image features onto + # the placeholder-token embeddings (residual), which means we index `inputs_embeds` directly. + # Boolean indexing does not broadcast, so the mask must be expanded to the hidden dim first. + special_image_mask = self.get_placeholder_mask(input_ids, inputs_embeds, image_features) + special_image_mask = special_image_mask.expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter( + special_image_mask, + inputs_embeds[special_image_mask] + image_features.reshape(-1), + ) + + # An already prepared 4D mask (e.g. from `generate`) is returned as is by `create_causal_mask` + mask_kwargs = { + "config": self.config.get_text_config(), + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + is_prefill = past_key_values is None or not past_key_values.is_initialized or image_features is not None + if mm_token_type_ids is not None and is_prefill: + # every image or frame token attends to every other one, so they all share a single block + mask_kwargs["block_sequence_ids"] = torch.where(mm_token_type_ids.to(inputs_embeds.device) == 1, 0, -1) + causal_mask = create_causal_mask(**mask_kwargs) + + outputs = self.language_model( + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return Molmo2ModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features, + ) + + @can_return_tuple + @auto_docstring( + custom_intro="Obtains pooled video features from the vision tower and the adapter: the `pooler_output` is a " + "tuple with one `(num_video_tokens, hidden_size)` tensor per video." + ) + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_token_pooling: torch.Tensor, + video_grids: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + # Video frames only go through the low-res path: `num_frames` frames of `rows x cols` tokens, no high-res crops. + num_frames, rows, cols = video_grids.unbind(-1) + image_grids = torch.stack([num_frames * rows, cols, torch.zeros_like(cols), torch.zeros_like(cols)], dim=-1) + return self.get_image_features(pixel_values_videos, video_token_pooling, image_grids, **kwargs) + + +@auto_docstring( + custom_intro=""" + The Molmo2 model which consists of a vision backbone, a pooling adapter and a language model. + """ +) +class Molmo2ForConditionalGeneration(Molmo2PreTrainedModel, GenerationMixin): + # the loss must not be normalized by `num_items_in_batch`: logits/labels are filtered before computing it + accepts_loss_kwargs = False + config: Molmo2Config + + def __init__(self, config: Molmo2Config): + super().__init__(config) + + self.model = Molmo2Model(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.vocab_size = config.text_config.vocab_size + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_token_pooling: torch.Tensor, + image_grids: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values, image_token_pooling, image_grids, **kwargs) + + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_token_pooling: torch.Tensor, + video_grids: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_video_features(pixel_values_videos, video_token_pooling, video_grids, **kwargs) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.Tensor | None = None, + image_token_pooling: torch.Tensor | None = None, + image_grids: torch.Tensor | None = None, + image_num_crops: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_token_pooling: torch.Tensor | None = None, + video_grids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + mm_token_type_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Molmo2CausalLMOutputWithPast: + r""" + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, Molmo2ForConditionalGeneration + + >>> model = Molmo2ForConditionalGeneration.from_pretrained("allenai/Molmo2-8B") + >>> processor = AutoProcessor.from_pretrained("allenai/Molmo2-8B") + + >>> prompt = "What's the content of the image?" + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> messages = [{"role": "user", "content": [{"type": "text", "text": prompt}, {"type": "image", "image": image}]}] + + >>> inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True) + + >>> # Generate + >>> generated_ids = model.generate(**inputs, max_new_tokens=15) + >>> generated_tokens = generated_ids[:, inputs['input_ids'].size(1):] + >>> processor.post_process_image_text_to_text(generated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "The image shows a bustling street scene in what appears to be a Chinatown area. There's ..." + ```""" + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + image_token_pooling=image_token_pooling, + image_grids=image_grids, + image_num_crops=image_num_crops, + pixel_values_videos=pixel_values_videos, + video_token_pooling=video_token_pooling, + video_grids=video_grids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + mm_encoder_outputs=mm_encoder_outputs, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.vocab_size, **kwargs) + + return Molmo2CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + @staticmethod + def create_masks_for_generate( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.Tensor | None, + mm_token_type_ids: torch.Tensor | None = None, + **kwargs, + ) -> dict: + mask_kwargs = { + "config": config.get_text_config(), + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + if mm_token_type_ids is not None and inputs_embeds.shape[1] != 1: + mask_kwargs["block_sequence_ids"] = torch.where(mm_token_type_ids.to(inputs_embeds.device) == 1, 0, -1) + + return create_masks_for_generate(**mask_kwargs) + + +__all__ = [ + "Molmo2Adapter", + "Molmo2ForConditionalGeneration", + "Molmo2Model", + "Molmo2PreTrainedModel", + "Molmo2TextModel", + "Molmo2VisionModel", +] diff --git a/src/transformers/models/molmo2/modular_molmo2.py b/src/transformers/models/molmo2/modular_molmo2.py new file mode 100644 index 000000000000..1dc3b1854a31 --- /dev/null +++ b/src/transformers/models/molmo2/modular_molmo2.py @@ -0,0 +1,1824 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PyTorch Molmo2 model.""" + +import itertools +import math +from collections.abc import Callable + +import torch +from huggingface_hub.dataclasses import strict +from torch import nn +from torch.nn import functional as F + +from ... import initialization as init +from ...backbone_utils import filter_output_hidden_states +from ...cache_utils import Cache, DynamicCache +from ...configuration_utils import PreTrainedConfig, SubConfigSpec +from ...generation import GenerationMixin +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...masking_utils import create_bidirectional_mask, create_causal_mask, create_masks_for_generate +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, RopeParameters +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import ( + ImagesKwargs, + MultiModalData, + ProcessingKwargs, + ProcessorMixin, + Unpack, + VideosKwargs, +) +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import ( + TensorType, + TransformersKwargs, + auto_docstring, + can_return_tuple, + logging, +) +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ...video_processing_utils import BaseVideoProcessor +from ...video_utils import VideoInput, VideoMetadata +from ..cohere_asr.modeling_cohere_asr import CohereAsrCrossAttention +from ..llama.modeling_llama import ( + LlamaDecoderLayer, + LlamaMLP, + LlamaModel, + LlamaPreTrainedModel, + LlamaRMSNorm, + LlamaRotaryEmbedding, + apply_rotary_pos_emb, + eager_attention_forward, +) +from ..llava.modeling_llava import ( + LlavaCausalLMOutputWithPast, + LlavaModel, + LlavaModelOutputWithPast, +) +from ..olmo.modeling_olmo import OlmoMLP +from ..olmo2.modeling_olmo2 import Olmo2Attention +from ..siglip2.modeling_siglip2 import ( + Siglip2EncoderLayer, + Siglip2MLP, +) + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="allenai/Molmo2-8B") +@strict +class Molmo2VisionConfig(PreTrainedConfig): + r""" + image_size (`list[int]`, *optional*): + Input image size as (height, width), `[378, 378]` when not provided. + num_position_embeddings (`int`, *optional*): + Number of positional embeddings, one per image patch. Derived from `image_size` and `patch_size` when not + provided. + """ + + model_type = "molmo2_vision" + base_config_key = "vision_config" + + hidden_size: int = 1152 + intermediate_size: int = 4304 + num_hidden_layers: int = 25 + num_attention_heads: int = 16 + num_key_value_heads: int = 16 + head_dim: int = 72 + hidden_act: str = "gelu_pytorch_tanh" + layer_norm_eps: float = 1e-6 + image_size: list[int] | None = None + patch_size: int = 14 + num_position_embeddings: int | None = None + num_channels: int = 3 + attention_dropout: float = 0.0 + attention_bias: bool = True + initializer_range: float = 0.02 + + def __post_init__(self, **kwargs): + # Keys of the released checkpoints' `config.json`. + kwargs.pop("model_type", None) + self.image_size = kwargs.pop("image_default_input_size", self.image_size) + self.patch_size = kwargs.pop("image_patch_size", self.patch_size) + self.num_position_embeddings = kwargs.pop("image_num_pos", self.num_position_embeddings) + if self.image_size is None: + self.image_size = [378, 378] + if self.num_position_embeddings is None: + self.num_position_embeddings = (self.image_size[0] // self.patch_size) * ( + self.image_size[1] // self.patch_size + ) + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="allenai/Molmo2-8B") +@strict +class Molmo2AdapterConfig(PreTrainedConfig): + r""" + vision_feature_layer (`list[int]`, *optional*): + Indices of the ViT layers whose outputs are concatenated and pooled, `[24, 18]` when not provided. Negative + indices count from the last layer. + text_hidden_size (`int`, *optional*, defaults to 3584): + Hidden size of the text model (used for projection). + """ + + model_type = "molmo2_adapter" + base_config_key = "adapter_config" + + vision_feature_layer: list[int] | None = None + hidden_size: int = 1152 + num_attention_heads: int = 16 + num_key_value_heads: int = 16 + head_dim: int = 72 + attention_dropout: float = 0.0 + attention_bias: bool = True + hidden_act: str = "silu" + intermediate_size: int = 18944 + text_hidden_size: int = 3584 + mlp_bias: bool = False + initializer_range: float = 0.02 + + def __post_init__(self, **kwargs): + # Keys of the released checkpoints' `config.json`. + kwargs.pop("model_type", None) + self.vision_feature_layer = kwargs.pop("vit_layers", self.vision_feature_layer) + if self.vision_feature_layer is None: + self.vision_feature_layer = [24, 18] + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="allenai/Molmo2-8B") +@strict +class Molmo2TextConfig(PreTrainedConfig): + r""" + additional_vocab_size (`int`, *optional*, defaults to 128): + Number of additional vocabulary tokens beyond the base vocabulary. + qk_norm_type (`str`, *optional*, defaults to `"qwen3"`): + Query/key normalization layout used by the checkpoint. `"qwen3"` normalizes per head; `"olmo"` normalizes the + full projected query/key tensors. + rope_scaling_layers (`list[int]`, *optional*): + Indices of the layers that apply the scaled RoPE described by `rope_parameters`. The remaining layers use an + unscaled RoPE with the same theta. All layers are scaled when not provided. + norm_after (`bool`, *optional*, defaults to `False`): + Whether to apply layer normalization after the attention/FFN blocks instead of before. + """ + + model_type = "molmo2_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"qkv_bias": "attention_bias", "layer_norm_eps": "rms_norm_eps"} + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + hidden_size: int = 4096 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + head_dim: int = 128 + vocab_size: int = 151936 + additional_vocab_size: int = 128 + attention_bias: bool = False + qk_norm_type: str = "qwen3" + num_hidden_layers: int = 36 + intermediate_size: int = 12288 + hidden_act: str = "silu" + attention_dropout: float = 0.0 + max_position_embeddings: int = 36864 + rope_parameters: RopeParameters | dict | None = None + rope_scaling_layers: list[int] | None = None + rms_norm_eps: float = 1e-6 + norm_after: bool = False + initializer_range: float = 0.02 + use_cache: bool = True + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + if self.rope_scaling_layers is None: + self.rope_scaling_layers = list(range(self.num_hidden_layers)) + super().__post_init__(**kwargs) + + def validate_architecture(self): + super().validate_architecture() + if self.qk_norm_type not in ("qwen3", "olmo"): + raise ValueError(f"Unsupported `qk_norm_type`: {self.qk_norm_type}") + if self.additional_vocab_size < 0: + raise ValueError(f"`additional_vocab_size` must be >= 0, got {self.additional_vocab_size}") + + +@auto_docstring(checkpoint="allenai/Molmo2-8B") +@strict +class Molmo2Config(PreTrainedConfig): + r""" + vision_config (`Molmo2VisionConfig`, *optional*): + Configuration for the vision transformer backbone. + adapter_config (`Molmo2AdapterConfig`, *optional*): + Configuration for the vision-to-language adapter. + image_start_token_id (`int`, *optional*, defaults to 151936): + Token ID marking the start of an image region. + low_res_image_start_token_id (`int`, *optional*, defaults to 151940): + Token ID marking the start of a low-resolution image crop. + image_end_token_id (`int`, *optional*, defaults to 151937): + Token ID marking the end of an image region. + image_patch_id (`int`, *optional*, defaults to 151938): + Token ID for image patches. + frame_start_token_id (`int`, *optional*, defaults to 151943): + Token ID marking the start of a video frame. + frame_end_token_id (`int`, *optional*, defaults to 151944): + Token ID marking the end of a video frame. + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether the model's input and output word embeddings should be tied. + """ + + model_type = "molmo2" + attribute_map = {"image_token_id": "image_patch_id", "video_token_id": "image_patch_id"} + sub_configs_defaults = { + "text_config": SubConfigSpec(config_class=Molmo2TextConfig), + "vision_config": SubConfigSpec(config_class=Molmo2VisionConfig), + "adapter_config": SubConfigSpec(config_class=Molmo2AdapterConfig), + } + + vision_config: dict | PreTrainedConfig | None = None + adapter_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + image_start_token_id: int = 151936 + low_res_image_start_token_id: int = 151940 + image_end_token_id: int = 151937 + image_patch_id: int = 151938 + frame_start_token_id: int = 151943 + frame_end_token_id: int = 151944 + initializer_range: float = 0.02 + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + # Checkpoints serialized before the rename store the vision sub-config under `vit_config`. + legacy_vision_config = kwargs.pop("vit_config", None) + if self.vision_config is None and legacy_vision_config is not None: + self.vision_config = legacy_vision_config + + super().__post_init__(**kwargs) + + # The hub `config.json` declares a 27-layer ViT, but the released weights stop at layer 24, the deepest one + # `vit_layers` reads. Without this trim, layers 25-26 are built with random weights. Do not remove. + if legacy_vision_config is not None: + num_vit_layers = self.vision_config.num_hidden_layers + self.adapter_config.vision_feature_layer = [ + layer % num_vit_layers for layer in self.adapter_config.vision_feature_layer + ] + self.vision_config.num_hidden_layers = max(self.adapter_config.vision_feature_layer) + 1 + + def validate_architecture(self): + super().validate_architecture() + num_vit_layers = self.vision_config.num_hidden_layers + if not all(-num_vit_layers <= layer < num_vit_layers for layer in self.adapter_config.vision_feature_layer): + raise ValueError( + f"`adapter_config.vision_feature_layer` {self.adapter_config.vision_feature_layer} must index the " + f"{num_vit_layers} layers of `vision_config`." + ) + + +def select_tiling(height: int, width: int, patch_size: int, max_num_crops: int) -> tuple[int, int]: + """Same as `get_optimal_tiled_canvas` in cohere2_vision, with the candidates in (height, width) order and ties + broken on height, as in the original Molmo2 processor.""" + tilings = [ + (tile_height, tile_width) + for tile_height, tile_width in itertools.product(range(1, max_num_crops + 1), repeat=2) + if tile_height * tile_width <= max_num_crops + ] + tilings.sort(key=lambda x: (x[0] * x[1], x[0])) + + candidate_resolutions = torch.tensor(tilings, dtype=torch.int32) * patch_size + original_size = torch.tensor([height, width], dtype=torch.float32) + + required_scales = candidate_resolutions.to(torch.float32) / original_size + required_scale = required_scales.amin(dim=-1, keepdim=True) + + if torch.all(required_scale < 1): + return tilings[int(required_scale.argmax())] + + required_scale = torch.where(required_scale < 1.0, 10e9, required_scale) + return tilings[int(required_scale.argmin())] + + +def compute_crop_geometry( + height: int, + width: int, + max_crops: int, + overlap_margins: list[int], + crop_size: int, + patch_size: int, +) -> tuple[int, int, int, int]: + """Tiling `(tiling_h, tiling_w)` and canvas pixel size `(canvas_height, canvas_width)` for one image.""" + left_margin, right_margin = overlap_margins + crop_patches = crop_size // patch_size + window_patches = crop_patches - (left_margin + right_margin) + window_size = window_patches * patch_size + margin_size = (left_margin + right_margin) * patch_size + tiling_h, tiling_w = select_tiling(height - margin_size, width - margin_size, window_size, max_crops) + canvas_height = tiling_h * window_size + margin_size + canvas_width = tiling_w * window_size + margin_size + return tiling_h, tiling_w, canvas_height, canvas_width + + +def batch_pixels_to_patches(images: torch.Tensor, patch_size: int) -> torch.Tensor: + """Reshape images of [num_images, h, w, channels] -> [num_images, n_patches, pixels_per_patch]""" + num_crops, height, width, channels = images.shape + num_patches_height = height // patch_size + num_patches_width = width // patch_size + images = images.reshape(num_crops, num_patches_height, patch_size, num_patches_width, patch_size, channels) + images = images.transpose(2, 3) + images = images.reshape(num_crops, num_patches_height * num_patches_width, patch_size * patch_size * channels) + return images + + +def arange_for_pooling( + index_grid: torch.Tensor, + pool_h: int, + pool_w: int, +) -> tuple[torch.Tensor, int, int]: + """Group a [h, w] patch-index grid into pooling windows: `[num_rows * num_cols, pool_h * pool_w]`, padded + with -1. Also returns the pooled grid's `(num_rows, num_cols)`.""" + height_padding = pool_h * ((index_grid.shape[0] + pool_h - 1) // pool_h) - index_grid.shape[0] + width_padding = pool_w * ((index_grid.shape[1] + pool_w - 1) // pool_w) - index_grid.shape[1] + index_grid = F.pad( + index_grid, + (width_padding // 2, (width_padding + 1) // 2, height_padding // 2, (height_padding + 1) // 2), + mode="constant", + value=-1, + ) + num_rows, num_cols = index_grid.shape[0] // pool_h, index_grid.shape[1] // pool_w + pooling_indices = ( + index_grid.reshape(num_rows, pool_h, num_cols, pool_w) + .permute(0, 2, 1, 3) + .reshape(num_rows * num_cols, pool_h * pool_w) + ) + return pooling_indices, num_rows, num_cols + + +def resize_and_normalize_image( + backend: "TorchvisionBackend", + image: torch.Tensor, + output_size: list[int], + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], +) -> torch.Tensor: + # Under `torch.compile`, torchvision's uint8 bilinear resize can produce values slightly outside 0-255, + # which overflow and wrap around (-1 becomes 255, so a near-black pixel becomes white). Resizing in + # float, then rounding and clamping back, avoids that and gives the exact same bytes in eager mode. + is_integer_input = not image.dtype.is_floating_point + resized = backend.resize( + image.float() if is_integer_input else image, + size=SizeDict(height=output_size[0], width=output_size[1]), + resample=resample, + antialias=False, + ) + if is_integer_input: + resized = resized.round().clamp(0, 255).to(image.dtype) + return backend.rescale_and_normalize(resized, do_rescale, rescale_factor, do_normalize, image_mean, image_std) + + +def build_resized_image( + backend: "TorchvisionBackend", + images: torch.Tensor, + crop_size: int, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + image_patch_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # `images`: a batch of same-shape images `[N, C, H, W]`; resize the whole batch at once. + resized = resize_and_normalize_image( + backend, + images, + [crop_size, crop_size], + resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + ) + # [N, C, S, S] -> [N, 1, S, S, C]: one global (low-res) view per image. + resized = resized.permute(0, 2, 3, 1).unsqueeze(1) + # The per-patch index grid depends only on the (shared) shape, so it is built once. + patches_per_side = crop_size // image_patch_size + resized_index_grid = torch.arange(patches_per_side**2, dtype=torch.int32, device=images.device).reshape( + patches_per_side, patches_per_side + ) + return resized, resized_index_grid + + +class Molmo2ImagesKwargs(ImagesKwargs, total=False): + """ + max_crops (`int`, *optional*, defaults to 8): + Maximum number of crops to use per image. + overlap_margins (`list[int]`, *optional*, defaults to `[4, 4]`): + Overlap margins (in patches) for overlapping crop extraction. + patch_size (`int`, *optional*, defaults to 14): + The spatial patch size of the vision encoder. + pooling_size (`list[int]`, *optional*, defaults to `[2, 2]`): + The pooling size of the vision adapter. + """ + + max_crops: int + overlap_margins: list[int] + patch_size: int + pooling_size: list[int] + + +@auto_docstring +class Molmo2ImageProcessor(TorchvisionBackend): + valid_kwargs = Molmo2ImagesKwargs + model_input_names = ["pixel_values", "image_token_pooling", "image_grids", "image_num_crops"] + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 378, "width": 378} + do_resize = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + max_crops = 8 + overlap_margins = [4, 4] + patch_size = 14 + pooling_size = [2, 2] + + def __init__(self, **kwargs: Unpack[Molmo2ImagesKwargs]): + super().__init__(**kwargs) + + def _build_overlapping_crops( + self, + images: torch.Tensor, + max_crops: int, + overlap_margins: list[int], + crop_size: int, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + image_patch_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Tile a batch of same-shape images `[N, C, H, W]` into overlapping square crops. The tiling and the per-patch index grid depend only on the (shared) shape, so they are computed once and the resize/unfold are batched over N.""" + left_margin, right_margin = overlap_margins + crop_patches = crop_size // image_patch_size + window_patches = crop_patches - (left_margin + right_margin) + window_size = window_patches * image_patch_size + + num_images, num_channels, original_height, original_width = images.shape + tiling_h, tiling_w, canvas_height, canvas_width = compute_crop_geometry( + original_height, + original_width, + max_crops=max_crops, + overlap_margins=overlap_margins, + crop_size=crop_size, + patch_size=image_patch_size, + ) + + canvas = resize_and_normalize_image( + self, + images, + [canvas_height, canvas_width], + resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + ) + + # [N, C, canvas_height, canvas_width] -> unfold spatial dims -> [N, C, tiling_h, tiling_w, crop, crop] + crops = canvas.unfold(2, crop_size, window_size).unfold(3, crop_size, window_size) + crops = ( + crops.permute(0, 2, 3, 4, 5, 1) + .reshape(num_images, tiling_h * tiling_w, crop_size, crop_size, num_channels) + .contiguous() + ) + + patch_index_grid = torch.arange( + tiling_h * tiling_w * crop_patches * crop_patches, dtype=torch.int32, device=images.device + ).reshape(tiling_h, tiling_w, crop_patches, crop_patches) + if left_margin: + patch_index_grid[1:, :, :left_margin, :] = -1 + patch_index_grid[:, 1:, :, :left_margin] = -1 + if right_margin: + patch_index_grid[:-1, :, -right_margin:, :] = -1 + patch_index_grid[:, :-1, :, -right_margin:] = -1 + + patch_index_grid = patch_index_grid.permute(0, 2, 1, 3).reshape(-1) + patch_index_grid = patch_index_grid[patch_index_grid >= 0].reshape( + canvas_height // image_patch_size, canvas_width // image_patch_size + ) + + return crops, patch_index_grid + + def _image_batch_to_patches_and_grids( + self, + images: torch.Tensor, + max_crops: int, + overlap_margins: list[int], + crop_size: int, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + image_patch_size: int, + image_pooling_w: int, + image_pooling_h: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Process a batch of same-shape images `[N, C, H, W]`. The grid and pooling indices are + shape-dependent (shared across the batch) and returned expanded to `[N, ...]` so the caller + can reorder them per-image alongside the batched patch tensor.""" + patches_per_crop = (crop_size // image_patch_size) ** 2 + + crops, patch_index_grid = self._build_overlapping_crops( + images, + max_crops=max_crops, + overlap_margins=overlap_margins, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=image_patch_size, + ) + pooling_indices, num_patch_rows, num_patch_cols = arange_for_pooling( + patch_index_grid, image_pooling_h, image_pooling_w + ) + + resized, resized_index_grid = build_resized_image( + self, + images, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=image_patch_size, + ) + # [N, 1, S, S, C] + [N, ncrops, S, S, C] -> [N, 1 + ncrops, S, S, C] + crops = torch.cat([resized, crops], dim=1) + + resized_pooling_indices, resized_h, resized_w = arange_for_pooling( + resized_index_grid, image_pooling_h, image_pooling_w + ) + + pooling_indices = torch.where(pooling_indices >= 0, pooling_indices + patches_per_crop, -1) + pooling_indices = torch.cat([resized_pooling_indices, pooling_indices]) + image_grid = torch.tensor( + [[resized_h, resized_w, num_patch_rows, num_patch_cols]], dtype=torch.int64, device=images.device + ) + + # [N, total_crops, S, S, C] -> patches [N, total_crops, n_patch, pixels_per_patch] + num_images, total_crops = crops.shape[0], crops.shape[1] + pixels_per_patch = image_patch_size * image_patch_size * crops.shape[-1] + patches = batch_pixels_to_patches( + crops.reshape(num_images * total_crops, *crops.shape[2:]), image_patch_size + ).reshape(num_images, total_crops, -1, pixels_per_patch) + + # The grid and pooling indices are shared by all images of this shape; expand them to one row + # per image so the batch can later be restored to the original image order. + return image_grid.expand(num_images, -1), patches, pooling_indices.unsqueeze(0).expand(num_images, -1, -1) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + **kwargs: Unpack[Molmo2ImagesKwargs], + ) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + do_convert_rgb: bool, + max_crops: int, + overlap_margins: list[int], + patch_size: int, + pooling_size: list[int], + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + if size.height != size.width: + raise ValueError(f"Molmo2 only supports a square `size`, got height={size.height}, width={size.width}.") + crop_size = size.height + image_pooling_h, image_pooling_w = pooling_size + + # Group images by shape and process each unique shape as a single batch (the tiling and all + # patch/pooling indices are shape-dependent only), then restore the original order. + device = images[0].device + grouped_images, grouped_index = group_images_by_shape(images, disable_grouping=disable_grouping) + + grids_grouped: dict = {} + patches_grouped: dict = {} + pooled_grouped: dict = {} + for shape, stacked_images in grouped_images.items(): + image_grid, patches, pooled_indices = self._image_batch_to_patches_and_grids( + stacked_images, + max_crops=max_crops, + overlap_margins=overlap_margins, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=patch_size, + image_pooling_w=image_pooling_w, + image_pooling_h=image_pooling_h, + ) + grids_grouped[shape] = image_grid + patches_grouped[shape] = patches + pooled_grouped[shape] = pooled_indices + + grids = reorder_images(grids_grouped, grouped_index) + patches = reorder_images(patches_grouped, grouped_index) + pooled = reorder_images(pooled_grouped, grouped_index) + + all_crops: list[torch.Tensor] = [] + all_pooled: list[torch.Tensor] = [] + patch_offset = 0 + for crops, pooled_indices in zip(patches, pooled): + all_pooled.append(torch.where(pooled_indices >= 0, pooled_indices + patch_offset, pooled_indices)) + all_crops.append(crops) + patch_offset += crops.shape[0] * crops.shape[1] + + data = { + "pixel_values": torch.cat(all_crops, dim=0), + "image_token_pooling": torch.cat(all_pooled, dim=0), + "image_grids": torch.stack(grids, dim=0), + "image_num_crops": torch.tensor([crops.shape[0] for crops in patches], dtype=torch.int64, device=device), + } + return BatchFeature(data=data, tensor_type=return_tensors) + + def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None) -> int: + image_grid, _ = self.get_image_grid_and_crops(height, width, images_kwargs) + resized_h, resized_w, high_res_rows, high_res_cols = image_grid + return resized_h * resized_w + high_res_rows * high_res_cols + + def get_image_grid_and_crops(self, height: int, width: int, images_kwargs=None) -> tuple[list[int], int]: + """Return the `[low_res_h, low_res_w, high_res_h, high_res_w]` pooled grid and the crop count for one image.""" + if images_kwargs is None: + images_kwargs = {} + max_crops = images_kwargs.get("max_crops", self.max_crops) + overlap_margins = images_kwargs.get("overlap_margins", self.overlap_margins) + patch_size = images_kwargs.get("patch_size", self.patch_size) + pooling_size = images_kwargs.get("pooling_size", self.pooling_size) + size = images_kwargs.get("size", self.size) + + base_h = size["height"] if isinstance(size, dict) else size.height + base_w = size["width"] if isinstance(size, dict) else size.width + pooling_h, pooling_w = pooling_size + + tiling_h, tiling_w, canvas_height, canvas_width = compute_crop_geometry( + height, + width, + max_crops=max_crops, + overlap_margins=overlap_margins, + crop_size=base_h, + patch_size=patch_size, + ) + num_patch_rows_high = math.ceil((canvas_height // patch_size) / pooling_h) + num_patch_cols_high = math.ceil((canvas_width // patch_size) / pooling_w) + + resized_h = math.ceil((base_h // patch_size) / pooling_h) + resized_w = math.ceil((base_w // patch_size) / pooling_w) + + return [resized_h, resized_w, num_patch_rows_high, num_patch_cols_high], tiling_h * tiling_w + 1 + + +class Molmo2VideosKwargs(VideosKwargs, total=False): + """ + patch_size (`int`, *optional*): + Side length in pixels of each ViT patch for video frames. + pooling_size (`list[int]`, *optional*): + `[pool_h, pool_w]` pooling window applied to video patch features. + max_fps (`int`, *optional*): + Maximum sampling rate in frames per second for short videos. + """ + + patch_size: int + pooling_size: list[int] + max_fps: int + + +@auto_docstring +class Molmo2VideoProcessor(BaseVideoProcessor): + resample = PILImageResampling.BILINEAR + size = {"height": 378, "width": 378} + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + patch_size = 14 + pooling_size = [3, 3] + num_frames = 64 + do_sample_frames = True + max_fps = 2 + valid_kwargs = Molmo2VideosKwargs + model_input_names = ["pixel_values_videos", "video_token_pooling", "video_grids"] + + def __init__(self, **kwargs: Unpack[Molmo2VideosKwargs]): + super().__init__(**kwargs) + if self.size is not None and (self.size.get("height", None) is None or self.size.get("width", None) is None): + raise ValueError("size must contain 'height' and 'width' keys.") + + def sample_frames( + self, + metadata: VideoMetadata, + num_frames: int | None = None, + fps: int | float | None = None, + max_fps: int | float | None = None, + **kwargs, + ): + if fps is not None and num_frames is not None: + raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!") + + max_fps = max_fps if max_fps is not None else self.max_fps + + if metadata.fps is None: + metadata.fps = fps if fps is not None else max_fps + logger.warning_once( + "Molmo2 inserts frame timestamps into video prompts, but the input video's `fps` was not provided " + f"or could not be inferred. Defaulting to `fps={metadata.fps}`. Please provide `video_metadata` " + "for more accurate timestamps." + ) + if metadata.duration is None: + metadata.duration = metadata.total_num_frames / metadata.fps + + if fps is not None: + target_num_frames = int(metadata.duration * fps) + else: + target_num_frames = num_frames if num_frames is not None else self.num_frames + if max_fps is not None and metadata.fps > max_fps: + target_num_frames = min(target_num_frames, int(metadata.duration * max_fps)) + + target_num_frames = max(min(target_num_frames, metadata.total_num_frames), 1) + total = metadata.total_num_frames + return torch.arange(0, total, total / target_num_frames).int() + + def _build_video_patches( + self, + video: torch.Tensor, + crop_size: int, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + image_patch_size: int, + image_pooling_h: int, + image_pooling_w: int, + ) -> tuple[list[int], torch.Tensor, torch.Tensor]: + # `build_resized_image` takes a whole `[N, C, H, W]` batch, so all frames of a video are resized in one call. + resized_frames, resized_index_grid = build_resized_image( + self, + video, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=image_patch_size, + ) + # The pooling index grid is shape-dependent only, hence shared by every frame. + pooling_indices, num_patch_rows, num_patch_cols = arange_for_pooling( + resized_index_grid, image_pooling_h, image_pooling_w + ) + # [T, 1, S, S, C] -> [T, n_patch, pixels_per_patch] + patches = batch_pixels_to_patches(resized_frames.squeeze(1), image_patch_size) + return [num_patch_rows, num_patch_cols], patches, pooling_indices + + def _preprocess( + self, + videos: list["torch.Tensor"], + size: SizeDict, + resample: PILImageResampling | None, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + patch_size: int, + pooling_size: list[int], + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchFeature: + if size.height != size.width: + raise ValueError(f"Molmo2 only supports a square `size`, got height={size.height}, width={size.width}.") + crop_size = size.height + image_pooling_h, image_pooling_w = pooling_size + + all_crops: list[torch.Tensor] = [] + all_pooled: list[torch.Tensor] = [] + all_grids: list[torch.Tensor] = [] + patch_offset = 0 + + for video in videos: + image_grid, patches, pooling_indices = self._build_video_patches( + video, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=patch_size, + image_pooling_h=image_pooling_h, + image_pooling_w=image_pooling_w, + ) + num_frames, patches_per_frame = patches.shape[:2] + frame_offsets = ( + patch_offset + torch.arange(num_frames, device=pooling_indices.device) * patches_per_frame + ).view(-1, 1, 1) + pooled_indices = torch.where(pooling_indices >= 0, pooling_indices + frame_offsets, pooling_indices) + all_pooled.append(pooled_indices.reshape(-1, pooling_indices.shape[-1])) + all_crops.append(patches) + patch_offset += num_frames * patches_per_frame + + all_grids.append(torch.tensor([num_frames, image_grid[0], image_grid[1]], dtype=torch.int64)) + + data = { + "pixel_values_videos": torch.cat(all_crops, dim=0), + "video_token_pooling": torch.cat(all_pooled, dim=0), + "video_grids": torch.stack(all_grids, dim=0), + } + return BatchFeature(data, tensor_type=return_tensors) + + +class Molmo2ProcessorKwargs(ProcessingKwargs, total=False): + """Molmo2 processor kwargs""" + + images_kwargs: Molmo2ImagesKwargs + videos_kwargs: Molmo2VideosKwargs + _defaults = { + "text_kwargs": { + "padding": False, + "return_mm_token_type_ids": True, + }, + "videos_kwargs": {"return_metadata": True}, + } + + +@auto_docstring +class Molmo2Processor(ProcessorMixin): + valid_processor_kwargs = Molmo2ProcessorKwargs + + @property + def model_input_names(self): + return super().model_input_names + ["mm_token_type_ids"] + + def __init__( + self, + image_processor=None, + video_processor=None, + tokenizer=None, + chat_template: str | None = None, + image_use_col_tokens: bool | None = True, + use_single_crop_col_tokens: bool | None = None, + use_single_crop_start_token: bool | None = True, + video_use_col_tokens: bool | None = False, + use_frame_special_tokens: bool | None = True, + **kwargs, + ) -> None: + r""" + image_use_col_tokens (`bool`, *optional*, defaults to `True`): + Whether to append column-separator tokens (``) after each patch row of the high-resolution image + view. + use_single_crop_col_tokens (`bool`, *optional*): + Whether to append column-separator tokens after each patch row of the low-resolution (single-crop) image + view. If `None`, falls back to `image_use_col_tokens`. + use_single_crop_start_token (`bool`, *optional*, defaults to `True`): + Whether to start the low-resolution image view with `` instead of the regular + ``. + video_use_col_tokens (`bool`, *optional*, defaults to `False`): + Whether to append column-separator tokens after each patch row of video frames. + use_frame_special_tokens (`bool`, *optional*, defaults to `True`): + Whether to wrap each video frame with `` / `` tokens. If `False`, falls back to + `` / ``. + """ + self.image_use_col_tokens = image_use_col_tokens + self.use_single_crop_col_tokens = use_single_crop_col_tokens + self.use_single_crop_start_token = use_single_crop_start_token + self.video_use_col_tokens = video_use_col_tokens + self.use_frame_special_tokens = use_frame_special_tokens + self.image_token = getattr(tokenizer, "image_token", "<|image|>") + self.video_token = getattr(tokenizer, "video_token", "<|video|>") + self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) + self.video_token_id = tokenizer.convert_tokens_to_ids(self.video_token) + self.image_token_ids = tokenizer.convert_tokens_to_ids( + [ + "", + "", + "", + "", + "", + "", + "", + "", + ] + ) + super().__init__(image_processor, video_processor, tokenizer, chat_template=chat_template) + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + videos: VideoInput | None = None, + **kwargs: Unpack[Molmo2ProcessorKwargs], + ): + # TODO(molbap): remove once the allenai chat templates emit `bos_token` + # (hub PRs allenai/Molmo2-8B#11, allenai/Molmo2-4B#4, allenai/Molmo2-O-7B#2) + bos_token = self.tokenizer.bos_token or self.tokenizer.eos_token + if isinstance(text, str): + text = [text] + if text is not None: + text = [prompt if prompt.startswith(bos_token) else bos_token + prompt for prompt in text] + return super().__call__(images=images, text=text, videos=videos, **kwargs) + + def get_video_string(self, video_grid, timestamps) -> str: + start_token = "" if self.use_frame_special_tokens else "" + end_token = "" if self.use_frame_special_tokens else "" + + num_frames, num_patch_rows, num_patch_cols = video_grid.tolist() + video_string = "" + for frame_idx, frame_time in enumerate(timestamps): + prev_space = " " if frame_idx > 0 else "" + video_string += prev_space + f"{frame_time:.1f} " + tokens_per_row = [""] * num_patch_cols + if self.video_use_col_tokens: + tokens_per_row = tokens_per_row + [""] + video_string += "".join([start_token] + tokens_per_row * num_patch_rows + [end_token]) + + return video_string + + def validate_inputs( + self, + images=None, + text=None, + videos=None, + audio=None, + **kwargs, + ): + super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs) + if videos is not None and text is not None: + for sample in text: + if sample.count(self.video_token) > 1: + raise ValueError("At most one video is supported per sample.") + + def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> str: + image_grid = image_inputs["image_grids"][image_idx] + if not isinstance(image_grid, list): + image_grid = image_grid.tolist() + resized_h, resized_w, height, width = image_grid + + tokens_per_row = [""] * width + if self.image_use_col_tokens: + tokens_per_row = tokens_per_row + [""] + high_res_tokens = [""] + tokens_per_row * height + [""] + + tokens_per_row = [""] * resized_w + use_single_crop_col_tokens = ( + self.image_use_col_tokens if self.use_single_crop_col_tokens is None else self.use_single_crop_col_tokens + ) + image_start_token = "" if self.use_single_crop_start_token else "" + if use_single_crop_col_tokens: + tokens_per_row = tokens_per_row + [""] + low_res_tokens = [image_start_token] + tokens_per_row * resized_h + [""] + + return "".join(low_res_tokens + high_res_tokens) + + def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + + Args: + image_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (height, width) per each image. + + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + vision_data = {} + if image_sizes is not None: + images_kwargs = Molmo2ProcessorKwargs._defaults.get("images_kwargs", {}) + images_kwargs.update(kwargs) + + num_image_tokens, num_image_patches = [], [] + for image_size in image_sizes: + image_grid, num_crops = self.image_processor.get_image_grid_and_crops(*image_size, images_kwargs) + image_string = self.replace_image_token({"image_grids": [image_grid]}, image_idx=0) + num_image_tokens.append(len(self.tokenizer.tokenize(image_string))) + num_image_patches.append(num_crops) + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + return MultiModalData(**vision_data) + + def replace_video_token(self, video_inputs: dict, video_idx: int, **kwargs) -> str: + video_grid = video_inputs["video_grids"][video_idx] + metadata = video_inputs["video_metadata"][video_idx] + + frames_indices = metadata.frames_indices + if frames_indices is None: + frames_indices = range(int(video_grid[0].item())) + fps = metadata.fps + if fps is None: + fps = self.video_processor.max_fps + logger.warning_once( + "Molmo2 inserts frame timestamps into video prompts, but the input video's `fps` was not " + f"provided or could not be inferred. Defaulting to `fps={fps}`. Please provide " + "`video_metadata` for more accurate timestamps." + ) + timestamps = [frame_idx / fps for frame_idx in frames_indices] + return self.get_video_string(video_grid, timestamps) + + +class Molmo2CausalLMOutputWithPast(LlavaCausalLMOutputWithPast): + pass + + +class Molmo2ModelOutputWithPast(LlavaModelOutputWithPast): + pass + + +class Molmo2RotaryEmbedding(LlamaRotaryEmbedding): + """RoPE from the flat `rope_parameters`; `scaled=False` keeps the theta but skips the scaling, for the + layers outside `rope_scaling_layers`.""" + + def __init__(self, config: Molmo2TextConfig, scaled: bool = True): + nn.Module.__init__(self) + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + + # `rope_type = "default"` also keeps the generic `_init_weights` rotary re-init on the unscaled table. + self.rope_type = self.config.rope_parameters["rope_type"] if scaled else "default" + rope_init_fn: Callable = ( + self.compute_default_rope_parameters + if self.rope_type == "default" + else ROPE_INIT_FUNCTIONS[self.rope_type] + ) + inv_freq, self.attention_scaling = rope_init_fn(self.config) + + self.inv_freq = nn.Buffer(inv_freq, persistent=False) + self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False) + + +class Molmo2RMSNorm(LlamaRMSNorm): + pass + + +class Molmo2Attention(Olmo2Attention): + """Olmo2 attention whose q/k RMSNorm runs either over the full projection (`olmo`) or per head (`qwen3`).""" + + def __init__(self, config: Molmo2TextConfig, layer_idx: int) -> None: + super().__init__(config, layer_idx) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + self.qk_norm_type = config.qk_norm_type + self.q_norm = Molmo2RMSNorm( + self.head_dim if config.qk_norm_type == "qwen3" else config.num_attention_heads * self.head_dim, + eps=config.rms_norm_eps, + ) + self.k_norm = Molmo2RMSNorm( + self.head_dim if config.qk_norm_type == "qwen3" else config.num_key_value_heads * self.head_dim, + eps=config.rms_norm_eps, + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + if self.qk_norm_type == "olmo": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.view(hidden_shape) + key_states = key_states.view(hidden_shape) + + if self.qk_norm_type == "qwen3": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Molmo2MLP(OlmoMLP): + pass + + +class Molmo2DecoderLayer(LlamaDecoderLayer): + def __init__(self, config: Molmo2TextConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.norm_after = config.norm_after + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + # `norm_after=True` normalizes each sublayer's output (post-norm) instead of its input (pre-norm) + residual = hidden_states + if not self.norm_after: + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + if self.norm_after: + hidden_states = self.input_layernorm(hidden_states) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + if not self.norm_after: + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + if self.norm_after: + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class Molmo2PreTrainedModel(LlamaPreTrainedModel): + config: Molmo2Config + input_modalities = ("image", "video", "text") + + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + if isinstance(module, Molmo2VisionModel): + init.normal_(module.positional_embedding, mean=0.0, std=self.config.initializer_range) + + +class Molmo2VisionMLP(Siglip2MLP): + pass + + +class Molmo2VisionAttention(CohereAsrCrossAttention): + def __init__(self, config: Molmo2VisionConfig | Molmo2AdapterConfig, hidden_size: int | None = None): + super().__init__(config, layer_idx=None) + del self.layer_idx + self.q_proj = nn.Linear( + hidden_size or config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + hidden_size or config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + hidden_size or config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + encoder_hidden_states = hidden_states if encoder_hidden_states is None else encoder_hidden_states + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + key_value_shape = (*encoder_hidden_states.shape[:-1], -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(encoder_hidden_states).view(key_value_shape).transpose(1, 2) + value_states = self.v_proj(encoder_hidden_states).view(key_value_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + is_causal=self.is_causal, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + +class Molmo2VisionEncoderLayer(Siglip2EncoderLayer): + def __init__(self, config: Molmo2VisionConfig): + super().__init__(config) + self.self_attn = Molmo2VisionAttention(config) + self.mlp = Molmo2VisionMLP(config) + + +@auto_docstring +class Molmo2VisionModel(Molmo2PreTrainedModel): + config_class = Molmo2VisionConfig + main_input_name = "pixel_values" + input_modalities = ("image",) + _no_split_modules = ["Molmo2VisionEncoderLayer"] + _can_record_outputs = { + "hidden_states": Molmo2VisionEncoderLayer, + "attentions": Molmo2VisionAttention, + } + + def __init__(self, config: Molmo2VisionConfig): + super().__init__(config) + self.positional_embedding = nn.Parameter(torch.zeros(config.num_position_embeddings, config.hidden_size)) + self.patch_embedding = nn.Linear( + config.patch_size * config.patch_size * config.num_channels, config.hidden_size + ) + # trf-ignore: TRF034 (false positive: Siglip2EncoderLayer subclasses GradientCheckpointingLayer) + self.layers = nn.ModuleList([Molmo2VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.post_init() + + @merge_with_config_defaults + @capture_outputs(tie_last_hidden_states=False) + @auto_docstring + def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutputWithPooling: + hidden_states = self.patch_embedding(pixel_values.to(dtype=self.dtype)) + # patch count == num_position_embeddings, locked by config; only retraining with a different grid breaks this. + hidden_states = hidden_states + self.positional_embedding[None, :, :].to(hidden_states.dtype) + + for encoder_layer in self.layers: + hidden_states = encoder_layer(hidden_states, None, **kwargs) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=hidden_states.mean(dim=1), + ) + + +class Molmo2ImageProjectorMLP(LlamaMLP): + def __init__(self, config: Molmo2AdapterConfig): + super().__init__(config) + self.down_proj = nn.Linear(config.intermediate_size, config.text_hidden_size, bias=False) + + +@auto_docstring( + custom_intro=""" + The Molmo2 vision adapter: pools ViT patch features with a cross-attention layer and projects them into the + language model's embedding space. + """ +) +class Molmo2Adapter(Molmo2PreTrainedModel): + config_class = Molmo2AdapterConfig + input_modalities = ("image",) + _no_split_modules = ["Molmo2VisionAttention"] + _can_record_outputs = {"attentions": Molmo2VisionAttention} + + def __init__(self, config: Molmo2AdapterConfig): + super().__init__(config) + pooling_input_dim = config.hidden_size * len(config.vision_feature_layer) + self.image_pooling_2d = Molmo2VisionAttention(config, hidden_size=pooling_input_dim) + self.image_projector = Molmo2ImageProjectorMLP(config) + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + image_features: torch.Tensor, + pooled_patches_idx: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + r""" + image_features (`torch.Tensor` of shape `(num_crops, num_patches, hidden_size * len(vision_feature_layer))`): + Concatenated intermediate ViT features of every crop. + pooled_patches_idx (`torch.Tensor` of shape `(num_tokens, pool_h * pool_w)`): + Indices into the flattened patch sequence pooled by each output token; `-1` marks padding slots. + """ + flat_features = image_features.reshape(-1, image_features.shape[-1]) + + valid_mask = pooled_patches_idx >= 0 + valid_token_mask = torch.any(valid_mask, -1) + + patches_to_pool = flat_features[torch.clip(pooled_patches_idx, 0)] + patches_to_pool = patches_to_pool * valid_mask.to(patches_to_pool.dtype)[..., None] + + num_valid_patches = valid_mask.float().sum(-1) + num_valid_patches = torch.where(num_valid_patches == 0, 1, num_valid_patches) + query = patches_to_pool.sum(-2, keepdim=True) / num_valid_patches[:, None, None].to(patches_to_pool.dtype) + + attention_mask = create_bidirectional_mask( + config=self.config, inputs_embeds=query, attention_mask=valid_mask, encoder_hidden_states=patches_to_pool + ) + pooled_features, _ = self.image_pooling_2d(query, patches_to_pool, attention_mask=attention_mask, **kwargs) + pooled_features = pooled_features.squeeze(1) + pooled_features = self.image_projector(pooled_features) + return BaseModelOutput(last_hidden_state=pooled_features[valid_token_mask]) + + +class Molmo2TextModel(LlamaModel): + config: Molmo2TextConfig + + def __init__(self, config: Molmo2TextConfig): + Molmo2PreTrainedModel.__init__(self, config) + self.vocab_size = config.vocab_size + # The checkpoint's extra-vocabulary table is concatenated onto the base one at load time + # (see `conversion_mapping.py`), so the embedding covers `vocab_size + additional_vocab_size`. + self.embed_tokens = nn.Embedding(config.vocab_size + config.additional_vocab_size, config.hidden_size) + # trf-ignore: TRF034 (false positive: LlamaDecoderLayer subclasses GradientCheckpointingLayer) + self.layers = nn.ModuleList( + [Molmo2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Molmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Molmo2RotaryEmbedding(config) + self.rotary_emb_unscaled = Molmo2RotaryEmbedding(config, scaled=False) + self.rope_types = [ + "scaled" if layer_idx in config.rope_scaling_layers else "unscaled" + for layer_idx in range(config.num_hidden_layers) + ] + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = { + "scaled": self.rotary_emb(hidden_states, position_ids=position_ids), + "unscaled": self.rotary_emb_unscaled(hidden_states, position_ids=position_ids), + } + + for layer_idx, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings[self.rope_types[layer_idx]], + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring( + custom_intro=""" + The Molmo2 model which consists of a vision backbone, a pooling adapter and a language model, without a language + modeling head. + """ +) +class Molmo2Model(LlavaModel): + config: Molmo2Config + + def __init__(self, config: Molmo2Config): + Molmo2PreTrainedModel.__init__(self, config) + self.vision_tower = Molmo2VisionModel(config.vision_config) + self.multi_modal_projector = Molmo2Adapter(config.adapter_config) + self.language_model = Molmo2TextModel(config.text_config) + self.post_init() + + @can_return_tuple + @filter_output_hidden_states + @auto_docstring( + custom_intro="Obtains pooled image features from the vision tower and the adapter: the `pooler_output` is a " + "tuple with one `(num_image_tokens, hidden_size)` tensor per image." + ) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_token_pooling: torch.Tensor, + image_grids: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + image_shape = None + if pixel_values.dim() == 4: + image_shape = pixel_values.shape[:3] + pixel_values = pixel_values.reshape(-1, pixel_values.shape[-2], pixel_values.shape[-1]) + + image_outputs: BaseModelOutputWithPooling = self.vision_tower(pixel_values, **kwargs) + image_features = torch.cat( + [ + image_outputs.hidden_states[layer + 1 if layer >= 0 else layer] + for layer in self.config.adapter_config.vision_feature_layer + ], + dim=-1, + ) + + if image_shape is not None: + image_features = image_features.reshape(*image_shape, -1) + pooled_features = self.multi_modal_projector(image_features, image_token_pooling).last_hidden_state + split_sizes = (image_grids[:, 0] * image_grids[:, 1] + image_grids[:, 2] * image_grids[:, 3]).tolist() + image_outputs.pooler_output = torch.split(pooled_features, split_sizes) + return image_outputs + + @can_return_tuple + @auto_docstring( + custom_intro="Obtains pooled video features from the vision tower and the adapter: the `pooler_output` is a " + "tuple with one `(num_video_tokens, hidden_size)` tensor per video." + ) + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_token_pooling: torch.Tensor, + video_grids: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + # Video frames only go through the low-res path: `num_frames` frames of `rows x cols` tokens, no high-res crops. + num_frames, rows, cols = video_grids.unbind(-1) + image_grids = torch.stack([num_frames * rows, cols, torch.zeros_like(cols), torch.zeros_like(cols)], dim=-1) + return self.get_image_features(pixel_values_videos, video_token_pooling, image_grids, **kwargs) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + image_token_pooling: torch.Tensor | None = None, + image_grids: torch.Tensor | None = None, + image_num_crops: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_token_pooling: torch.Tensor | None = None, + video_grids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + past_key_values: Cache | None = None, + mm_token_type_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Molmo2ModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if pixel_values is not None and pixel_values_videos is not None: + raise ValueError("pixel_values and pixel_values_videos are provided at the same time") + if (pixel_values is not None or pixel_values_videos is not None) and mm_encoder_outputs is not None: + raise ValueError( + "You cannot specify both pixel_values/pixel_values_videos and mm_encoder_outputs at the same time" + ) + if self.training and mm_token_type_ids is None: + raise ValueError("`mm_token_type_ids` is required as a model input when training") + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + mm_encoder_outputs = mm_encoder_outputs if mm_encoder_outputs is not None else {} + if mm_encoder_outputs.get("image") is None and pixel_values is not None: + mm_encoder_outputs["image"] = self.get_image_features(pixel_values, image_token_pooling, image_grids) + if mm_encoder_outputs.get("video") is None and pixel_values_videos is not None: + mm_encoder_outputs["video"] = self.get_video_features( + pixel_values_videos, video_token_pooling, video_grids + ) + + image_features: torch.FloatTensor | None = None + if mm_encoder_outputs.get("image") is not None: + image_features = mm_encoder_outputs["image"].pooler_output + elif mm_encoder_outputs.get("video") is not None: + image_features = mm_encoder_outputs["video"].pooler_output + + if image_features is not None: + image_features = torch.cat(image_features, dim=0) + # `get_placeholder_mask` returns a [batch, seq, 1] mask; Molmo2 *adds* the image features onto + # the placeholder-token embeddings (residual), which means we index `inputs_embeds` directly. + # Boolean indexing does not broadcast, so the mask must be expanded to the hidden dim first. + special_image_mask = self.get_placeholder_mask(input_ids, inputs_embeds, image_features) + special_image_mask = special_image_mask.expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter( + special_image_mask, + inputs_embeds[special_image_mask] + image_features.reshape(-1), + ) + + # An already prepared 4D mask (e.g. from `generate`) is returned as is by `create_causal_mask` + mask_kwargs = { + "config": self.config.get_text_config(), + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + is_prefill = past_key_values is None or not past_key_values.is_initialized or image_features is not None + if mm_token_type_ids is not None and is_prefill: + # every image or frame token attends to every other one, so they all share a single block + mask_kwargs["block_sequence_ids"] = torch.where(mm_token_type_ids.to(inputs_embeds.device) == 1, 0, -1) + causal_mask = create_causal_mask(**mask_kwargs) + + outputs = self.language_model( + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + return Molmo2ModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features, + ) + + +@auto_docstring( + custom_intro=""" + The Molmo2 model which consists of a vision backbone, a pooling adapter and a language model. + """ +) +class Molmo2ForConditionalGeneration(Molmo2PreTrainedModel, GenerationMixin): + # the loss must not be normalized by `num_items_in_batch`: logits/labels are filtered before computing it + accepts_loss_kwargs = False + config: Molmo2Config + + def __init__(self, config: Molmo2Config): + super().__init__(config) + + self.model = Molmo2Model(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.vocab_size = config.text_config.vocab_size + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_token_pooling: torch.Tensor, + image_grids: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values, image_token_pooling, image_grids, **kwargs) + + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_token_pooling: torch.Tensor, + video_grids: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_video_features(pixel_values_videos, video_token_pooling, video_grids, **kwargs) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + pixel_values: torch.Tensor | None = None, + image_token_pooling: torch.Tensor | None = None, + image_grids: torch.Tensor | None = None, + image_num_crops: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_token_pooling: torch.Tensor | None = None, + video_grids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + mm_token_type_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Molmo2CausalLMOutputWithPast: + r""" + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, Molmo2ForConditionalGeneration + + >>> model = Molmo2ForConditionalGeneration.from_pretrained("allenai/Molmo2-8B") + >>> processor = AutoProcessor.from_pretrained("allenai/Molmo2-8B") + + >>> prompt = "What's the content of the image?" + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> messages = [{"role": "user", "content": [{"type": "text", "text": prompt}, {"type": "image", "image": image}]}] + + >>> inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True) + + >>> # Generate + >>> generated_ids = model.generate(**inputs, max_new_tokens=15) + >>> generated_tokens = generated_ids[:, inputs['input_ids'].size(1):] + >>> processor.post_process_image_text_to_text(generated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "The image shows a bustling street scene in what appears to be a Chinatown area. There's ..." + ```""" + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + image_token_pooling=image_token_pooling, + image_grids=image_grids, + image_num_crops=image_num_crops, + pixel_values_videos=pixel_values_videos, + video_token_pooling=video_token_pooling, + video_grids=video_grids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + mm_encoder_outputs=mm_encoder_outputs, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.vocab_size, **kwargs) + + return Molmo2CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) + + @staticmethod + def create_masks_for_generate( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.Tensor | None, + mm_token_type_ids: torch.Tensor | None = None, + **kwargs, + ) -> dict: + mask_kwargs = { + "config": config.get_text_config(), + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + if mm_token_type_ids is not None and inputs_embeds.shape[1] != 1: + mask_kwargs["block_sequence_ids"] = torch.where(mm_token_type_ids.to(inputs_embeds.device) == 1, 0, -1) + + return create_masks_for_generate(**mask_kwargs) + + +__all__ = [ + "Molmo2AdapterConfig", + "Molmo2Config", + "Molmo2TextConfig", + "Molmo2VisionConfig", + "Molmo2Adapter", + "Molmo2ForConditionalGeneration", + "Molmo2ImageProcessor", + "Molmo2Model", + "Molmo2PreTrainedModel", + "Molmo2Processor", + "Molmo2TextModel", + "Molmo2VideoProcessor", + "Molmo2VisionModel", +] diff --git a/src/transformers/models/molmo2/processing_molmo2.py b/src/transformers/models/molmo2/processing_molmo2.py new file mode 100644 index 000000000000..5af78c17d80e --- /dev/null +++ b/src/transformers/models/molmo2/processing_molmo2.py @@ -0,0 +1,253 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/molmo2/modular_molmo2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_molmo2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...image_utils import ImageInput +from ...processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring, logging +from ...video_utils import VideoInput + + +logger = logging.get_logger(__name__) + + +class Molmo2ImagesKwargs(ImagesKwargs, total=False): + """ + max_crops (`int`, *optional*, defaults to 8): + Maximum number of crops to use per image. + overlap_margins (`list[int]`, *optional*, defaults to `[4, 4]`): + Overlap margins (in patches) for overlapping crop extraction. + patch_size (`int`, *optional*, defaults to 14): + The spatial patch size of the vision encoder. + pooling_size (`list[int]`, *optional*, defaults to `[2, 2]`): + The pooling size of the vision adapter. + """ + + max_crops: int + overlap_margins: list[int] + patch_size: int + pooling_size: list[int] + + +class Molmo2VideosKwargs(VideosKwargs, total=False): + """ + patch_size (`int`, *optional*): + Side length in pixels of each ViT patch for video frames. + pooling_size (`list[int]`, *optional*): + `[pool_h, pool_w]` pooling window applied to video patch features. + max_fps (`int`, *optional*): + Maximum sampling rate in frames per second for short videos. + """ + + patch_size: int + pooling_size: list[int] + max_fps: int + + +class Molmo2ProcessorKwargs(ProcessingKwargs, total=False): + """Molmo2 processor kwargs""" + + images_kwargs: Molmo2ImagesKwargs + videos_kwargs: Molmo2VideosKwargs + _defaults = { + "text_kwargs": { + "padding": False, + "return_mm_token_type_ids": True, + }, + "videos_kwargs": {"return_metadata": True}, + } + + +@auto_docstring +class Molmo2Processor(ProcessorMixin): + valid_processor_kwargs = Molmo2ProcessorKwargs + + @property + def model_input_names(self): + return super().model_input_names + ["mm_token_type_ids"] + + def __init__( + self, + image_processor=None, + video_processor=None, + tokenizer=None, + chat_template: str | None = None, + image_use_col_tokens: bool | None = True, + use_single_crop_col_tokens: bool | None = None, + use_single_crop_start_token: bool | None = True, + video_use_col_tokens: bool | None = False, + use_frame_special_tokens: bool | None = True, + **kwargs, + ) -> None: + r""" + image_use_col_tokens (`bool`, *optional*, defaults to `True`): + Whether to append column-separator tokens (``) after each patch row of the high-resolution image + view. + use_single_crop_col_tokens (`bool`, *optional*): + Whether to append column-separator tokens after each patch row of the low-resolution (single-crop) image + view. If `None`, falls back to `image_use_col_tokens`. + use_single_crop_start_token (`bool`, *optional*, defaults to `True`): + Whether to start the low-resolution image view with `` instead of the regular + ``. + video_use_col_tokens (`bool`, *optional*, defaults to `False`): + Whether to append column-separator tokens after each patch row of video frames. + use_frame_special_tokens (`bool`, *optional*, defaults to `True`): + Whether to wrap each video frame with `` / `` tokens. If `False`, falls back to + `` / ``. + """ + self.image_use_col_tokens = image_use_col_tokens + self.use_single_crop_col_tokens = use_single_crop_col_tokens + self.use_single_crop_start_token = use_single_crop_start_token + self.video_use_col_tokens = video_use_col_tokens + self.use_frame_special_tokens = use_frame_special_tokens + self.image_token = getattr(tokenizer, "image_token", "<|image|>") + self.video_token = getattr(tokenizer, "video_token", "<|video|>") + self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) + self.video_token_id = tokenizer.convert_tokens_to_ids(self.video_token) + self.image_token_ids = tokenizer.convert_tokens_to_ids( + [ + "", + "", + "", + "", + "", + "", + "", + "", + ] + ) + super().__init__(image_processor, video_processor, tokenizer, chat_template=chat_template) + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None, + videos: VideoInput | None = None, + **kwargs: Unpack[Molmo2ProcessorKwargs], + ): + # TODO(molbap): remove once the allenai chat templates emit `bos_token` + # (hub PRs allenai/Molmo2-8B#11, allenai/Molmo2-4B#4, allenai/Molmo2-O-7B#2) + bos_token = self.tokenizer.bos_token or self.tokenizer.eos_token + if isinstance(text, str): + text = [text] + if text is not None: + text = [prompt if prompt.startswith(bos_token) else bos_token + prompt for prompt in text] + return super().__call__(images=images, text=text, videos=videos, **kwargs) + + def get_video_string(self, video_grid, timestamps) -> str: + start_token = "" if self.use_frame_special_tokens else "" + end_token = "" if self.use_frame_special_tokens else "" + + num_frames, num_patch_rows, num_patch_cols = video_grid.tolist() + video_string = "" + for frame_idx, frame_time in enumerate(timestamps): + prev_space = " " if frame_idx > 0 else "" + video_string += prev_space + f"{frame_time:.1f} " + tokens_per_row = [""] * num_patch_cols + if self.video_use_col_tokens: + tokens_per_row = tokens_per_row + [""] + video_string += "".join([start_token] + tokens_per_row * num_patch_rows + [end_token]) + + return video_string + + def validate_inputs( + self, + images=None, + text=None, + videos=None, + audio=None, + **kwargs, + ): + super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs) + if videos is not None and text is not None: + for sample in text: + if sample.count(self.video_token) > 1: + raise ValueError("At most one video is supported per sample.") + + def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> str: + image_grid = image_inputs["image_grids"][image_idx] + if not isinstance(image_grid, list): + image_grid = image_grid.tolist() + resized_h, resized_w, height, width = image_grid + + tokens_per_row = [""] * width + if self.image_use_col_tokens: + tokens_per_row = tokens_per_row + [""] + high_res_tokens = [""] + tokens_per_row * height + [""] + + tokens_per_row = [""] * resized_w + use_single_crop_col_tokens = ( + self.image_use_col_tokens if self.use_single_crop_col_tokens is None else self.use_single_crop_col_tokens + ) + image_start_token = "" if self.use_single_crop_start_token else "" + if use_single_crop_col_tokens: + tokens_per_row = tokens_per_row + [""] + low_res_tokens = [image_start_token] + tokens_per_row * resized_h + [""] + + return "".join(low_res_tokens + high_res_tokens) + + def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs): + """ + Computes the number of placeholder tokens needed for multimodal inputs with the given sizes. + + Args: + image_sizes (`list[list[int]]`, *optional*): + The input sizes formatted as (height, width) per each image. + + Returns: + `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided + input modalities, along with other useful data. + """ + vision_data = {} + if image_sizes is not None: + images_kwargs = Molmo2ProcessorKwargs._defaults.get("images_kwargs", {}) + images_kwargs.update(kwargs) + + num_image_tokens, num_image_patches = [], [] + for image_size in image_sizes: + image_grid, num_crops = self.image_processor.get_image_grid_and_crops(*image_size, images_kwargs) + image_string = self.replace_image_token({"image_grids": [image_grid]}, image_idx=0) + num_image_tokens.append(len(self.tokenizer.tokenize(image_string))) + num_image_patches.append(num_crops) + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + return MultiModalData(**vision_data) + + def replace_video_token(self, video_inputs: dict, video_idx: int, **kwargs) -> str: + video_grid = video_inputs["video_grids"][video_idx] + metadata = video_inputs["video_metadata"][video_idx] + + frames_indices = metadata.frames_indices + if frames_indices is None: + frames_indices = range(int(video_grid[0].item())) + fps = metadata.fps + if fps is None: + fps = self.video_processor.max_fps + logger.warning_once( + "Molmo2 inserts frame timestamps into video prompts, but the input video's `fps` was not " + f"provided or could not be inferred. Defaulting to `fps={fps}`. Please provide " + "`video_metadata` for more accurate timestamps." + ) + timestamps = [frame_idx / fps for frame_idx in frames_indices] + return self.get_video_string(video_grid, timestamps) + + +__all__ = ["Molmo2Processor"] diff --git a/src/transformers/models/molmo2/video_processing_molmo2.py b/src/transformers/models/molmo2/video_processing_molmo2.py new file mode 100644 index 000000000000..c1794ec87e78 --- /dev/null +++ b/src/transformers/models/molmo2/video_processing_molmo2.py @@ -0,0 +1,279 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/molmo2/modular_molmo2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_molmo2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +from torch.nn import functional as F + +from ...image_processing_backends import TorchvisionBackend +from ...image_processing_utils import BatchFeature +from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, PILImageResampling, SizeDict +from ...processing_utils import Unpack +from ...utils import TensorType, auto_docstring, logging +from ...video_processing_utils import BaseVideoProcessor +from ...video_utils import VideoMetadata +from .processing_molmo2 import Molmo2VideosKwargs + + +logger = logging.get_logger(__name__) + + +def batch_pixels_to_patches(images: torch.Tensor, patch_size: int) -> torch.Tensor: + """Reshape images of [num_images, h, w, channels] -> [num_images, n_patches, pixels_per_patch]""" + num_crops, height, width, channels = images.shape + num_patches_height = height // patch_size + num_patches_width = width // patch_size + images = images.reshape(num_crops, num_patches_height, patch_size, num_patches_width, patch_size, channels) + images = images.transpose(2, 3) + images = images.reshape(num_crops, num_patches_height * num_patches_width, patch_size * patch_size * channels) + return images + + +def arange_for_pooling( + index_grid: torch.Tensor, + pool_h: int, + pool_w: int, +) -> tuple[torch.Tensor, int, int]: + """Group a [h, w] patch-index grid into pooling windows: `[num_rows * num_cols, pool_h * pool_w]`, padded + with -1. Also returns the pooled grid's `(num_rows, num_cols)`.""" + height_padding = pool_h * ((index_grid.shape[0] + pool_h - 1) // pool_h) - index_grid.shape[0] + width_padding = pool_w * ((index_grid.shape[1] + pool_w - 1) // pool_w) - index_grid.shape[1] + index_grid = F.pad( + index_grid, + (width_padding // 2, (width_padding + 1) // 2, height_padding // 2, (height_padding + 1) // 2), + mode="constant", + value=-1, + ) + num_rows, num_cols = index_grid.shape[0] // pool_h, index_grid.shape[1] // pool_w + pooling_indices = ( + index_grid.reshape(num_rows, pool_h, num_cols, pool_w) + .permute(0, 2, 1, 3) + .reshape(num_rows * num_cols, pool_h * pool_w) + ) + return pooling_indices, num_rows, num_cols + + +def resize_and_normalize_image( + backend: "TorchvisionBackend", + image: torch.Tensor, + output_size: list[int], + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], +) -> torch.Tensor: + # Under `torch.compile`, torchvision's uint8 bilinear resize can produce values slightly outside 0-255, + # which overflow and wrap around (-1 becomes 255, so a near-black pixel becomes white). Resizing in + # float, then rounding and clamping back, avoids that and gives the exact same bytes in eager mode. + is_integer_input = not image.dtype.is_floating_point + resized = backend.resize( + image.float() if is_integer_input else image, + size=SizeDict(height=output_size[0], width=output_size[1]), + resample=resample, + antialias=False, + ) + if is_integer_input: + resized = resized.round().clamp(0, 255).to(image.dtype) + return backend.rescale_and_normalize(resized, do_rescale, rescale_factor, do_normalize, image_mean, image_std) + + +def build_resized_image( + backend: "TorchvisionBackend", + images: torch.Tensor, + crop_size: int, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + image_patch_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # `images`: a batch of same-shape images `[N, C, H, W]`; resize the whole batch at once. + resized = resize_and_normalize_image( + backend, + images, + [crop_size, crop_size], + resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + ) + # [N, C, S, S] -> [N, 1, S, S, C]: one global (low-res) view per image. + resized = resized.permute(0, 2, 3, 1).unsqueeze(1) + # The per-patch index grid depends only on the (shared) shape, so it is built once. + patches_per_side = crop_size // image_patch_size + resized_index_grid = torch.arange(patches_per_side**2, dtype=torch.int32, device=images.device).reshape( + patches_per_side, patches_per_side + ) + return resized, resized_index_grid + + +@auto_docstring +class Molmo2VideoProcessor(BaseVideoProcessor): + resample = PILImageResampling.BILINEAR + size = {"height": 378, "width": 378} + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + patch_size = 14 + pooling_size = [3, 3] + num_frames = 64 + do_sample_frames = True + max_fps = 2 + valid_kwargs = Molmo2VideosKwargs + model_input_names = ["pixel_values_videos", "video_token_pooling", "video_grids"] + + def __init__(self, **kwargs: Unpack[Molmo2VideosKwargs]): + super().__init__(**kwargs) + if self.size is not None and (self.size.get("height", None) is None or self.size.get("width", None) is None): + raise ValueError("size must contain 'height' and 'width' keys.") + + def sample_frames( + self, + metadata: VideoMetadata, + num_frames: int | None = None, + fps: int | float | None = None, + max_fps: int | float | None = None, + **kwargs, + ): + if fps is not None and num_frames is not None: + raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!") + + max_fps = max_fps if max_fps is not None else self.max_fps + + if metadata.fps is None: + metadata.fps = fps if fps is not None else max_fps + logger.warning_once( + "Molmo2 inserts frame timestamps into video prompts, but the input video's `fps` was not provided " + f"or could not be inferred. Defaulting to `fps={metadata.fps}`. Please provide `video_metadata` " + "for more accurate timestamps." + ) + if metadata.duration is None: + metadata.duration = metadata.total_num_frames / metadata.fps + + if fps is not None: + target_num_frames = int(metadata.duration * fps) + else: + target_num_frames = num_frames if num_frames is not None else self.num_frames + if max_fps is not None and metadata.fps > max_fps: + target_num_frames = min(target_num_frames, int(metadata.duration * max_fps)) + + target_num_frames = max(min(target_num_frames, metadata.total_num_frames), 1) + total = metadata.total_num_frames + return torch.arange(0, total, total / target_num_frames).int() + + def _build_video_patches( + self, + video: torch.Tensor, + crop_size: int, + resample: PILImageResampling, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: list[float], + image_std: list[float], + image_patch_size: int, + image_pooling_h: int, + image_pooling_w: int, + ) -> tuple[list[int], torch.Tensor, torch.Tensor]: + # `build_resized_image` takes a whole `[N, C, H, W]` batch, so all frames of a video are resized in one call. + resized_frames, resized_index_grid = build_resized_image( + self, + video, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=image_patch_size, + ) + # The pooling index grid is shape-dependent only, hence shared by every frame. + pooling_indices, num_patch_rows, num_patch_cols = arange_for_pooling( + resized_index_grid, image_pooling_h, image_pooling_w + ) + # [T, 1, S, S, C] -> [T, n_patch, pixels_per_patch] + patches = batch_pixels_to_patches(resized_frames.squeeze(1), image_patch_size) + return [num_patch_rows, num_patch_cols], patches, pooling_indices + + def _preprocess( + self, + videos: list["torch.Tensor"], + size: SizeDict, + resample: PILImageResampling | None, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + patch_size: int, + pooling_size: list[int], + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchFeature: + if size.height != size.width: + raise ValueError(f"Molmo2 only supports a square `size`, got height={size.height}, width={size.width}.") + crop_size = size.height + image_pooling_h, image_pooling_w = pooling_size + + all_crops: list[torch.Tensor] = [] + all_pooled: list[torch.Tensor] = [] + all_grids: list[torch.Tensor] = [] + patch_offset = 0 + + for video in videos: + image_grid, patches, pooling_indices = self._build_video_patches( + video, + crop_size=crop_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + image_patch_size=patch_size, + image_pooling_h=image_pooling_h, + image_pooling_w=image_pooling_w, + ) + num_frames, patches_per_frame = patches.shape[:2] + frame_offsets = ( + patch_offset + torch.arange(num_frames, device=pooling_indices.device) * patches_per_frame + ).view(-1, 1, 1) + pooled_indices = torch.where(pooling_indices >= 0, pooling_indices + frame_offsets, pooling_indices) + all_pooled.append(pooled_indices.reshape(-1, pooling_indices.shape[-1])) + all_crops.append(patches) + patch_offset += num_frames * patches_per_frame + + all_grids.append(torch.tensor([num_frames, image_grid[0], image_grid[1]], dtype=torch.int64)) + + data = { + "pixel_values_videos": torch.cat(all_crops, dim=0), + "video_token_pooling": torch.cat(all_pooled, dim=0), + "video_grids": torch.stack(all_grids, dim=0), + } + return BatchFeature(data, tensor_type=return_tensors) + + +__all__ = ["Molmo2VideoProcessor"] diff --git a/src/transformers/utils/auto_docstring.py b/src/transformers/utils/auto_docstring.py index ae403d382a21..597d17a08d2e 100644 --- a/src/transformers/utils/auto_docstring.py +++ b/src/transformers/utils/auto_docstring.py @@ -2065,6 +2065,46 @@ class ModelArgs: "shape": "of shape `(batch_size, sequence_length)`", } + image_token_pooling = { + "description": """ + Indices into the flattened image patch sequence gathered by each pooled image token, with `-1` marking padding + slots. Can be obtained using [`AutoProcessor`]. See [`ProcessorMixin.__call__`] for details. + """, + "shape": "of shape `(num_image_tokens, pool_h * pool_w)`", + } + + image_grids = { + "description": """ + Per-image `[low_res_h, low_res_w, high_res_h, high_res_w]` pooled token grid. Can be obtained using + [`AutoProcessor`]. See [`ProcessorMixin.__call__`] for details. + """, + "shape": "of shape `(num_images, 4)`", + } + + image_num_crops = { + "description": """ + Number of crops per image, low-resolution view included. Can be obtained using [`AutoProcessor`]. See + [`ProcessorMixin.__call__`] for details. + """, + "shape": "of shape `(num_images,)`", + } + + video_token_pooling = { + "description": """ + Indices into the flattened video frame-patch sequence gathered by each pooled video token, with `-1` marking + padding slots. Can be obtained using [`AutoProcessor`]. See [`ProcessorMixin.__call__`] for details. + """, + "shape": "of shape `(num_video_tokens, pool_h * pool_w)`", + } + + video_grids = { + "description": """ + Per-video `[num_frames, pooled_h, pooled_w]` token grid. Can be obtained using [`AutoProcessor`]. See + [`ProcessorMixin.__call__`] for details. + """, + "shape": "of shape `(num_videos, 3)`", + } + position_ids = { "description": """ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.n_positions - 1]`. diff --git a/tests/models/molmo2/__init__.py b/tests/models/molmo2/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/models/molmo2/test_image_processing_molmo2.py b/tests/models/molmo2/test_image_processing_molmo2.py new file mode 100644 index 000000000000..0138a7a2f5d5 --- /dev/null +++ b/tests/models/molmo2/test_image_processing_molmo2.py @@ -0,0 +1,135 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np + +from transformers.testing_utils import require_torch, require_torchvision, require_vision +from transformers.utils import is_torch_available, is_torchvision_available, is_vision_available + +from ...test_image_processing_common import ImageProcessingTester, ImageProcessingTestMixin + + +if is_torch_available(): + import torch + +if is_vision_available() and is_torchvision_available(): + from PIL import Image + + from transformers import Molmo2ImageProcessor + + +class Molmo2ImageProcessingTester(ImageProcessingTester): + def __init__(self, **kwargs): + kwargs.setdefault("patch_size", 14) + super().__init__(**kwargs) + + +@require_torch +@require_vision +@require_torchvision +class Molmo2ImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase): + image_processing_class = Molmo2ImageProcessor if (is_vision_available() and is_torchvision_available()) else None + image_processor_tester_class = Molmo2ImageProcessingTester + + def test_image_processor_properties(self): + image_processor = self.image_processing_class(**self.image_processor_dict) + self.assertTrue(hasattr(image_processor, "do_resize")) + self.assertTrue(hasattr(image_processor, "size")) + self.assertTrue(hasattr(image_processor, "do_normalize")) + self.assertTrue(hasattr(image_processor, "image_mean")) + self.assertTrue(hasattr(image_processor, "image_std")) + self.assertTrue(hasattr(image_processor, "do_convert_rgb")) + self.assertTrue(hasattr(image_processor, "max_crops")) + self.assertTrue(hasattr(image_processor, "overlap_margins")) + self.assertTrue(hasattr(image_processor, "patch_size")) + self.assertTrue(hasattr(image_processor, "pooling_size")) + + def test_image_processor_from_dict_with_kwargs(self): + image_processor = self.image_processing_class.from_dict(self.image_processor_dict) + self.assertEqual(image_processor.size, {"height": 378, "width": 378}) + self.assertEqual(image_processor.do_normalize, True) + + image_processor = self.image_processing_class.from_dict( + self.image_processor_dict, size={"height": 400, "width": 400}, do_normalize=False + ) + self.assertEqual(image_processor.size, {"height": 400, "width": 400}) + self.assertEqual(image_processor.do_normalize, False) + + def _assert_patchified_output(self, outputs, expected_num_images): + pixel_values = outputs.pixel_values + self.assertEqual(pixel_values.ndim, 3) + pixels_per_patch = self.image_processor_tester.patch_size**2 * self.image_processor_tester.num_channels + self.assertEqual(pixel_values.shape[-1], pixels_per_patch) + image_num_crops = outputs.image_num_crops + self.assertEqual(image_num_crops.shape[0], expected_num_images) + self.assertEqual(pixel_values.shape[0], int(image_num_crops.sum().item())) + + def test_call_pil(self): + for image_processing_class in [self.image_processing_class]: + image_processing = image_processing_class(**self.image_processor_dict) + image_inputs = self.image_processor_tester.prepare_image_inputs(equal_resolution=False) + for image in image_inputs: + self.assertIsInstance(image, Image.Image) + + outputs = image_processing(image_inputs[0], return_tensors="pt") + self._assert_patchified_output(outputs, 1) + + outputs = image_processing([[image] for image in image_inputs], return_tensors="pt") + self._assert_patchified_output(outputs, self.image_processor_tester.batch_size) + + def test_call_numpy(self): + for image_processing_class in [self.image_processing_class]: + image_processing = image_processing_class(**self.image_processor_dict) + image_inputs = self.image_processor_tester.prepare_image_inputs(equal_resolution=False, numpify=True) + for image in image_inputs: + self.assertIsInstance(image, np.ndarray) + + outputs = image_processing(image_inputs[0], return_tensors="pt") + self._assert_patchified_output(outputs, 1) + + outputs = image_processing([[image] for image in image_inputs], return_tensors="pt") + self._assert_patchified_output(outputs, self.image_processor_tester.batch_size) + + def test_call_nested_images(self): + for image_processing_class in [self.image_processing_class]: + image_processing = image_processing_class(**self.image_processor_dict) + image_inputs = self.image_processor_tester.prepare_image_inputs(equal_resolution=False) + nested_images = [[], [image_inputs[0]], [image_inputs[1], image_inputs[2]]] + + outputs = image_processing(nested_images, return_tensors="pt") + + self._assert_patchified_output(outputs, 3) + for crops in outputs.image_num_crops.tolist(): + self.assertGreater(crops, 0) + + def test_call_pytorch(self): + for image_processing_class in [self.image_processing_class]: + image_processing = image_processing_class(**self.image_processor_dict) + image_inputs = self.image_processor_tester.prepare_image_inputs(equal_resolution=False, torchify=True) + for image in image_inputs: + self.assertIsInstance(image, torch.Tensor) + + outputs = image_processing(image_inputs[0], return_tensors="pt") + self._assert_patchified_output(outputs, 1) + + outputs = image_processing([[image] for image in image_inputs], return_tensors="pt") + self._assert_patchified_output(outputs, self.image_processor_tester.batch_size) + + @unittest.skip( + reason="Molmo2ImageProcessor always converts to RGB before processing; 4-channel images are not supported." + ) + def test_call_numpy_4_channels(self): + pass diff --git a/tests/models/molmo2/test_modeling_molmo2.py b/tests/models/molmo2/test_modeling_molmo2.py new file mode 100644 index 000000000000..0b3d0c3e8561 --- /dev/null +++ b/tests/models/molmo2/test_modeling_molmo2.py @@ -0,0 +1,713 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Testing suite for the PyTorch Molmo2 model.""" + +import copy +import unittest +from unittest.mock import patch + +from parameterized import parameterized + +from transformers import ( + Molmo2Config, + Molmo2ForConditionalGeneration, + Molmo2Model, + Molmo2Processor, + is_torch_available, + is_vision_available, +) +from transformers.models.molmo2.configuration_molmo2 import ( + Molmo2AdapterConfig, + Molmo2TextConfig, + Molmo2VisionConfig, +) +from transformers.testing_utils import ( + Expectations, + require_torch, + require_torch_large_accelerator, + require_vision, + slow, + torch_device, +) +from transformers.video_utils import load_video + +from ...test_memory_cleanup_mixin import MemoryCleanupMixin +from ...test_modeling_common import floats_tensor +from ...test_processing_common import url_to_local_path +from ...vlm_tester import VLMModelTest, VLMModelTester + + +if is_torch_available(): + import torch + + +if is_vision_available(): + from transformers.image_utils import load_image + + +class Molmo2VisionText2TextModelTester(VLMModelTester): + base_model_class = Molmo2Model + config_class = Molmo2Config + text_config_class = Molmo2TextConfig + vision_config_class = Molmo2VisionConfig + conditional_generation_class = Molmo2ForConditionalGeneration + + def __init__(self, parent, **kwargs): + kwargs.setdefault("image_size", 378) + kwargs.setdefault("patch_size", 14) + kwargs.setdefault("num_image_tokens", 32) + kwargs.setdefault("seq_length", 7 + kwargs["num_image_tokens"]) + kwargs.setdefault("hidden_size", 32) + kwargs.setdefault("intermediate_size", 37) + kwargs.setdefault("num_attention_heads", 4) + kwargs.setdefault("num_key_value_heads", 2) + kwargs.setdefault("head_dim", 128) + kwargs.setdefault("num_hidden_layers", 2) + kwargs.setdefault("hidden_act", "silu") + kwargs.setdefault("max_position_embeddings", 512) + kwargs.setdefault("bos_token_id", 0) + kwargs.setdefault("eos_token_id", 1) + kwargs.setdefault("pad_token_id", 2) + kwargs.setdefault("image_start_token_id", 3) + kwargs.setdefault("image_end_token_id", 4) + kwargs.setdefault("image_patch_id", 5) + # Alias so base helpers (special-token clearing, mismatch tests) protect image patch tokens. + kwargs.setdefault("image_token_id", kwargs["image_patch_id"]) + super().__init__(parent, **kwargs) + + def create_pixel_values(self, batch_size: int | None = None): + batch_size = batch_size if batch_size is not None else self.batch_size + num_patches = (self.image_size // self.patch_size) ** 2 + return floats_tensor( + [ + batch_size, + num_patches, + self.patch_size * self.patch_size * self.num_channels, + ] + ) + + def place_image_tokens(self, input_ids, config): + input_ids = input_ids.clone() + input_ids[:, -1] = self.pad_token_id + input_ids[input_ids == self.image_patch_id] = self.pad_token_id + input_ids[:, : self.num_image_tokens] = self.image_patch_id + return input_ids + + def get_additional_inputs(self, config, input_ids, pixel_values, batch_size: int | None = None): + batch_size = batch_size if batch_size is not None else input_ids.shape[0] + num_patches = (self.image_size // self.patch_size) ** 2 + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[input_ids == self.image_patch_id] = 1 + pooled = torch.randint( + 0, + num_patches, + (batch_size, self.num_image_tokens, 4), + device=torch_device, + ) + sample_offsets = (torch.arange(batch_size, device=torch_device) * num_patches).view(-1, 1, 1) + image_token_pooling = (pooled + sample_offsets).view(-1, 4) + return { + "image_token_pooling": image_token_pooling, + "image_grids": torch.tensor([[4, 4, 4, 4]] * batch_size, device=torch_device), + "image_num_crops": torch.ones(batch_size, dtype=torch.long, device=torch_device), + "mm_token_type_ids": mm_token_type_ids, + } + + def get_config(self): + text_config = Molmo2TextConfig( + bos_token_id=self.bos_token_id, + eos_token_id=self.eos_token_id, + pad_token_id=self.pad_token_id, + hidden_act=self.hidden_act, + head_dim=self.head_dim, + hidden_size=self.hidden_size, + vocab_size=self.vocab_size, + intermediate_size=self.intermediate_size, + max_position_embeddings=self.max_position_embeddings, + num_attention_heads=self.num_attention_heads, + num_hidden_layers=self.num_hidden_layers, + num_key_value_heads=self.num_key_value_heads, + rope_theta=10000.0, + tie_word_embeddings=self.tie_word_embeddings, + rms_norm_eps=1e-6, + ) + vision_config = Molmo2VisionConfig( + hidden_size=32, + intermediate_size=37, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=8, + hidden_act="gelu_pytorch_tanh", + layer_norm_eps=1e-6, + image_size=[self.image_size, self.image_size], + patch_size=self.patch_size, + attention_dropout=0.0, + ) + adapter_config = Molmo2AdapterConfig( + vision_feature_layer=[1], + hidden_size=32, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=8, + intermediate_size=37, + text_hidden_size=32, + hidden_act="silu", + ) + return Molmo2Config( + text_config=text_config, + vision_config=vision_config, + adapter_config=adapter_config, + image_start_token_id=self.image_start_token_id, + image_end_token_id=self.image_end_token_id, + image_patch_id=self.image_patch_id, + tie_word_embeddings=self.tie_word_embeddings, + ) + + +@require_torch +class Molmo2ModelTest(VLMModelTest, unittest.TestCase): + """ + Model tester for `Molmo2ForConditionalGeneration`. + """ + + model_tester_class = Molmo2VisionText2TextModelTester + additional_model_inputs = ["mm_token_type_ids"] + + @unittest.skip( + "The test slices every input tensor in half along dim 0, but `image_token_pooling` has no batch " + "dimension: pooled patches of all images are concatenated along one axis." + ) + def test_generate_compile_model_forward_fullgraph(self): + pass + + def _video_features_prepare_config_and_inputs(self): + # The generic helper only renames `pixel_values`; Molmo2's `get_video_features` also needs the + # pooling index tensor under its video name and a `[num_frames, rows, cols]` grid per video. + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + batch_size = inputs_dict["image_grids"].shape[0] + inputs_dict = { + "pixel_values_videos": inputs_dict["pixel_values"], + "video_token_pooling": inputs_dict["image_token_pooling"], + "video_grids": torch.tensor([[2, 4, 4]] * batch_size, device=torch_device), + } + return config, inputs_dict + + @unittest.skip( + reason="Molmo2 always builds an attention mask in the forward (block-sequence mask for images, like " + "gemma3/mllama), and SDPA cannot dispatch masked attention to the flash kernel" + ) + def test_sdpa_can_dispatch_on_flash(self): + pass + + def flash_attn_from_config(self, attn_implementation: str, test_fwd_in_train: bool = True): + super().flash_attn_from_config(attn_implementation, test_fwd_in_train=False) + + def flash_attn_inference_equivalence( + self, attn_implementation: str, padding_side: str, atol: float = 4e-2, rtol: float = 4e-2 + ): + # The common test keeps one sample with `[:1]`, which cannot slice the flat `image_token_pooling`, + # so the check runs on text-only inputs. + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + input_ids = inputs_dict["input_ids"].clone() + input_ids[input_ids == config.image_patch_id] = self.model_tester.pad_token_id + text_inputs = {"input_ids": input_ids, "attention_mask": inputs_dict["attention_mask"]} + with patch.object( + self.model_tester, + "prepare_config_and_inputs_for_common", + side_effect=lambda: (copy.deepcopy(config), copy.deepcopy(text_inputs)), + ): + super().flash_attn_inference_equivalence(attn_implementation, padding_side, atol=atol, rtol=rtol) + + @unittest.skip( + reason="Multimodal special tokens live in the extra-vocab rows beyond `vocab_size`; standard resize is ill-defined" + ) + def test_resize_tokens_embeddings(self): + pass + + @unittest.skip( + reason="Multimodal special tokens live in the extra-vocab rows beyond `vocab_size`; standard resize is ill-defined" + ) + def test_resize_embeddings_untied(self): + pass + + @parameterized.expand([("greedy", 1), ("beam_search", 2)]) + def test_generate_from_inputs_embeds_textonly(self, _, num_beams): + """Pure-LLM path: drop the image inputs so generation runs from plain text `inputs_embeds`.""" + for model_class in self.all_generative_model_classes: + config, inputs_dict = self.prepare_config_and_inputs_for_generate() + config.is_decoder = True + model = model_class(config).to(torch_device).eval() + input_ids = inputs_dict.pop("input_ids") + for key in ( + "pixel_values", + "image_token_pooling", + "image_grids", + "image_num_crops", + "pixel_values_videos", + "video_token_pooling", + "video_grids", + ): + inputs_dict.pop(key, None) + gen_kwargs = { + "return_dict_in_generate": True, + "output_scores": True, + "num_beams": num_beams, + "do_sample": False, + "max_new_tokens": 5, + "min_new_tokens": 5, + "use_cache": True, + } + inputs_embeds = model.get_input_embeddings()(input_ids) + out = model.generate(inputs_embeds=inputs_embeds, **gen_kwargs, **inputs_dict) + # inputs_embeds-only generation returns only the newly generated tokens + self.assertEqual(out.sequences.shape[0], input_ids.shape[0]) + self.assertEqual(out.sequences.shape[1], 5) + + def test_mismatching_num_image_tokens(self): + """ + Tests that VLMs handle single-batch image inputs correctly. + """ + config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() + for model_class in self.all_model_classes: + model = model_class(config).to(torch_device) + model.eval() + _ = model(**input_dict) # successful forward with no modifications + curr_input_dict = copy.deepcopy(input_dict) + + num_image_tokens = self.model_tester.num_image_tokens + curr_input_dict["input_ids"] = curr_input_dict["input_ids"][:1, ...] + curr_input_dict["attention_mask"] = curr_input_dict["attention_mask"][:1, ...] + curr_input_dict["pixel_values"] = curr_input_dict["pixel_values"][:1, ...] + curr_input_dict["image_token_pooling"] = curr_input_dict["image_token_pooling"][:num_image_tokens] + curr_input_dict["image_grids"] = curr_input_dict["image_grids"][:1, ...] + curr_input_dict["image_num_crops"] = curr_input_dict["image_num_crops"][:1, ...] + _ = model(**curr_input_dict) + + # Image features get cached in KV cache like other VLMs; no need to skip. + + def test_retain_grad_hidden_states_attentions(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.output_hidden_states = True + config.output_attentions = self.has_attentions + self._set_subconfig_attributes(config, "output_hidden_states", True) + self._set_subconfig_attributes(config, "output_attentions", self.has_attentions) + + for model_class in self.all_model_classes: + model = model_class._from_config(config, attn_implementation="eager").to(torch_device) + outputs = model(**inputs_dict) + + output = outputs[0] + hidden_states = outputs.hidden_states[0] + hidden_states.retain_grad() + + if self.has_attentions: + attentions = outputs.attentions[0] + attentions.retain_grad() + + output.flatten()[0].backward(retain_graph=True) + + self.assertIsNotNone(hidden_states.grad) + if self.has_attentions: + self.assertIsNotNone(attentions.grad) + + +IMAGE_URL = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg" + + +@slow +@require_torch +@require_vision +class Molmo2_4BIntegrationTest(MemoryCleanupMixin, unittest.TestCase): + model_id = "allenai/Molmo2-4B" + + def setUp(self): + super().setUp() + self.processor = Molmo2Processor.from_pretrained(self.model_id) + self.image = load_image(url_to_local_path(IMAGE_URL)) + self.messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image."}, + {"type": "image", "image": self.image}, + ], + } + ] + + def build_inputs(self): + return self.processor.apply_chat_template( + self.messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ) + + def test_preprocessing(self): + inputs = self.build_inputs() + self.assertEqual(inputs["pixel_values"].shape, torch.Size([7, 729, 588])) + self.assertEqual(inputs["image_token_pooling"].shape, torch.Size([955, 4])) + self.assertEqual(inputs["image_grids"].shape, torch.Size([1, 4])) + self.assertEqual(inputs["image_num_crops"].tolist(), [7]) + self.assertEqual(inputs["mm_token_type_ids"].sum().item(), 982) + self.assertEqual(inputs["input_ids"].shape[0], 1) + # 4B uses the Qwen tokenizer; `<|im_end|>` (151645) is the leading BOS. + self.assertEqual(inputs["input_ids"][0, 0].item(), 151645) + + expected_pixel_slices = Expectations( + { + (None, None): torch.tensor( + [ + [-0.07450979948043823, -0.05098038911819458, 0.019607901573181152], + [-0.7019608020782471, -0.6784313917160034, -0.6078431606292725], + [-0.8745098114013672, -0.8823529481887817, -0.843137264251709], + ], + dtype=torch.float32, + ), + } + ) + torch.testing.assert_close( + inputs["pixel_values"][0, :3, :3].float().cpu(), + expected_pixel_slices.get_expectation(), + atol=1e-2, + rtol=1e-4, + ) + + def test_forward_logits(self): + inputs = self.build_inputs() + + model = Molmo2ForConditionalGeneration.from_pretrained( + self.model_id, + dtype=torch.float32, + device_map="auto", + ) + + device_inputs = inputs.to(torch_device) + + with torch.no_grad(): + outputs = model(**device_inputs) + + logits = outputs.logits + self.assertEqual(logits.shape[0], 1) + self.assertEqual(logits.shape[1], device_inputs["input_ids"].shape[1]) + + expected_last_logits = Expectations( + { + ("cuda", (8, 0)): [-10.407500, -5.903657, -10.977587, -10.325406, -16.847645, -14.505170, -11.184648, -9.696571, -11.637183, -9.205433], + ("cuda", (8, 6)): [-10.407710, -5.903650, -10.977592, -10.325156, -16.847317, -14.505305, -11.184698, -9.697040, -11.636929, -9.205041], + } + ) # fmt: skip + torch.testing.assert_close( + logits[0, -1, :10].cpu().float(), + torch.tensor(expected_last_logits.get_expectation(), dtype=torch.float32), + atol=3e-1, + rtol=5e-2, + ) + self.assertEqual(logits[0, -1].argmax().item(), 641) + + def test_generation(self): + inputs = self.build_inputs() + + model = Molmo2ForConditionalGeneration.from_pretrained( + self.model_id, + dtype=torch.float32, + device_map="auto", + ) + + device_inputs = inputs.to(torch_device) + + with torch.no_grad(): + generated_ids = model.generate(**device_inputs, max_new_tokens=10, do_sample=False) + + input_len = device_inputs["input_ids"].shape[1] + generated_text = self.processor.batch_decode(generated_ids[:, input_len:], skip_special_tokens=True)[0] + expected_texts = Expectations( + { + ("cuda", (8, 0)): "In this captivating image, a large, chubby cat", + ("cuda", (8, 6)): "In this captivating image, a large, chubby cat", + } + ) # fmt: skip + self.assertEqual(generated_text.strip(), expected_texts.get_expectation()) + + +@slow +@require_torch +@require_vision +class Molmo2_O7BIntegrationTest(MemoryCleanupMixin, unittest.TestCase): + model_id = "allenai/Molmo2-O-7B" + + def setUp(self): + super().setUp() + self.processor = Molmo2Processor.from_pretrained(self.model_id) + self.image = load_image(url_to_local_path(IMAGE_URL)) + self.messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image."}, + {"type": "image", "image": self.image}, + ], + } + ] + + def build_inputs(self): + return self.processor.apply_chat_template( + self.messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ) + + def test_preprocessing(self): + inputs = self.build_inputs() + self.assertEqual(inputs["pixel_values"].shape, torch.Size([7, 729, 588])) + self.assertEqual(inputs["image_token_pooling"].shape, torch.Size([955, 4])) + self.assertEqual(inputs["image_grids"].shape, torch.Size([1, 4])) + self.assertEqual(inputs["input_ids"].shape[0], 1) + # O-7B uses the OLMo tokenizer; `<|endoftext|>` (100257) is the leading BOS. + self.assertEqual(inputs["input_ids"][0, 0].item(), 100257) + + expected_pixel_slices = Expectations( + { + (None, None): torch.tensor( + [ + [-0.07450979948043823, -0.05098038911819458, 0.019607901573181152], + [-0.7019608020782471, -0.6784313917160034, -0.6078431606292725], + [-0.8745098114013672, -0.8823529481887817, -0.843137264251709], + ], + dtype=torch.float32, + ), + } + ) + torch.testing.assert_close( + inputs["pixel_values"][0, :3, :3].float().cpu(), + expected_pixel_slices.get_expectation(), + atol=1e-2, + rtol=1e-4, + ) + + def test_forward_logits(self): + inputs = self.build_inputs() + + model = Molmo2ForConditionalGeneration.from_pretrained( + self.model_id, + dtype=torch.bfloat16, + device_map="auto", + ) + + device_inputs = inputs.to(torch_device) + + with torch.no_grad(): + outputs = model(**device_inputs) + + logits = outputs.logits + self.assertEqual(logits.shape[0], 1) + self.assertEqual(logits.shape[1], device_inputs["input_ids"].shape[1]) + + expected_last_logits = Expectations( + { + ("cuda", (8, 0)): [-13.0625, -5.9375, -11.75, -11.0, -12.6875, -16.25, -10.375, -12.3125, -12.6875, -10.625], + ("cuda", (8, 6)): [-13.0625, -5.875, -11.6875, -11.0, -12.6875, -16.25, -10.3125, -12.25, -12.6875, -10.625], + } + ) # fmt: skip + torch.testing.assert_close( + logits[0, -1, :10].cpu().float(), + torch.tensor(expected_last_logits.get_expectation(), dtype=torch.float32), + atol=3e-1, + rtol=5e-2, + ) + self.assertEqual(logits[0, -1].argmax().item(), 644) + + def test_generation(self): + inputs = self.build_inputs() + + model = Molmo2ForConditionalGeneration.from_pretrained( + self.model_id, + dtype=torch.bfloat16, + device_map="auto", + ) + + device_inputs = inputs.to(torch_device) + + with torch.no_grad(): + generated_ids = model.generate(**device_inputs, max_new_tokens=10, do_sample=False) + + input_len = device_inputs["input_ids"].shape[1] + generated_text = self.processor.batch_decode(generated_ids[:, input_len:], skip_special_tokens=True)[0] + expected_texts = Expectations( + { + ("cuda", (8, 0)): "In this captivating image, a small, chubby cat", + ("cuda", (8, 6)): "In this captivating image, a small, chubby cat", + } + ) # fmt: skip + self.assertEqual(generated_text.strip(), expected_texts.get_expectation()) + + +@slow +@require_torch +@require_vision +class Molmo2_8BIntegrationTest(MemoryCleanupMixin, unittest.TestCase): + model_id = "allenai/Molmo2-8B" + + def setUp(self): + super().setUp() + self.processor = Molmo2Processor.from_pretrained(self.model_id) + self.image = load_image(url_to_local_path(IMAGE_URL)) + self.messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image."}, + {"type": "image", "image": self.image}, + ], + } + ] + + def build_inputs(self): + return self.processor.apply_chat_template( + self.messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ) + + def test_preprocessing(self): + inputs = self.build_inputs() + self.assertEqual(inputs["pixel_values"].shape, torch.Size([7, 729, 588])) + self.assertEqual(inputs["image_token_pooling"].shape, torch.Size([955, 4])) + self.assertEqual(inputs["image_grids"].shape, torch.Size([1, 4])) + self.assertEqual(inputs["input_ids"].shape[0], 1) + # 8B uses the Qwen tokenizer; `<|im_end|>` (151645) is the leading BOS. + self.assertEqual(inputs["input_ids"][0, 0].item(), 151645) + + expected_pixel_slices = Expectations( + { + (None, None): torch.tensor( + [ + [-0.07450979948043823, -0.05098038911819458, 0.019607901573181152], + [-0.7019608020782471, -0.6784313917160034, -0.6078431606292725], + [-0.8745098114013672, -0.8823529481887817, -0.843137264251709], + ], + dtype=torch.float32, + ), + } + ) + torch.testing.assert_close( + inputs["pixel_values"][0, :3, :3].float().cpu(), + expected_pixel_slices.get_expectation(), + atol=1e-2, + rtol=1e-4, + ) + + def test_forward_logits(self): + inputs = self.build_inputs() + + model = Molmo2ForConditionalGeneration.from_pretrained( + self.model_id, + dtype=torch.bfloat16, + device_map="auto", + ) + + device_inputs = inputs.to(torch_device) + + with torch.no_grad(): + outputs = model(**device_inputs) + + logits = outputs.logits + self.assertEqual(logits.shape[0], 1) + self.assertEqual(logits.shape[1], device_inputs["input_ids"].shape[1]) + + expected_last_logits = Expectations( + { + ("cuda", (8, 0)): [-15.875, -7.875, -15.5625, -15.0, -16.5, -18.25, -14.4375, -15.8125, -15.4375, -12.4375], + ("cuda", (8, 6)): [-15.875, -7.875, -15.625, -15.0, -16.5, -18.25, -14.5, -15.75, -15.4375, -12.5], + } + ) # fmt: skip + torch.testing.assert_close( + logits[0, -1, :10].cpu().float(), + torch.tensor(expected_last_logits.get_expectation(), dtype=torch.float32), + atol=3e-1, + rtol=5e-2, + ) + self.assertEqual(logits[0, -1].argmax().item(), 641) + + def test_generation(self): + inputs = self.build_inputs() + + model = Molmo2ForConditionalGeneration.from_pretrained( + self.model_id, + dtype=torch.bfloat16, + device_map="auto", + ) + + device_inputs = inputs.to(torch_device) + + with torch.no_grad(): + generated_ids = model.generate(**device_inputs, max_new_tokens=10, do_sample=False) + + input_len = device_inputs["input_ids"].shape[1] + generated_text = self.processor.batch_decode(generated_ids[:, input_len:], skip_special_tokens=True)[0] + expected_texts = Expectations( + { + ("cuda", (8, 0)): "In this captivating image, a snow leopard is captured", + ("cuda", (8, 6)): "In this captivating image, a snow leopard is captured", + } + ) # fmt: skip + self.assertEqual(generated_text.strip(), expected_texts.get_expectation()) + + @require_torch_large_accelerator(memory=30) + def test_generation_video_qa(self): + """Test video question answering for Molmo2-8B.""" + video_url = "https://storage.googleapis.com/oe-training-public/demo_videos/many_penguins.mp4" + video, metadata = load_video(url_to_local_path(video_url)) + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Which animal appears in the video?"}, + {"type": "video", "video": video}, + ], + } + ] + inputs = self.processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + processor_kwargs={"video_metadata": [metadata]}, + ) + + model = Molmo2ForConditionalGeneration.from_pretrained( + self.model_id, + dtype=torch.bfloat16, + device_map="auto", + ) + + device_inputs = inputs.to(torch_device) + + with torch.no_grad(): + generated_ids = model.generate(**device_inputs, max_new_tokens=64, do_sample=False) + + input_len = device_inputs["input_ids"].shape[1] + generated_text = self.processor.batch_decode(generated_ids[:, input_len:], skip_special_tokens=True)[0] + expected_texts = Expectations( + { + ("cuda", (8, 0)): "Penguins appear in the video.", + } + ) # fmt: skip + self.assertEqual(generated_text.strip(), expected_texts.get_expectation()) diff --git a/tests/models/molmo2/test_processing_molmo2.py b/tests/models/molmo2/test_processing_molmo2.py new file mode 100644 index 000000000000..86a134e2ea7f --- /dev/null +++ b/tests/models/molmo2/test_processing_molmo2.py @@ -0,0 +1,279 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import torch +from PIL import Image + +from transformers import Molmo2ImageProcessor, Molmo2Processor, Molmo2VideoProcessor +from transformers.testing_utils import require_torch, require_torchvision, require_vision + +from ...test_processing_common import ProcessorTesterMixin + + +@require_vision +@require_torch +@require_torchvision +class Molmo2ProcessorTest(ProcessorTesterMixin, unittest.TestCase): + processor_class = Molmo2Processor + model_id = "allenai/Molmo2-8B" + + @unittest.skip( + "Processor prepends the BOS token as text, which shifts the offsets the assistant mask is computed from" + ) + def test_apply_chat_template_assistant_mask(self): + pass + + @classmethod + def _setup_image_processor(cls): + return Molmo2ImageProcessor(size={"height": 56, "width": 56}, max_crops=1, overlap_margins=[1, 1]) + + @classmethod + def _setup_video_processor(cls): + return Molmo2VideoProcessor(size={"height": 56, "width": 56}, num_frames=4) + + @staticmethod + def prepare_processor_dict(): + # Override the chat template to support the "system" role used by the base test harness. + # The original Molmo2 template enforces strict user/assistant alternation without system. + return { + "chat_template": ( + "{{ bos_token }}" + "{%- if messages[0]['role'] == 'system' -%}" + " {%- set system_message = messages[0]['content'][0]['text'] -%}" + " {%- set loop_messages = messages[1:] -%}" + "{%- else -%}" + " {%- set system_message = '' -%}" + " {%- set loop_messages = messages -%}" + "{%- endif -%}" + "{%- for message in loop_messages -%}" + " {%- if message['role'] == 'user' -%}" + " {{ '<|im_start|>user\\n' }}" + " {%- if message['content'] is string -%}" + " {{ message['content'] }}" + " {%- else -%}" + " {%- for item in message['content'] -%}" + " {%- if item['type'] == 'image' -%}" + " {{ '<|image|>' }}" + " {%- elif item['type'] == 'video' -%}" + " {{ '<|video|>' }}" + " {%- elif item['type'] == 'text' -%}" + " {{ item['text'] }}" + " {%- endif -%}" + " {%- endfor -%}" + " {%- endif -%}" + " {{ '<|im_end|>\\n' }}" + " {%- elif message['role'] == 'assistant' -%}" + " {{ '<|im_start|>assistant\\n' }}" + " {%- if message['content'] is string -%}" + " {{ message['content'] }}" + " {%- else -%}" + " {%- for item in message['content'] -%}" + " {%- if item['type'] == 'text' -%}" + " {{ item['text'] }}" + " {%- endif -%}" + " {%- endfor -%}" + " {%- endif -%}" + " {%- endif -%}" + "{%- endfor -%}" + "{%- if add_generation_prompt -%}" + " {{ '<|im_start|>assistant\\n' }}" + "{%- endif -%}" + ), + } + + @require_torch + def _test_apply_chat_template( + self, + modality: str, + batch_size: int, + return_tensors: str, + input_name: str, + processor_name: str, + input_data: list, + ): + processor = self.get_processor() + if processor.chat_template is None: + self.skipTest("Processor has no chat template") + + if processor_name not in self.processor_class.get_attributes(): + self.skipTest(f"{processor_name} attribute not present in {self.processor_class}") + + if getattr(processor, processor_name).__class__.__name__.endswith("Fast"): + return_tensors = "pt" + + batch_messages = [ + [ + {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]}, + {"role": "user", "content": [{"type": "text", "text": "Describe this."}]}, + ] + ] * batch_size + + formatted_prompt = processor.apply_chat_template(batch_messages, add_generation_prompt=True, tokenize=False) + self.assertEqual(len(formatted_prompt), batch_size) + + formatted_prompt_tokenized = processor.apply_chat_template( + batch_messages, add_generation_prompt=True, tokenize=True, return_tensors=return_tensors + ) + add_special_tokens = True + if processor.tokenizer.bos_token is not None and formatted_prompt[0].startswith(processor.tokenizer.bos_token): + add_special_tokens = False + tok_output = processor.tokenizer( + formatted_prompt, return_tensors=return_tensors, add_special_tokens=add_special_tokens + ) + self.assertListEqual(tok_output.input_ids.tolist(), formatted_prompt_tokenized.tolist()) + + tokenized_prompt_100 = processor.apply_chat_template( + batch_messages, + add_generation_prompt=True, + tokenize=True, + return_tensors=return_tensors, + processor_kwargs={ + "padding": "max_length", + "truncation": True, + "max_length": self.chat_template_max_length, + }, + ) + self.assertEqual(len(tokenized_prompt_100[0]), self.chat_template_max_length) + + out_dict_text = processor.apply_chat_template( + batch_messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors=return_tensors, + ) + self.assertTrue(all(key in out_dict_text for key in ["input_ids", "attention_mask"])) + self.assertEqual(len(out_dict_text["input_ids"]), batch_size) + self.assertEqual(len(out_dict_text["attention_mask"]), batch_size) + + for idx, url in enumerate(input_data[:batch_size]): + batch_messages[idx][1]["content"] = [batch_messages[idx][1]["content"][0], {"type": modality, "url": url}] + + out_dict = processor.apply_chat_template( + batch_messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors=return_tensors, + processor_kwargs={"num_frames": 2}, + ) + input_name = getattr(self, input_name) + self.assertTrue(input_name in out_dict) + self.assertEqual(len(out_dict["input_ids"]), batch_size) + self.assertEqual(len(out_dict["attention_mask"]), batch_size) + if modality == "video": + expected_len = int(out_dict["video_grids"][:, 0].sum().item()) + else: + expected_len = int(out_dict["image_num_crops"].sum().item()) + self.assertEqual(len(out_dict[input_name]), expected_len) + + return_tensor_to_type = {"pt": torch.Tensor, "np": np.ndarray, None: list} + for k in out_dict: + self.assertIsInstance(out_dict[k], return_tensor_to_type[return_tensors]) + + assistant_message = {"role": "assistant", "content": [{"type": "text", "text": "It is the sound of"}]} + for idx, url in enumerate(input_data[:batch_size]): + batch_messages[idx] = batch_messages[idx] + [assistant_message] + continue_prompt = processor.apply_chat_template(batch_messages, continue_final_message=True, tokenize=False) + for prompt in continue_prompt: + self.assertTrue(prompt.endswith("It is the sound of")) + + def test_apply_chat_template_video_frame_sampling(self): + processor = self.get_processor() + if processor.chat_template is None: + self.skipTest("Processor has no chat template") + + video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/tiny_video.mp4" + messages = [ + [ + { + "role": "user", + "content": [ + {"type": "video", "url": video_url}, + {"type": "text", "text": "What is shown in this video?"}, + ], + }, + ] + ] + + # Default `max_fps=2` caps `num_frames=3` for the ~1s tiny_video to a single frame. + out_capped = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + processor_kwargs={"num_frames": 3, "fps": None}, + ) + self.assertIn(self.videos_input_name, out_capped) + self.assertEqual(out_capped["video_grids"].shape[0], 1) + self.assertEqual(int(out_capped["video_grids"][0, 0].item()), 1) + + # Raising the cap above the video's native fps restores the requested `num_frames`. + out_uncapped = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + processor_kwargs={"num_frames": 3, "fps": None, "max_fps": 1000}, + ) + self.assertEqual(int(out_uncapped["video_grids"][0, 0].item()), 3) + + # `fps` mode bypasses the cap entirely (different code path in sample_frames). + out_fps = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + processor_kwargs={"fps": 10, "num_frames": None}, + ) + self.assertGreaterEqual(int(out_fps["video_grids"][0, 0].item()), 1) + + # `fps` and `num_frames` are mutually exclusive. + with self.assertRaises(ValueError): + processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + processor_kwargs={"fps": 10, "num_frames": 3}, + ) + + # `num_image_tokens` is covered by the inherited `test_get_num_multimodal_tokens_matches_processor_call`; + # `num_image_patches` is not, and vLLM's transformers backend splits `pixel_values` per image with it. + @unittest.skip("processor prepends BOS to the text until the hub chat templates emit it") + def test_subprocessor_defaults_0_text(self): + pass + + def test_get_num_multimodal_tokens_matches_num_patches(self): + processor = self.get_processor() + + for height, width in ((354, 536), (768, 1024), (100, 100)): + image = Image.fromarray(np.random.randint(255, size=(height, width, 3), dtype=np.uint8)) + inputs = processor(text=[processor.image_token], images=[image], return_tensors="pt") + predicted = processor._get_num_multimodal_tokens(image_sizes=[(height, width)]) + + self.assertEqual(predicted["num_image_patches"][0], inputs["pixel_values"].shape[0]) + + @unittest.skip( + reason="Hub processor_config.json has use_single_crop_col_tokens=False, which differs from the __init__ " + "default of None when building from components." + ) + def test_processor_from_pretrained_vs_from_components(self): + pass diff --git a/tests/models/molmo2/test_video_processing_molmo2.py b/tests/models/molmo2/test_video_processing_molmo2.py new file mode 100644 index 000000000000..8c67778b12a5 --- /dev/null +++ b/tests/models/molmo2/test_video_processing_molmo2.py @@ -0,0 +1,225 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np + +from transformers.image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD +from transformers.testing_utils import require_torch, require_torchvision, require_vision +from transformers.utils import is_torchvision_available, is_vision_available + +from ...test_video_processing_common import VideoProcessingTestMixin, prepare_video_inputs + + +if is_vision_available() and is_torchvision_available(): + from transformers import Molmo2VideoProcessor + + +class Molmo2VideoProcessingTester: + def __init__( + self, + parent, + batch_size=5, + num_frames=8, + num_channels=3, + min_resolution=32, + max_resolution=80, + do_resize=True, + size=None, + do_normalize=True, + image_mean=IMAGENET_STANDARD_MEAN, + image_std=IMAGENET_STANDARD_STD, + do_convert_rgb=True, + patch_size=14, + pooling_size=[3, 3], + do_sample_frames=True, + max_fps=2, + ): + size = size if size is not None else {"height": 378, "width": 378} + self.parent = parent + self.batch_size = batch_size + self.num_frames = num_frames + self.num_channels = num_channels + self.min_resolution = min_resolution + self.max_resolution = max_resolution + self.do_resize = do_resize + self.size = size + self.do_normalize = do_normalize + self.image_mean = image_mean + self.image_std = image_std + self.do_convert_rgb = do_convert_rgb + self.patch_size = patch_size + self.pooling_size = pooling_size + self.do_sample_frames = do_sample_frames + self.max_fps = max_fps + + def prepare_video_processor_dict(self): + return { + "do_resize": self.do_resize, + "size": self.size, + "do_normalize": self.do_normalize, + "image_mean": self.image_mean, + "image_std": self.image_std, + "do_convert_rgb": self.do_convert_rgb, + "patch_size": self.patch_size, + "pooling_size": self.pooling_size, + "do_sample_frames": False, + "max_fps": self.max_fps, + } + + def prepare_video_inputs(self, equal_resolution=False, numpify=False, torchify=False, return_tensors="pil"): + if numpify: + return_tensors = "np" + elif torchify: + return_tensors = "torch" + return prepare_video_inputs( + self.batch_size, + self.num_frames, + self.num_channels, + self.min_resolution, + self.max_resolution, + equal_resolution=equal_resolution, + return_tensors=return_tensors, + ) + + +@require_torch +@require_vision +@require_torchvision +class Molmo2VideoProcessingTest(VideoProcessingTestMixin, unittest.TestCase): + fast_video_processing_class = ( + Molmo2VideoProcessor if (is_vision_available() and is_torchvision_available()) else None + ) + video_processing_class = Molmo2VideoProcessor if (is_vision_available() and is_torchvision_available()) else None + + def setUp(self): + super().setUp() + self.video_processor_tester = Molmo2VideoProcessingTester(self) + + @property + def video_processor_dict(self): + return self.video_processor_tester.prepare_video_processor_dict() + + def test_video_processor_from_dict_with_kwargs(self): + video_processor = self.fast_video_processing_class.from_dict(self.video_processor_dict) + self.assertEqual(video_processor.size, {"height": 378, "width": 378}) + self.assertEqual(video_processor.patch_size, 14) + + video_processor = self.fast_video_processing_class.from_dict( + self.video_processor_dict, size={"height": 252, "width": 252}, patch_size=18 + ) + self.assertEqual(video_processor.size, {"height": 252, "width": 252}) + self.assertEqual(video_processor.patch_size, 18) + + def test_video_processor_properties(self): + video_processor = self.video_processing_class(**self.video_processor_dict) + self.assertTrue(hasattr(video_processor, "do_resize")) + self.assertTrue(hasattr(video_processor, "size")) + self.assertTrue(hasattr(video_processor, "do_normalize")) + self.assertTrue(hasattr(video_processor, "image_mean")) + self.assertTrue(hasattr(video_processor, "image_std")) + self.assertTrue(hasattr(video_processor, "do_convert_rgb")) + self.assertTrue(hasattr(video_processor, "patch_size")) + self.assertTrue(hasattr(video_processor, "pooling_size")) + self.assertTrue(hasattr(video_processor, "do_sample_frames")) + + def _assert_patchified_output(self, outputs, expected_num_videos): + pixel_values = outputs[self.input_name] + self.assertEqual(pixel_values.ndim, 3) + pixels_per_patch = self.video_processor_tester.patch_size**2 * self.video_processor_tester.num_channels + self.assertEqual(pixel_values.shape[-1], pixels_per_patch) + self.assertEqual(outputs["video_grids"].shape[0], expected_num_videos) + pool_h, pool_w = self.video_processor_tester.pooling_size + self.assertEqual(outputs["video_token_pooling"].ndim, 2) + self.assertEqual(outputs["video_token_pooling"].shape[-1], pool_h * pool_w) + + def test_call_numpy(self): + for video_processing_class in self.video_processor_list: + video_processing = video_processing_class(**self.video_processor_dict) + video_inputs = self.video_processor_tester.prepare_video_inputs(equal_resolution=False, numpify=True) + for video in video_inputs: + self.assertIsInstance(video, np.ndarray) + + outputs = video_processing(video_inputs[0], return_tensors="pt") + self._assert_patchified_output(outputs, 1) + + outputs = video_processing(video_inputs, return_tensors="pt") + self._assert_patchified_output(outputs, self.video_processor_tester.batch_size) + + # Molmo2 video processor expects channels-last numpy input, not channels-first torch tensors + def test_call_pytorch(self): + pass + + def test_call_pil(self): + for video_processing_class in self.video_processor_list: + video_processing = video_processing_class(**self.video_processor_dict) + video_inputs = self.video_processor_tester.prepare_video_inputs( + equal_resolution=False, return_tensors="pil" + ) + + outputs = video_processing(video_inputs[0], return_tensors="pt", input_data_format="channels_last") + self._assert_patchified_output(outputs, 1) + + outputs = video_processing(video_inputs, return_tensors="pt", input_data_format="channels_last") + self._assert_patchified_output(outputs, self.video_processor_tester.batch_size) + + def test_call_sample_frames(self): + for video_processing_class in self.video_processor_list: + video_processing = video_processing_class(**self.video_processor_dict) + video_inputs = self.video_processor_tester.prepare_video_inputs(equal_resolution=False, numpify=True) + + outputs = video_processing(video_inputs[0], return_tensors="pt", num_frames=3) + self._assert_patchified_output(outputs, 1) + + outputs = video_processing(video_inputs, return_tensors="pt", num_frames=3) + self._assert_patchified_output(outputs, self.video_processor_tester.batch_size) + + def test_call_sample_frames_short_batched_videos(self): + for video_processing_class in self.video_processor_list: + video_processing = video_processing_class(**{**self.video_processor_dict, "do_sample_frames": True}) + videos = [np.zeros((15, 32, 32, 3), dtype=np.uint8), np.zeros((7, 48, 40, 3), dtype=np.uint8)] + + outputs = video_processing(videos, return_tensors="pt", num_frames=384) + + self._assert_patchified_output(outputs, 2) + self.assertEqual(outputs["video_grids"][:, 0].tolist(), [15, 7]) + + def test_nested_input(self): + for video_processing_class in self.video_processor_list: + video_processing = video_processing_class(**self.video_processor_dict) + video_inputs = self.video_processor_tester.prepare_video_inputs( + equal_resolution=False, return_tensors="np" + ) + video_inputs = [list(video) for video in video_inputs] + + outputs = video_processing(video_inputs[0], return_tensors="pt") + self._assert_patchified_output(outputs, 1) + + outputs = video_processing(video_inputs, return_tensors="pt") + self._assert_patchified_output(outputs, self.video_processor_tester.batch_size) + + @unittest.skip( + reason="Molmo2VideoProcessor always converts to RGB before processing; 4-channel videos are not supported." + ) + def test_call_numpy_4_channels(self): + pass + + @unittest.skip( + reason="Molmo2 resizes without antialiasing to match the reference processor; under inductor the bilinear " + "kernel rounds a few pixels to the neighboring uint8 value, and that 1/255 difference exceeds the " + "test's 1e-4 tolerance." + ) + def test_can_compile_fast_video_processor(self): + pass diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 309879ab5091..09ccc3b8c5c3 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -35,6 +35,7 @@ MergeModulelist, PermuteForRope, PrefixChange, + Split, VisionFuseAndPermuteForRope, VisionUnfuseAndPermuteForRope, WeightConverter, @@ -183,11 +184,12 @@ def __init__(self, shape): class DummySelfAttn(nn.Module): - def __init__(self): + def __init__(self, attention_shapes=((1, 2), (1, 2), (1, 2))): super().__init__() - self.q_proj = DummyParamModule((1, 2)) - self.k_proj = DummyParamModule((1, 2)) - self.v_proj = DummyParamModule((1, 2)) + q_shape, k_shape, v_shape = attention_shapes + self.q_proj = DummyParamModule(q_shape) + self.k_proj = DummyParamModule(k_shape) + self.v_proj = DummyParamModule(v_shape) class DummyExperts(nn.Module): @@ -198,18 +200,20 @@ def __init__(self): class DummyLayer(nn.Module): - def __init__(self, add_extra_moe=False): + def __init__(self, add_extra_moe=False, attention_shapes=((1, 2), (1, 2), (1, 2))): super().__init__() - self.self_attn = DummySelfAttn() + self.self_attn = DummySelfAttn(attention_shapes) self.experts = DummyExperts() if add_extra_moe: self.extra_experts = DummyExperts() class DummyTopModel(nn.Module): - def __init__(self, add_extra_moe=False): + def __init__(self, add_extra_moe=False, attention_shapes=((1, 2), (1, 2), (1, 2))): super().__init__() - self.layers = nn.ModuleList([DummyLayer(add_extra_moe), DummyLayer(add_extra_moe)]) + self.layers = nn.ModuleList( + [DummyLayer(add_extra_moe, attention_shapes), DummyLayer(add_extra_moe, attention_shapes)] + ) class DummyMLP(nn.Module): @@ -222,9 +226,9 @@ class DummyModelWithNorm(PreTrainedModel): base_model_prefix = "model" config: PreTrainedConfig - def __init__(self, config, add_extra_moe=False): + def __init__(self, config, add_extra_moe=False, attention_shapes=((1, 2), (1, 2), (1, 2))): super().__init__(config) - self.model = DummyTopModel(add_extra_moe) + self.model = DummyTopModel(add_extra_moe, attention_shapes) self.norm1 = nn.LayerNorm(16) self.norm2 = nn.LayerNorm(16) self.post_init() @@ -1239,6 +1243,35 @@ def __init__(self, config, fused_qkv: bool = False): model_state = model_fused.state_dict() torch.testing.assert_close(model_state["layers.0.self_attn.qkv_proj.weight"], qkv_weight) + def test_split_to_target_parameter_sizes(self): + model = DummyModelWithNorm(PreTrainedConfig(), attention_shapes=((4, 2), (2, 2), (2, 2))) + fused = torch.arange(16.0).reshape(8, 2) + weight_mapping = [ + WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + operations=[Split(dim=0)], + ) + ] + state_dict = {f"model.layers.{index}.self_attn.qkv_proj.weight": fused.clone() for index in range(2)} + + loading_info, _ = convert_and_load_state_dict_in_model( + model, state_dict, LoadStateDictConfig(weight_mapping=weight_mapping) + ) + + self.assertEqual(loading_info.unexpected_keys, set()) + self.assertEqual(loading_info.conversion_errors, {}) + model_state = model.state_dict() + for index in range(2): + prefix = f"model.layers.{index}.self_attn" + torch.testing.assert_close(model_state[f"{prefix}.q_proj.weight"], fused[:4]) + torch.testing.assert_close(model_state[f"{prefix}.k_proj.weight"], fused[4:6]) + torch.testing.assert_close(model_state[f"{prefix}.v_proj.weight"], fused[6:]) + + reversed_state_dict = revert_weight_conversion(model, model_state) + reversed_qkv = {key: value for key, value in reversed_state_dict.items() if "qkv_proj" in key} + self.assertTrue(compare_state_dicts(reversed_qkv, state_dict)) + class TestConversionMapping(unittest.TestCase): def test_group_weight_rename(self): diff --git a/utils/check_repo.py b/utils/check_repo.py index fcefdffb84dc..8db0e503ccad 100644 --- a/utils/check_repo.py +++ b/utils/check_repo.py @@ -254,6 +254,10 @@ "Qwen3OmniMoeThinkerTextModel", "Qwen3OmniMoeForConditionalGeneration", # Bigger model tested through Qwen3OmniMoeForConditionalGenerationIntegrationTest. "Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration", # Building part of bigger (tested) model. Tested implicitly through Qwen3OmniMoeForConditionalGenerationIntegrationTest. + "Molmo2TextModel", # Building part of bigger (tested) model. Tested implicitly through Molmo2ForConditionalGeneration. + "Molmo2VisionModel", # Building part of bigger (tested) model. Tested implicitly through Molmo2ForConditionalGeneration. + "Molmo2VisionBackbone", # Building part of bigger (tested) model. Tested implicitly through Molmo2ForConditionalGeneration. + "Molmo2Adapter", # Building part of bigger (tested) model. Tested implicitly through Molmo2ForConditionalGeneration. "MllamaTextModel", # Building part of bigger (tested) model. # TODO: add tests "MllamaVisionModel", # Building part of bigger (tested) model. # TODO: add tests "Llama4TextModel", # Building part of bigger (tested) model. # TODO: add tests @@ -505,6 +509,9 @@ "MetaClip2TextModelWithProjection", "MetaClip2VisionModel", "MetaClip2VisionModelWithProjection", + "Molmo2VisionModel", + "Molmo2VisionBackbone", + "Molmo2Adapter", # Building part of bigger (tested) model "MoshiForConditionalGeneration", # no auto class for speech-to-speech "Emu3VQVAE", # no autoclass for VQ-VAE models "Emu3TextModel", # Building part of bigger (tested) model