Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

chat_template.jinja

@@ -0,0 +1 @@
+{% set DEMO_STYLES = ['point_count','pointing','cosyn_point','user_qa','long_caption','short_caption','video_long_caption','video_short_caption','video_point_track_per_frame','video_point_track_start_end','video_point_track_all_frames','video_single_point_track_start_end','video_transcript','video_clip_caption_start_end','video_clip_caption_start_end_in_seconds','video_clip_transcript_start_end','video_clip_transcript_start_end_in_seconds','video_frame_caption_timestamp','video_frame_caption_timestamp_in_seconds','correction_qa','text_sft','video_point','video_point_count','video_count','video_count_point','multi_image_pointing','multi_image_counting','multi_image_point_then_count','multi_image_count_then_point','demo','a_okvqa_mc','ai2_diagram_no_letter','ai2_diagram','science_qa','multi_image_mc','multi_image_mc_exp','mantis_instruct_mc','video_multiple_choice','video_multiple_choice_count_without_pointing','video_multiple_choice_multiple_correct','video_multiple_choice_w_subtitle'] %}{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% set has_subtitle = messages and messages[0]['role'].lower() == 'subtitle' %}{% for message in messages %}{% if message['content'] is not string %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% elif content['type'] == 'video' or 'video' in content or 'video_url' in content %}{% set video_count.value = video_count.value + 1 %}{% endif %}{% endfor %}{% endif %}{% endfor %}{% if image_count.value == 1 %}{{ '<|image|>' }}{% elif image_count.value > 1 %}{% for i in range(image_count.value) %}{{ 'Image ' ~ (i + 1) ~ '<|image|>' }}{% endfor %}{% endif %}{% for _ in range(video_count.value) %}{{ '<|video|>' }}{% endfor %}{% if has_subtitle %}{{ messages[0]['content'] }}{% endif %}{% for message in messages %}{% set role = message['role'].lower() %}{% if role == 'subtitle' %}{% continue %}{% endif %}{% set conv_index = loop.index - (1 if has_subtitle else 0) %}{%- if (conv_index % 2 == 1 and role != 'user') or (conv_index % 2 == 0 and role != 'assistant') -%}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{%- endif -%}{% if message['content'] is string %}{% set text_content = message['content'] %}{% else %}{% set m = namespace(text='') %}{% for content in message['content'] %}{% if content['type'] == 'text' %}{% if content['style'] is defined and content['style'] not in DEMO_STYLES %}{% set seg = content['style'] ~ ': ' ~ content['text'] %}{% else %}{% set seg = content['text'] %}{% endif %}{% set m.text = m.text ~ ('' if not m.text else ' ') ~ seg %}{% endif %}{% endfor %}{% set text_content = m.text %}{% endif %}{% if role == 'user' %}{% if not (has_subtitle and loop.index == 2) and not (not has_subtitle and loop.first) %}{{ '<|im_end|>\n' }}{% endif %}{{ '<|im_start|>user\n' }}{{ text_content }}{{ '<|im_end|>\n' }}{% else %} {# assistant #}{{ '<|im_start|>assistant\n' }}{{ text_content }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}

config.json

@@ -0,0 +1,111 @@
+{
+  "adapter_config": {
+    "attention_dropout": 0.0,
+    "attn_implementation": "sdpa",
+    "float32_attention": true,
+    "head_dim": 72,
+    "hidden_act": "silu",
+    "hidden_size": 1152,
+    "image_feature_dropout": 0,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "model_type": "molmo_point",
+    "num_attention_heads": 16,
+    "num_key_value_heads": 16,
+    "pooling_attention_mask": true,
+    "positional_embeddings": null,
+    "residual_dropout": 0.0,
+    "text_hidden_size": 4096,
+    "vit_layers": [
+      -3,
+      -9
+    ]
+  },
+  "architectures": [
+    "MolmoPointForConditionalGeneration"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_molmo_point.MolmoPointConfig",
+    "AutoModelForImageTextToText": "modeling_molmo_point.MolmoPointForConditionalGeneration"
+  },
+  "dtype": "float32",
+  "embed_location": false,
+  "embed_selected_vit_patch": "linear",
+  "frame_end_token_id": 151944,
+  "frame_start_token_id": 151943,
+  "image_col_id": 151939,
+  "image_end_token_id": 151937,
+  "image_high_res_id": 151938,
+  "image_patch_id": 151938,
+  "image_start_token_id": 151936,
+  "initializer_range": 0.02,
+  "layer_norm_x": true,
+  "location_token_id": 151949,
+  "low_res_image_start_token_id": null,
+  "mask_patches": "always",
+  "mask_repeats": "inference",
+  "mask_subpatches": "inference",
+  "model_type": "molmo_point",
+  "no_more_points_class": true,
+  "norm_logits": true,
+  "patch_embed_dim": 512,
+  "patch_embedding_kind": "image_feature0",
+  "patch_location": "3x3",
+  "patch_token_id": 151947,
+  "subpatch_token_id": 151948,
+  "text_config": {
+    "additional_vocab_size": 128,
+    "attention_dropout": 0.0,
+    "attn_implementation": "sdpa",
+    "embedding_dropout": 0.0,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 12288,
+    "layer_norm_eps": 1e-06,
+    "max_position_embeddings": 37376,
+    "model_type": "molmo2_text",
+    "norm_after": false,
+    "num_attention_heads": 32,
+    "num_hidden_layers": 36,
+    "num_key_value_heads": 8,
+    "qk_norm_type": "qwen3",
+    "qkv_bias": false,
+    "residual_dropout": 0.0,
+    "rope_scaling": null,
+    "rope_scaling_layers": null,
+    "rope_theta": 1000000.0,
+    "use_cache": true,
+    "use_qk_norm": true,
+    "vocab_size": 151936
+  },
+  "tie_word_embeddings": false,
+  "token_prediction_rotary": "one_d",
+  "token_prediction_rotary_theta": 50000.0,
+  "transformers_version": "4.57.6",
+  "use_cache": true,
+  "use_frame_special_tokens": true,
+  "vit_config": {
+    "attention_dropout": 0.0,
+    "attn_implementation": "sdpa",
+    "float32_attention": true,
+    "head_dim": 72,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "image_default_input_size": [
+      378,
+      378
+    ],
+    "image_num_pos": 729,
+    "image_patch_size": 14,
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "layer_norm_eps": 1e-06,
+    "model_type": "molmo2",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 27,
+    "num_key_value_heads": 16,
+    "residual_dropout": 0.0
+  }
+}

configuration_molmo2.py

@@ -0,0 +1,391 @@
+"""
+Molmo2 configuration
+"""
+
+from typing import Optional, Any
+
+from transformers import PretrainedConfig
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class Molmo2VitConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Molmo2VisionTransformer`].
+    It is used to instantiate a `Molmo2VisionTransformer` according to the specified arguments,
+    defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Example:
+    ```python
+    >>> from transformers import Molmo2VitConfig, Molmo2VisionTransformer
+
+    >>> # Initializing a Molmo2VitConfig
+    >>> configuration = Molmo2VitConfig()
+
+    >>> # Initializing a Molmo2VisionTransformer (with random weights)
+    >>> model = Molmo2VisionTransformer(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "molmo2"
+    base_config_key = "vit_config"
+
+    def __init__(
+        self,
+        hidden_size: int = 1152,
+        intermediate_size: int = 4304,
+        num_hidden_layers: int = 27,
+        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_default_input_size: tuple[int, int] = (378, 378),
+        image_patch_size: int = 14,
+        image_num_pos: int = 577,
+        attention_dropout: float = 0.0,
+        residual_dropout: float = 0.0,
+        initializer_range: float = 0.02,
+        float32_attention: bool = True,
+        attn_implementation: str = "eager",
+        **kwargs,
+    ):
+        self.attn_implementation = attn_implementation
+        super().__init__(
+            attn_implementation=attn_implementation,
+            **kwargs
+        )
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.layer_norm_eps = layer_norm_eps
+        self.image_default_input_size = image_default_input_size
+        self.image_patch_size = image_patch_size
+        self.image_num_pos = image_num_pos
+        self.attention_dropout = attention_dropout
+        self.residual_dropout = residual_dropout
+        self.initializer_range = initializer_range
+        self.float32_attention = float32_attention
+
+    @property
+    def image_num_patch(self):
+        h, w = self.image_default_input_size
+        return h // self.image_patch_size, w // self.image_patch_size
+
+
+class Molmo2AdapterConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of Molmo2Adapter. With Molmo2VitConfig,
+    It is used to instantiate an Molmo2VisionBackbone according to the specified arguments,
+    defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Example:
+
+    ```python
+    >>> from transformers import Molmo2VitConfig, Molmo2AdapterConfig, Molmo2VisionBackbone
+
+    >>> # Initializing a Molmo2VitConfig and a Molmo2AdapterConfig
+    >>> vit_config = Molmo2VitConfig()
+    >>> adapter_config = MolmoPoolingConfig()
+
+    >>> # Initializing a Molmo2VisionBackbone (with random weights)
+    >>> model = Molmo2VisionBackbone(vit_config, adapter_config)
+
+    >>> # Accessing the model configuration
+    >>> vit_configuration = model.vit_config
+    >>> adapter_configuration = model.adapter_config
+    ```"""
+
+    model_type = "molmo2"
+    base_config_key = "adapter_config"
+
+    def __init__(
+        self,
+        vit_layers: tuple = (-3, -9),
+        pooling_attention_mask: bool = False,
+        hidden_size: int = 1152,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 16,
+        head_dim: int = 72,
+        float32_attention: bool = True,
+        attention_dropout: float = 0.0,
+        residual_dropout: float = 0.0,
+        hidden_act: str = "silu",
+        intermediate_size: int = 18944,
+        text_hidden_size: int = 3584,
+        image_feature_dropout: float = 0.0,
+        initializer_range: float = 0.02,
+        attn_implementation: str = "eager",
+        **kwargs,
+    ):
+        self.attn_implementation = attn_implementation
+        super().__init__(
+            attn_implementation=attn_implementation,
+            **kwargs
+        )
+        self.vit_layers = vit_layers
+        self.pooling_attention_mask = pooling_attention_mask
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.float32_attention = float32_attention
+        self.attention_dropout = attention_dropout
+        self.residual_dropout = residual_dropout
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.text_hidden_size = text_hidden_size
+        self.image_feature_dropout = image_feature_dropout
+        self.initializer_range = initializer_range
+
+
+class Molmo2TextConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Molmo2TextModel`]. It is used to instantiate a
+    `Molmo2TextModel` according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Example:
+    ```python
+    >>> from transformers import Molmo2TextConfig, Molmo2TextModel
+
+    >>> # Initializing a Molmo2TextConfig
+    >>> configuration = Molmo2TextConfig()
+
+    >>> # Initializing a Molmo2TextModel (with random weights)
+    >>> model = Molmo2TextModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "molmo2_text"
+    base_config_key = "text_config"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "blocks.*.self_attn.att_proj": "colwise",
+        "blocks.*.self_attn.attn_out": "rowwise",
+        "blocks.*.mlp.ff_proj": "colwise",
+        "blocks.*.mlp.ff_out": "rowwise",
+    }
+    base_model_pp_plan = {
+        "wte": (["input_ids"], ["inputs_embeds"]),
+        "blocks": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "ln_f": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        hidden_size: int = 3584,
+        num_attention_heads: int = 28,
+        num_key_value_heads: Optional[int] = 4,
+        head_dim: int = 128,
+        vocab_size: int = 152064,
+        additional_vocab_size: int = 128,
+        qkv_bias: bool = True,
+        num_hidden_layers: int = 48,
+        intermediate_size: int = 18944,
+        hidden_act: str = "silu",
+        embedding_dropout: float=0.0,
+        attention_dropout: float=0.0,
+        residual_dropout: float = 0.0,
+        max_position_embeddings: int = 4096,
+        rope_theta: float = 1000000.0,
+        rope_scaling: dict[str, Any] = None,
+        rope_scaling_layers: Optional[list[int]] = None,
+        use_qk_norm: bool = False,
+        qk_norm_type: str = "olmo",
+        layer_norm_eps: int = 1e-6,
+        norm_after: bool = False,
+        initializer_range: float = 0.02,
+        use_cache=True,
+        tie_word_embeddings=False,
+        attn_implementation: str = "eager",
+        **kwargs,
+    ):
+        self.attn_implementation = attn_implementation
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            attn_implementation=attn_implementation,
+            **kwargs
+        )
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.vocab_size = vocab_size
+        self.additional_vocab_size = additional_vocab_size
+        self.qkv_bias = qkv_bias
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.embedding_dropout = embedding_dropout
+        self.attention_dropout = attention_dropout
+        self.residual_dropout = residual_dropout
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.rope_scaling_layers = rope_scaling_layers
+        self.use_qk_norm = use_qk_norm
+        self.qk_norm_type = qk_norm_type
+        self.layer_norm_eps = layer_norm_eps
+        self.norm_after = norm_after
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+
+        # Validate the correctness of rotary position embeddings parameters
+        rope_config_validation(self)
+
+
+class Molmo2Config(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Molmo2ForConditionalGeneration`].
+    It is used to instantiate an Molmo2 model according to the specified arguments, defining the model architecture.
+
+    Example:
+
+    ```python
+    >>> from transformers import Molmo2Config, Molmo2VitConfig, Molmo2AdapterConfig, Molmo2TextConfig
+
+    >>> # Initializing a Molmo2VitConfig
+    >>> vit_config = Molmo2VitConfig()
+
+    >>> # Initializing a Molmo2AdapterConfig
+    >>> adapter_config = Molmo2AdapterConfig()
+
+    >>> # Initializing a Molmo2TextConfig
+    >>> text_config = Molmo2TextConfig()
+
+    >>> # Initializing a Molmo2Config
+    >>> configuration = Molmo2Config(
+    >>>     vit_config=vit_config,
+    >>>     adapter_config=adapter_config,
+    >>>     text_config=text_config,
+    >>>     image_start_token_id=151936,
+    >>>     image_end_token_id=151937,
+    >>>     image_patch_id=151938,
+    >>>     image_col_id=151939,
+    >>>     low_res_image_start_token_id=151940,
+    >>>     image_low_res_id=151942,
+    >>>     frame_start_token_id=151943,
+    >>>     frame_end_token_id=151944,
+    >>> )
+
+    >>> # Initializing a model
+    >>> model = Molmo2ForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "molmo2"
+    sub_configs = {
+        "text_config": Molmo2TextConfig,
+        "vit_config": Molmo2VitConfig,
+        "adapter_config": Molmo2AdapterConfig,
+    }
+
+    def __init__(
+        self,
+        vit_config: Molmo2VitConfig = None,
+        adapter_config: Molmo2AdapterConfig = None,
+        text_config: Molmo2TextConfig = None,
+        image_start_token_id: int = None,
+        low_res_image_start_token_id: int = None,
+        image_end_token_id: int = None,
+        image_low_res_id: int = None,
+        image_patch_id: int = None,
+        image_col_id: int = None,
+        frame_start_token_id: int = None,
+        frame_end_token_id: int = None,
+        use_frame_special_tokens: bool = True,
+        initializer_range: float = 0.02,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if vit_config is None:
+            self.vit_config = Molmo2VitConfig()
+        elif isinstance(vit_config, dict):
+            self.vit_config = Molmo2VitConfig(**vit_config)
+        else:
+            self.vit_config = vit_config
+        if adapter_config is None:
+            self.adapter_config = Molmo2AdapterConfig()
+        elif isinstance(adapter_config, dict):
+            self.adapter_config = Molmo2AdapterConfig(**adapter_config)
+        else:
+            self.adapter_config = adapter_config
+        if text_config is None:
+            self.text_config = Molmo2TextConfig()
+        elif isinstance(text_config, dict):
+            self.text_config = Molmo2TextConfig(**text_config)
+        else:
+            self.text_config = text_config
+        self.image_start_token_id = image_start_token_id
+        self.low_res_image_start_token_id = low_res_image_start_token_id
+        self.image_end_token_id = image_end_token_id
+        self.image_low_res_id = image_low_res_id
+        self.image_high_res_id = image_patch_id
+        self.image_patch_id = image_patch_id
+        self.image_col_id = image_col_id
+        self.frame_start_token_id = frame_start_token_id
+        self.frame_end_token_id = frame_end_token_id
+        self.use_frame_special_tokens = use_frame_special_tokens
+        self.initializer_range = initializer_range
+
+    @property
+    def image_num_patch(self):
+        assert self.vit_config is not None
+        return self.vit_config.image_num_patch
+    
+    @property
+    def num_attention_heads(self):
+        return self.text_config.num_attention_heads
+    
+    @property
+    def num_key_value_heads(self):
+        return self.text_config.num_key_value_heads
+
+    @property
+    def head_dim(self):
+        return self.text_config.head_dim
+
+    @property
+    def num_hidden_layers(self):
+        return self.text_config.num_hidden_layers
+    
+    @property
+    def hidden_size(self):
+        return self.text_config.hidden_size
+    
+    @property
+    def vocab_size(self):
+        return self.text_config.vocab_size
+    
+    @property
+    def max_position_embeddings(self):
+        return self.text_config.max_position_embeddings
+
+
+Molmo2VitConfig.register_for_auto_class()
+Molmo2AdapterConfig.register_for_auto_class()
+Molmo2TextConfig.register_for_auto_class()
+Molmo2Config.register_for_auto_class()

configuration_molmo_point.py

@@ -0,0 +1,251 @@
+"""
+Molmo2 configuration
+"""
+
+from typing import Optional
+
+from transformers import PretrainedConfig, LogitsProcessor
+from transformers.utils import logging
+
+from .configuration_molmo2 import Molmo2TextConfig, Molmo2VitConfig, \
+    Molmo2AdapterConfig
+
+logger = logging.get_logger(__name__)
+
+
+class MolmoPointAdapterConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of Molmo2Adapter. With Molmo2VitConfig,
+    It is used to instantiate an Molmo2VisionBackbone according to the specified arguments,
+    defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Example:
+
+    ```python
+    >>> from transformers import Molmo2VitConfig, Molmo2AdapterConfig, Molmo2VisionBackbone
+
+    >>> # Initializing a Molmo2VitConfig and a Molmo2AdapterConfig
+    >>> vit_config = Molmo2VitConfig()
+    >>> adapter_config = MolmoPoolingConfig()
+
+    >>> # Initializing a Molmo2VisionBackbone (with random weights)
+    >>> model = Molmo2VisionBackbone(vit_config, adapter_config)
+
+    >>> # Accessing the model configuration
+    >>> vit_configuration = model.vit_config
+    >>> adapter_configuration = model.adapter_config
+    ```"""
+
+    model_type = "molmo_point"
+    base_config_key = "adapter_config"
+
+    def __init__(
+        self,
+        vit_layers: tuple = (-3, -9),
+        pooling_attention_mask: bool = False,
+        hidden_size: int = 1152,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 16,
+        head_dim: int = 72,
+        float32_attention: bool = True,
+        attention_dropout: float = 0.0,
+        residual_dropout: float = 0.0,
+        hidden_act: str = "silu",
+        intermediate_size: int = 18944,
+        text_hidden_size: int = 3584,
+        image_feature_dropout: float = 0.0,
+        initializer_range: float = 0.02,
+        attn_implementation: str = "eager",
+        positional_embeddings: int = 16,
+        **kwargs,
+    ):
+        self.attn_implementation = attn_implementation
+        super().__init__(
+            attn_implementation=attn_implementation,
+            **kwargs
+        )
+        self.vit_layers = vit_layers
+        self.pooling_attention_mask = pooling_attention_mask
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.float32_attention = float32_attention
+        self.attention_dropout = attention_dropout
+        self.residual_dropout = residual_dropout
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.text_hidden_size = text_hidden_size
+        self.image_feature_dropout = image_feature_dropout
+        self.initializer_range = initializer_range
+        self.positional_embeddings = positional_embeddings
+
+
+class MolmoPointConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MolmoPointForConditionalGeneration`].
+    It is used to instantiate an Molmo2 model according to the specified arguments, defining the model architecture.
+
+    Example:
+
+    ```python
+    >>> from transformers import Molmo2Config, Molmo2VitConfig, Molmo2AdapterConfig, Molmo2TextConfig
+
+    >>> # Initializing a Molmo2VitConfig
+    >>> vit_config = Molmo2VitConfig()
+
+    >>> # Initializing a Molmo2AdapterConfig
+    >>> adapter_config = MolmoPointAdapterConfig()
+
+    >>> # Initializing a Molmo2TextConfig
+    >>> text_config = Molmo2TextConfig()
+
+    >>> # Initializing a Molmo2Config
+    >>> configuration = MolmoPointConfig(
+    >>>     vit_config=vit_config,
+    >>>     adapter_config=adapter_config,
+    >>>     text_config=text_config,
+    >>>     image_start_token_id=151936,
+    >>>     image_end_token_id=151937,
+    >>>     image_patch_id=151938,
+    >>>     image_col_id=151939,
+    >>>     low_res_image_start_token_id=151940,
+    >>>     image_low_res_id=151942,
+    >>>     frame_start_token_id=151943,
+    >>>     frame_end_token_id=151944,
+    >>> )
+
+    >>> # Initializing a model
+    >>> model = MolmoPointForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "molmo_point"
+    sub_configs = {
+        "text_config": Molmo2TextConfig,
+        "vit_config": Molmo2VitConfig,
+        "adapter_config": MolmoPointAdapterConfig,
+    }
+
+    def __init__(
+        self,
+        vit_config: Molmo2VitConfig = None,
+        adapter_config: MolmoPointAdapterConfig = None,
+        text_config: Molmo2TextConfig = None,
+        image_start_token_id: int = None,
+        low_res_image_start_token_id: int = None,
+        image_end_token_id: int = None,
+        image_patch_id: int = None,
+        image_col_id: int = None,
+        frame_start_token_id: int = None,
+        frame_end_token_id: int = None,
+        patch_token_id: int = None,
+        subpatch_token_id: int = None,
+        location_token_id: int = None,
+        use_frame_special_tokens: bool = True,
+        initializer_range: float = 0.02,
+
+        # point config
+        patch_location: Optional[str]="3x3",
+        no_more_points_class: bool=False,
+        patch_embed_dim: int=256,
+        patch_embedding_kind: str="linear",
+        embed_selected_vit_patch: Optional[str]="linear",
+        embed_location: bool=False,
+        layer_norm_x: bool=True,
+        norm_logits: bool=True,
+        # FIXME figure out how infernce params work
+        mask_patches: Optional[str]="always",
+        mask_subpatches: str="inference",
+        mask_repeats: Optional[str]="inference",
+        token_prediction_rotary: bool=True,
+        token_prediction_rotary_theta: Optional[float]=50000,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if vit_config is None:
+            self.vit_config = Molmo2VitConfig()
+        elif isinstance(vit_config, dict):
+            self.vit_config = Molmo2VitConfig(**vit_config)
+        else:
+            self.vit_config = vit_config
+        if adapter_config is None:
+            self.adapter_config = Molmo2AdapterConfig()
+        elif isinstance(adapter_config, dict):
+            self.adapter_config = Molmo2AdapterConfig(**adapter_config)
+        else:
+            self.adapter_config = adapter_config
+        if text_config is None:
+            self.text_config = Molmo2TextConfig()
+        elif isinstance(text_config, dict):
+            self.text_config = Molmo2TextConfig(**text_config)
+        else:
+            self.text_config = text_config
+        self.image_start_token_id = image_start_token_id
+        self.low_res_image_start_token_id = low_res_image_start_token_id
+        self.image_end_token_id = image_end_token_id
+        self.image_high_res_id = image_patch_id
+        self.image_patch_id = image_patch_id
+        self.image_col_id = image_col_id
+        self.frame_start_token_id = frame_start_token_id
+        self.frame_end_token_id = frame_end_token_id
+        self.patch_token_id = patch_token_id
+        self.subpatch_token_id = subpatch_token_id
+        self.location_token_id = location_token_id
+        self.use_frame_special_tokens = use_frame_special_tokens
+        self.initializer_range = initializer_range
+        self.patch_location = patch_location
+        self.no_more_points_class = no_more_points_class
+        self.patch_embed_dim = patch_embed_dim
+        self.patch_embedding_kind = patch_embedding_kind
+        self.embed_selected_vit_patch = embed_selected_vit_patch
+        self.embed_location = embed_location
+        self.layer_norm_x = layer_norm_x
+        self.norm_logits = norm_logits
+        self.mask_patches = mask_patches
+        self.mask_subpatches = mask_subpatches
+        self.mask_repeats = mask_repeats
+        self.token_prediction_rotary = token_prediction_rotary
+        self.token_prediction_rotary_theta = token_prediction_rotary_theta
+
+    @property
+    def image_num_patch(self):
+        assert self.vit_config is not None
+        return self.vit_config.image_num_patch
+    
+    @property
+    def num_attention_heads(self):
+        return self.text_config.num_attention_heads
+    
+    @property
+    def num_key_value_heads(self):
+        return self.text_config.num_key_value_heads
+
+    @property
+    def head_dim(self):
+        return self.text_config.head_dim
+
+    @property
+    def num_hidden_layers(self):
+        return self.text_config.num_hidden_layers
+    
+    @property
+    def hidden_size(self):
+        return self.text_config.hidden_size
+    
+    @property
+    def vocab_size(self):
+        return self.text_config.vocab_size
+    
+    @property
+    def max_position_embeddings(self):
+        return self.text_config.max_position_embeddings
+
+
+MolmoPointAdapterConfig.register_for_auto_class()
+MolmoPointConfig.register_for_auto_class()

convert_molmo2_to_hf.py

@@ -0,0 +1,511 @@
+import argparse
+import os
+import shutil
+import logging
+import json
+import gc
+from typing import Dict, Any, Optional
+
+import torch
+from transformers import GenerationConfig
+from transformers.image_utils import (
+    PILImageResampling,
+    IMAGENET_STANDARD_MEAN,
+    IMAGENET_STANDARD_STD,
+)
+
+from olmo.models.molmo2.molmo2 import Molmo2Config as ModelConfig
+from olmo.train.checkpointer import load_model_state
+from olmo.util import (
+    prepare_cli_environment,
+    resource_path,
+    select_checkpoint
+)
+
+from .configuration_molmo2 import Molmo2Config, Molmo2VitConfig, Molmo2AdapterConfig, Molmo2TextConfig
+from .modeling_molmo2 import Molmo2ForConditionalGeneration
+from .processing_molmo2 import Molmo2Processor
+from .image_processing_molmo2 import Molmo2ImageProcessor
+from .video_processing_molmo2 import Molmo2VideoProcessor
+
+
+logger = logging.getLogger(__name__)
+
+
+CHAT_TEMPLATE = (
+    "{% set DEMO_STYLES = ["
+    "'point_count','pointing','cosyn_point','user_qa','long_caption','short_caption',"
+    "'video_long_caption','video_short_caption','video_point_track_per_frame',"
+    "'video_point_track_start_end','video_point_track_all_frames','video_single_point_track_start_end',"
+    "'video_transcript','video_clip_caption_start_end','video_clip_caption_start_end_in_seconds',"
+    "'video_clip_transcript_start_end','video_clip_transcript_start_end_in_seconds',"
+    "'video_frame_caption_timestamp','video_frame_caption_timestamp_in_seconds',"
+    "'correction_qa','text_sft','video_point','video_point_count','video_count','video_count_point',"
+    "'multi_image_pointing','multi_image_counting','multi_image_point_then_count','multi_image_count_then_point','demo',"
+    "'a_okvqa_mc','ai2_diagram_no_letter','ai2_diagram','science_qa',"
+    "'multi_image_mc','multi_image_mc_exp','mantis_instruct_mc',"
+    "'video_multiple_choice','video_multiple_choice_count_without_pointing',"
+    "'video_multiple_choice_multiple_correct','video_multiple_choice_w_subtitle'"
+    "] %}"
+
+    "{% set image_count = namespace(value=0) %}"
+    "{% set video_count = namespace(value=0) %}"
+
+    "{% set has_subtitle = messages and messages[0]['role'].lower() == 'subtitle' %}"
+
+    "{% for message in messages %}"
+        "{% if message['content'] is not string %}"
+            "{% for content in message['content'] %}"
+                "{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}"
+                    "{% set image_count.value = image_count.value + 1 %}"
+                "{% elif content['type'] == 'video' or 'video' in content or 'video_url' in content %}"
+                    "{% set video_count.value = video_count.value + 1 %}"
+                "{% endif %}"
+            "{% endfor %}"
+        "{% endif %}"
+    "{% endfor %}"
+
+    "{% if image_count.value == 1 %}"
+        "{{ '<|image|>' }}"
+    "{% elif image_count.value > 1 %}"
+        "{% for i in range(image_count.value) %}"
+            "{{ 'Image ' ~ (i + 1) ~ '<|image|>' }}"
+        "{% endfor %}"
+    "{% endif %}"
+
+    "{% for _ in range(video_count.value) %}"
+        "{{ '<|video|>' }}"
+    "{% endfor %}"
+
+    "{% if has_subtitle %}"
+        "{{ messages[0]['content'] }}"
+    "{% endif %}"
+
+    "{% for message in messages %}"
+        "{% set role = message['role'].lower() %}"
+
+        "{% if role == 'subtitle' %}"
+            "{% continue %}"
+        "{% endif %}"
+
+        "{% set conv_index = loop.index - (1 if has_subtitle else 0) %}"
+
+        "{%- if (conv_index % 2 == 1 and role != 'user') "
+        "or (conv_index % 2 == 0 and role != 'assistant') -%}"
+            "{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}"
+        "{%- endif -%}"
+
+        "{% if message['content'] is string %}"
+            "{% set text_content = message['content'] %}"
+        "{% else %}"
+            "{% set m = namespace(text='') %}"
+            "{% for content in message['content'] %}"
+                "{% if content['type'] == 'text' %}"
+                    "{% if content['style'] is defined and content['style'] not in DEMO_STYLES %}"
+                        "{% set seg = content['style'] ~ ': ' ~ content['text'] %}"
+                    "{% else %}"
+                        "{% set seg = content['text'] %}"
+                    "{% endif %}"
+                    "{% set m.text = m.text ~ ('' if not m.text else ' ') ~ seg %}"
+                "{% endif %}"
+            "{% endfor %}"
+            "{% set text_content = m.text %}"
+        "{% endif %}"
+
+        "{% if role == 'user' %}"
+            "{% if not (has_subtitle and loop.index == 2) and not (not has_subtitle and loop.first) %}{{ '<|im_end|>\\n' }}{% endif %}"
+            "{{ '<|im_start|>user\\n' }}"
+            "{{ text_content }}"
+            "{{ '<|im_end|>\\n' }}"
+        "{% else %} {# assistant #}"
+            "{{ '<|im_start|>assistant\\n' }}"
+            "{{ text_content }}"
+        "{% endif %}"
+    "{% endfor %}"
+
+    "{% if add_generation_prompt %}"
+        "{{ '<|im_start|>assistant\\n' }}"
+    "{% endif %}"
+)
+
+
+def convert_config(
+    model_config: ModelConfig,
+    attn_implementation: str,
+    override_max_model_len: Optional[int],
+) -> Molmo2Config:
+    """Convert config to HF-compatible config"""
+    vision_backbone_cfg = model_config.vision_backbone
+    vit_config = vision_backbone_cfg.vit
+    llm_config = model_config.llm
+
+    molmo2_vit_config = Molmo2VitConfig(
+        hidden_size=vit_config.image_emb_dim,
+        intermediate_size=vit_config.image_mlp_dim,
+        num_hidden_layers=vit_config.image_num_layers,
+        num_attention_heads=vit_config.image_num_heads,
+        num_key_value_heads=vit_config.image_num_key_value_heads,
+        head_dim=vit_config.image_head_dim,
+        hidden_act=vit_config.image_mlp_activations,
+        layer_norm_eps=vit_config.image_norm_eps,
+        image_default_input_size=vit_config.image_default_input_size,
+        image_patch_size=vit_config.image_patch_size,
+        image_num_pos=vit_config.image_num_pos,
+        attention_dropout=0.0,
+        residual_dropout=0.0,
+        initializer_range=vit_config.initializer_range,
+        float32_attention=vit_config.float32_attention,
+        attn_implementation=attn_implementation,
+    )
+    adapter_hidden_act = "silu" if llm_config.activation_type == "swiglu" else llm_config.activation_type
+    adapter_intermediate_size = (
+        llm_config.mlp_hidden_size if llm_config.mlp_hidden_size is not None
+        else llm_config.mlp_ratio * llm_config.d_model
+    ) // 2
+    molmo2_adapter_config = Molmo2AdapterConfig(
+        vit_layers=vision_backbone_cfg.vit_layers,
+        pooling_attention_mask=vision_backbone_cfg.pooling_attention_mask,
+        hidden_size=vit_config.image_emb_dim,
+        num_attention_heads=vit_config.image_num_heads,
+        num_key_value_heads=vit_config.image_num_key_value_heads,
+        head_dim=vit_config.image_head_dim,
+        float32_attention=vit_config.float32_attention,
+        attention_dropout=0.0,
+        residual_dropout=0.0,
+        hidden_act=adapter_hidden_act,
+        intermediate_size=adapter_intermediate_size,
+        text_hidden_size=llm_config.d_model,
+        image_feature_dropout=vision_backbone_cfg.image_feature_dropout,
+        initializer_range=llm_config.initializer_range,
+        attn_implementation=attn_implementation,
+    )
+    llm_head_dim = llm_config.d_model // llm_config.n_heads if llm_config.head_dim is None else llm_config.head_dim
+    llm_intermediate_size = (
+        llm_config.mlp_hidden_size if llm_config.mlp_hidden_size is not None
+        else llm_config.mlp_ratio * llm_config.d_model
+    ) // 2
+    llm_hidden_act = "silu" if llm_config.activation_type == "swiglu" else llm_config.activation_type
+    rope_scaling: Optional[Dict[str, Any]] = None
+    if llm_config.rope_type != "default":
+        rope_scaling = dict(rope_type=llm_config.rope_type)
+        for key in [
+            "rope_factor",
+            "rope_high_freq_factor",
+            "rope_low_freq_factor",
+            "rope_attention_factor",
+            "rope_original_max_position_embeddings",
+            "rope_beta_fast",
+            "rope_beta_slow",
+            "rope_mscale",
+            "rope_mscale_all_dim",
+            "rope_truncate",
+        ]:
+            if getattr(llm_config, key) is not None:
+                rope_scaling[key[len("rope_"):]] = getattr(llm_config, key)
+    
+    max_position_embeddings = llm_config.max_position_embeddings or llm_config.max_sequence_length
+    if override_max_model_len is not None:
+        max_position_embeddings = override_max_model_len
+    rope_scaling_layers: list[int] | None = None
+    if llm_config.full_attention_layers is not None:
+        # HACK: The original Olmo3 applies scaling to full attention layers,
+        # while we applies scaling to slinding attention layers.
+        if llm_config.sliding_attention_rope_scaling:
+            rope_scaling_layers = [idx for idx in range(llm_config.n_layers) if idx not in llm_config.full_attention_layers]
+        else:
+            rope_scaling_layers = list(llm_config.full_attention_layers)
+    molmo2_text_config = Molmo2TextConfig(
+        hidden_size=llm_config.d_model,
+        num_attention_heads=llm_config.n_heads,
+        num_key_value_heads=llm_config.effective_n_kv_heads,
+        head_dim=llm_head_dim,
+        vocab_size=llm_config.embedding_size or llm_config.vocab_size,
+        additional_vocab_size=llm_config.additional_vocab_size,
+        qkv_bias=llm_config.qkv_bias,
+        num_hidden_layers=llm_config.n_layers,
+        intermediate_size=llm_intermediate_size,
+        hidden_act=llm_hidden_act,
+        embedding_dropout=0.0,
+        attention_dropout=0.0,
+        residual_dropout=0.0,
+        max_position_embeddings=max_position_embeddings,
+        rope_theta=llm_config.rope_theta,
+        rope_scaling=rope_scaling,
+        rope_scaling_layers=rope_scaling_layers,
+        use_qk_norm=llm_config.attention_layer_norm,
+        qk_norm_type=llm_config.attention_layer_norm_type,
+        layer_norm_eps=llm_config.layer_norm_eps,
+        norm_after=llm_config.norm_after,
+        initializer_range=llm_config.initializer_range,
+        attn_implementation=attn_implementation,
+    )
+
+    tokenizer = model_config.build_tokenizer()
+    image_start_token_id = tokenizer.image_start_token_id
+    image_end_token_id = tokenizer.image_end_token_id
+    low_res_image_start_token_id = tokenizer.low_res_image_start_token_id
+    image_low_res_id = tokenizer.image_low_res_token_id
+    image_patch_id = tokenizer.image_patch_token_id
+    image_col_id = tokenizer.image_col_token_id
+    frame_start_token_id = tokenizer.frame_start_token_id
+    frame_end_token_id = tokenizer.frame_end_token_id
+
+    use_frame_special_tokens = getattr(model_config.mm_preprocessor, "use_frame_special_tokens", False)
+
+    molmo2_config = Molmo2Config(
+        vit_config=molmo2_vit_config,
+        adapter_config=molmo2_adapter_config,
+        text_config=molmo2_text_config,
+        image_start_token_id=image_start_token_id,
+        low_res_image_start_token_id=low_res_image_start_token_id,
+        image_end_token_id=image_end_token_id,
+        image_low_res_id=image_low_res_id,
+        image_patch_id=image_patch_id,
+        image_col_id=image_col_id,
+        frame_start_token_id=frame_start_token_id,
+        frame_end_token_id=frame_end_token_id,
+        use_frame_special_tokens=use_frame_special_tokens,
+        initializer_range=llm_config.initializer_range,
+        use_cache=True,
+        tie_word_embeddings=False,  # Always false for Molmo2
+    )
+    return molmo2_config
+
+
+def convert_lm_head_and_prefix(
+    state_dict: dict[str, Any],
+    base_model_prefix: str,
+    weight_tying: bool
+) -> dict[str, Any]:
+    new_state_dict = {}
+    for key, val in state_dict.items():
+        if key == "transformer.ff_out.weight":
+            new_key = "lm_head.weight"
+        else:
+            new_key = f"{base_model_prefix}.{key}"
+        new_state_dict[new_key] = val
+    
+    if weight_tying:
+        new_state_dict["lm_head.weight"] = state_dict["transformer.wte.embedding"]
+    
+    return new_state_dict
+
+
+def convert_molmo2(
+    state_dict: dict[str, Any],
+    config: Molmo2Config,
+    weight_tying: bool,
+) -> dict[str, Any]:
+    base_model_prefix = Molmo2ForConditionalGeneration.base_model_prefix
+    new_state_dict = convert_lm_head_and_prefix(state_dict, base_model_prefix, weight_tying)
+    model_prefix = f"{base_model_prefix}.transformer"
+    qkv_bias = config.qkv_bias if isinstance(config, Molmo2TextConfig) else config.text_config.qkv_bias
+    use_qk_norm = config.use_qk_norm if isinstance(config, Molmo2TextConfig) else config.text_config.use_qk_norm
+    for layer_i in range(config.num_hidden_layers):
+        prefix = f"{model_prefix}.blocks.{layer_i}"
+
+        move_to_attn = ["att_proj.weight", "attn_out.weight"]
+        if qkv_bias:
+            move_to_attn.append("att_proj.bias")
+        if use_qk_norm:
+            move_to_attn += ["q_norm.weight", "k_norm.weight"]
+        
+        for k in move_to_attn:
+            assert f"{prefix}.self_attn.{k}" not in new_state_dict
+            new_state_dict[f"{prefix}.self_attn.{k}"] = new_state_dict.pop(f"{prefix}.{k}")
+        
+        move_to_mlp = ["ff_proj.weight", "ff_out.weight"]
+        for k in move_to_mlp:
+            assert f"{prefix}.mlp.{k}" not in new_state_dict
+            new_state_dict[f"{prefix}.mlp.{k}"] = new_state_dict.pop(f"{prefix}.{k}")
+    
+    return new_state_dict
+
+
+def convert_model(
+    checkpoint_dir: str,
+    model_config: ModelConfig,
+    hf_config: Molmo2Config,
+    use_bfloat16: bool,
+) -> Molmo2ForConditionalGeneration:
+    """Convert model to HF-compatible model"""
+    with torch.device("meta"):
+        model = model_config.build_model()
+        hf_model = Molmo2ForConditionalGeneration(hf_config)
+    model.to_empty(device=torch.device("cpu"))
+    hf_model.to_empty(device=torch.device("cpu"))
+
+    load_model_state(checkpoint_dir, model)
+    model.eval()
+    model = model.to(torch.float32)
+    state_dict = model.state_dict()
+
+    new_state_dict = convert_molmo2(state_dict, hf_config, model_config.llm.weight_tying)
+    hf_model.eval()
+    hf_model = hf_model.to(torch.bfloat16 if use_bfloat16 else torch.float32)
+    hf_model.load_state_dict(new_state_dict)
+    return hf_model
+
+
+def save(
+    checkpoint_dir: str,
+    output_dir: str,
+    use_bfloat16: bool,
+    attn_implementation: str,
+    override_max_model_len: Optional[int],
+) -> None:
+    logger.info(f"Loading model config from {checkpoint_dir}")
+    config_path = resource_path(select_checkpoint(checkpoint_dir), "config.yaml")
+    model_config: ModelConfig = ModelConfig.load(config_path, key="model", validate_paths=False)
+
+    hf_config = convert_config(model_config, attn_implementation, override_max_model_len)
+
+    logger.info(f"Save HF-compatible model config and checkpoint to {output_dir}")
+    logger.info(f"Save HF-compatible model config and checkpoint to {output_dir}")
+    hf_model = convert_model(checkpoint_dir, model_config, hf_config, use_bfloat16)
+
+    hf_model.save_pretrained(output_dir)
+
+    gc.collect()
+
+    model_file = os.path.join(output_dir, "modeling_molmo2.py")
+    if not os.path.exists(model_file):
+        logger.warning(f"Copying model file to {model_file} manually")
+        shutil.copyfile(
+            "olmo/hf_model/modeling_molmo2.py",
+            model_file,
+        )
+
+    with open(os.path.join(output_dir, "config.json")) as f:
+        config = json.load(f)
+
+    auto_map = config.get("auto_map", None)
+    if auto_map is None:
+        auto_map = {}
+    if "AutoModelForImageTextToText" not in auto_map:
+        logger.warning("Add AutoModelForImageTextToText to auto_map")
+        auto_map["AutoModelForImageTextToText"] = "modeling_molmo2.Molmo2ForConditionalGeneration"
+        with open(os.path.join(output_dir, "config.json"), "w") as f:
+            json.dump(config, f, indent=2)
+
+    tokenizer = model_config.build_tokenizer().tokenizer
+    if not tokenizer.bos_token:
+        tokenizer.bos_token = tokenizer.eos_token
+        tokenizer.bos_token_id = tokenizer.eos_token_id
+    tokenizer.padding_side = "left"
+
+    tokenizer.chat_template = CHAT_TEMPLATE
+
+    logger.info(f"Save tokenizer and processor to {output_dir}")
+
+    mm_cfg = model_config.mm_preprocessor
+    vit_cfg = model_config.vision_backbone.vit
+
+    img_cfg = mm_cfg.image
+    video_cfg = mm_cfg.video
+
+    assert vit_cfg.resize_mode == "siglip", "Only siglip resize is supported for now"
+    assert vit_cfg.normalize == "siglip", "Only siglip normalization is supported for now"
+    assert img_cfg.crop_mode == "overlap-and-resize-c2", "Only overlap-and-resize-c2 crop mode is supported for now"
+    assert img_cfg.max_crops == img_cfg.max_multi_image_crops, "max_crops and max_multi_image_crops must be the same"
+    assert img_cfg.pooling_w == img_cfg.multi_image_pooling_w, "pooling_w and multi_image_pooling_w must be the same"
+    assert img_cfg.pooling_h == img_cfg.multi_image_pooling_h, "pooling_h and multi_image_pooling_h must be the same"
+
+    image_processor = Molmo2ImageProcessor(
+        size={"height": vit_cfg.image_default_input_size[0], "width": vit_cfg.image_default_input_size[1]},
+        resample=PILImageResampling.BILINEAR,
+        image_mean=IMAGENET_STANDARD_MEAN,
+        image_std=IMAGENET_STANDARD_STD,
+        do_convert_rgb=True,
+        max_crops=img_cfg.max_crops,
+        overlap_margins=img_cfg.overlap_margins,
+        patch_size=vit_cfg.image_patch_size,
+        pooling_size=[img_cfg.pooling_h, img_cfg.pooling_w],
+    )
+
+    image_use_col_tokens = img_cfg.use_col_tokens
+    use_single_crop_col_tokens = img_cfg.use_single_crop_col_tokens
+    use_single_crop_start_token = img_cfg.use_single_crop_start_token
+
+    assert vit_cfg.resize_mode == "siglip", "Only siglip resize is supported for now"
+    assert vit_cfg.normalize == "siglip", "Only siglip normalization is supported for now"
+    assert video_cfg.time_mode == "per-frame-compact", "Only per-frame-compact time mode is supported for now"
+
+    max_fps = video_cfg.max_fps
+    if isinstance(max_fps, (tuple, list)):
+        assert len(max_fps) == 1, "Only one max_fps is supported for now"
+        max_fps = max_fps[0]
+    video_processor = Molmo2VideoProcessor(
+        size={"height": vit_cfg.image_default_input_size[0], "width": vit_cfg.image_default_input_size[1]},
+        resample=PILImageResampling.BILINEAR,
+        image_mean=IMAGENET_STANDARD_MEAN,
+        image_std=IMAGENET_STANDARD_STD,
+        do_convert_rgb=True,
+        patch_size=vit_cfg.image_patch_size,
+        pooling_size=[video_cfg.pooling_h, video_cfg.pooling_w],
+        frame_sample_mode=video_cfg.frame_sample_mode,
+        num_frames=video_cfg.max_frames,
+        max_fps=max_fps,
+        sampling_fps=2,
+    )
+
+    video_use_col_tokens = False
+    use_frame_special_tokens = video_cfg.use_frame_special_tokens
+
+    processor = Molmo2Processor(
+        image_processor,
+        video_processor,
+        tokenizer,
+        chat_template=CHAT_TEMPLATE,
+        image_use_col_tokens=image_use_col_tokens,
+        use_single_crop_col_tokens=use_single_crop_col_tokens,
+        use_single_crop_start_token=use_single_crop_start_token,
+        video_use_col_tokens=video_use_col_tokens,
+        use_frame_special_tokens=use_frame_special_tokens,
+    )
+    processor.audio_tokenizer = None
+    processor.save_pretrained(output_dir)
+
+    logger.info(f"Save generation config to {output_dir}")
+    generation_config = GenerationConfig(
+        bos_token_id=tokenizer.bos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+    )
+    generation_config.save_pretrained(output_dir)
+
+    del hf_model, processor, tokenizer, generation_config
+    gc.collect()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert Molmo checkpoint to HuggingFace format."
+    )
+    parser.add_argument("checkpoint_dir", help="Location of Molmo2 checkpoint.")
+    parser.add_argument("output_dir", help="Location to save the converted checkpoint.", default="./hf-ckpt")
+    parser.add_argument("--use_bfloat16", action="store_true", help="Use bfloat16 weights")
+    parser.add_argument(
+        "--attn_implementation", type=str, default="sdpa", help="Attention type",
+        choices=["eager", "sdpa", "flash_attention_2"],
+    )
+    parser.add_argument(
+        "--override_max_model_len",
+        type=int,
+        default=None,
+        help="Override the max model length",
+    )
+    args = parser.parse_args()
+    prepare_cli_environment()
+
+    save(
+        args.checkpoint_dir,
+        args.output_dir,
+        args.use_bfloat16,
+        args.attn_implementation,
+        args.override_max_model_len,
+    )
+
+
+if __name__ == "__main__":
+    main()

convert_molmo_point_to_hf.py

@@ -0,0 +1,431 @@
+import argparse
+import os
+import shutil
+import logging
+import json
+import gc
+from typing import Dict, Any, Optional
+
+import torch
+
+from olmo.hf_model.processing_molmo2 import Molmo2Processor
+from olmo.tokenizer import EXTRA_TOKENS
+from transformers import GenerationConfig
+from transformers.image_utils import (
+    PILImageResampling,
+    IMAGENET_STANDARD_MEAN,
+    IMAGENET_STANDARD_STD,
+)
+
+from .convert_molmo2_to_hf import CHAT_TEMPLATE
+from .image_processing_molmo2 import Molmo2ImageProcessor
+from .video_processing_molmo2 import Molmo2VideoProcessor
+from olmo.models.molmo_point.molmo_point import MolmoPointConfig as OlmoMolmoPointConfig
+from olmo.train.checkpointer import load_model_state
+from olmo.util import (
+    prepare_cli_environment,
+    resource_path,
+    select_checkpoint
+)
+
+from .configuration_molmo_point import MolmoPointConfig
+from .configuration_molmo2 import Molmo2VitConfig, Molmo2TextConfig
+from .modeling_molmo_point import MolmoPointForConditionalGeneration, MolmoPointAdapterConfig
+
+
+N_POINT_TOKENS = 31200
+logger = logging.getLogger(__name__)
+
+
+def convert_config(
+    model_config: OlmoMolmoPointConfig,
+    attn_implementation: str,
+    override_max_model_len: Optional[int],
+) -> Molmo2Processor:
+    """Convert config to HF-compatible config"""
+    vit_config = model_config.vit
+    llm_config = model_config.llm
+
+    molmo2_vit_config = Molmo2VitConfig(
+        hidden_size=vit_config.image_emb_dim,
+        intermediate_size=vit_config.image_mlp_dim,
+        num_hidden_layers=vit_config.image_num_layers,
+        num_attention_heads=vit_config.image_num_heads,
+        num_key_value_heads=vit_config.image_num_key_value_heads,
+        head_dim=vit_config.image_head_dim,
+        hidden_act=vit_config.image_mlp_activations,
+        layer_norm_eps=vit_config.image_norm_eps,
+        image_default_input_size=vit_config.image_default_input_size,
+        image_patch_size=vit_config.image_patch_size,
+        image_num_pos=vit_config.image_num_pos,
+        attention_dropout=0.0,
+        residual_dropout=0.0,
+        initializer_range=vit_config.initializer_range,
+        float32_attention=vit_config.float32_attention,
+        attn_implementation=attn_implementation,
+    )
+    adapter_hidden_act = "silu" if llm_config.activation_type == "swiglu" else llm_config.activation_type
+    adapter_intermediate_size = (
+        llm_config.mlp_hidden_size if llm_config.mlp_hidden_size is not None
+        else llm_config.mlp_ratio * llm_config.d_model
+    ) // 2
+    connector = model_config.connector
+    molmo2_adapter_config = MolmoPointAdapterConfig(
+        vit_layers=connector.vit_layers,
+        pooling_attention_mask=connector.pooling_attention_mask,
+        hidden_size=vit_config.image_emb_dim,
+        num_attention_heads=vit_config.image_num_heads,
+        num_key_value_heads=vit_config.image_num_key_value_heads,
+        head_dim=vit_config.image_head_dim,
+        float32_attention=vit_config.float32_attention,
+        attention_dropout=0.0,
+        residual_dropout=0.0,
+        hidden_act=adapter_hidden_act,
+        intermediate_size=adapter_intermediate_size,
+        text_hidden_size=llm_config.d_model,
+        image_feature_dropout=0,
+        initializer_range=llm_config.initializer_range,
+        attn_implementation=attn_implementation,
+        positional_embeddings=connector.positional_embeddings
+    )
+    llm_head_dim = llm_config.d_model // llm_config.n_heads if llm_config.head_dim is None else llm_config.head_dim
+    llm_intermediate_size = (
+        llm_config.mlp_hidden_size if llm_config.mlp_hidden_size is not None
+        else llm_config.mlp_ratio * llm_config.d_model
+    ) // 2
+    llm_hidden_act = "silu" if llm_config.activation_type == "swiglu" else llm_config.activation_type
+    rope_scaling: Optional[Dict[str, Any]] = None
+    if llm_config.rope_type != "default":
+        rope_scaling = dict(rope_type=llm_config.rope_type)
+        for key in [
+            "rope_factor",
+            "rope_high_freq_factor",
+            "rope_low_freq_factor",
+            "rope_attention_factor",
+            "rope_original_max_position_embeddings",
+            "rope_beta_fast",
+            "rope_beta_slow",
+            "rope_mscale",
+            "rope_mscale_all_dim",
+            "rope_truncate",
+        ]:
+            if getattr(llm_config, key) is not None:
+                rope_scaling[key[len("rope_"):]] = getattr(llm_config, key)
+    
+    max_position_embeddings = llm_config.max_position_embeddings or llm_config.max_sequence_length
+    if override_max_model_len is not None:
+        max_position_embeddings = override_max_model_len
+    rope_scaling_layers: list[int] | None = None
+    if llm_config.full_attention_layers is not None:
+        # HACK: The original Olmo3 applies scaling to full attention layers,
+        # while we applies scaling to slinding attention layers.
+        if llm_config.sliding_attention_rope_scaling:
+            rope_scaling_layers = [idx for idx in range(llm_config.n_layers) if idx not in llm_config.full_attention_layers]
+        else:
+            rope_scaling_layers = list(llm_config.full_attention_layers)
+    molmo2_text_config = Molmo2TextConfig(
+        hidden_size=llm_config.d_model,
+        num_attention_heads=llm_config.n_heads,
+        num_key_value_heads=llm_config.effective_n_kv_heads,
+        head_dim=llm_head_dim,
+        vocab_size=llm_config.embedding_size or llm_config.vocab_size,
+        additional_vocab_size=llm_config.additional_vocab_size,
+        qkv_bias=llm_config.qkv_bias,
+        num_hidden_layers=llm_config.n_layers,
+        intermediate_size=llm_intermediate_size,
+        hidden_act=llm_hidden_act,
+        embedding_dropout=0.0,
+        attention_dropout=0.0,
+        residual_dropout=0.0,
+        max_position_embeddings=max_position_embeddings,
+        rope_theta=llm_config.rope_theta,
+        rope_scaling=rope_scaling,
+        rope_scaling_layers=rope_scaling_layers,
+        use_qk_norm=llm_config.attention_layer_norm,
+        qk_norm_type=llm_config.attention_layer_norm_type,
+        layer_norm_eps=llm_config.layer_norm_eps,
+        norm_after=llm_config.norm_after,
+        initializer_range=llm_config.initializer_range,
+        attn_implementation=attn_implementation,
+    )
+
+    tokenizer = model_config.build_tokenizer()
+    image_start_token_id = tokenizer.image_start_token_id
+    image_end_token_id = tokenizer.image_end_token_id
+    low_res_image_start_token_id = tokenizer.low_res_image_start_token_id
+    image_low_res_id = tokenizer.image_low_res_token_id
+    image_patch_id = tokenizer.image_patch_token_id
+    image_col_id = tokenizer.image_col_token_id
+    frame_start_token_id = tokenizer.frame_start_token_id
+    frame_end_token_id = tokenizer.frame_end_token_id
+
+    molmo2_config = MolmoPointConfig(
+        vit_config=molmo2_vit_config,
+        adapter_config=molmo2_adapter_config,
+        text_config=molmo2_text_config,
+        image_start_token_id=image_start_token_id,
+        image_end_token_id=image_end_token_id,
+        image_patch_id=image_patch_id,
+        image_col_id=image_col_id,
+        patch_token_id=tokenizer.token_index_token_id,
+        location_token_id=tokenizer.subpatch_loc_token_id,
+        subpatch_token_id=tokenizer.subpatch_index_token_id,
+        frame_start_token_id=frame_start_token_id,
+        frame_end_token_id=frame_end_token_id,
+        use_frame_special_tokens=model_config.mm_preprocessor.video.use_frame_special_tokens,
+        initializer_range=llm_config.initializer_range,
+        use_cache=True,
+        tie_word_embeddings=False,  # Always false for Molmo2
+
+        # Pointing configs
+        patch_location=model_config.patch_location,
+        no_more_points_class=model_config.no_more_points_class,
+        patch_embed_dim=model_config.patch_embed_dim,
+        patch_embedding_kind=model_config.patch_embedding_kind,
+        embed_selected_vit_patch=model_config.embed_selected_vit_patch,
+        embed_location=model_config.embed_location,
+        layer_norm_x=model_config.layer_norm_x,
+        mask_patches=model_config.mask_patches,
+        mask_subpatches=model_config.mask_subpatches,
+        mask_repeats=model_config.mask_repeats,
+        token_prediction_rotary=model_config.token_prediction_rotary,
+        token_prediction_rotary_theta=model_config.token_prediction_rotary_theta,
+    )
+    return molmo2_config
+
+
+def convert_molmo2(
+    state_dict: dict[str, Any],
+    config: MolmoPointConfig,
+    weight_tying: bool,
+) -> dict[str, Any]:
+    base_model_prefix = MolmoPointForConditionalGeneration.base_model_prefix
+    new_state_dict = {}
+    for key, val in state_dict.items():
+        if key == "transformer.ff_out.new_weight":
+            new_key = "new_output_embeddings"
+        elif key == "transformer.ff_out.weight":
+            new_key = "output_embeddings"
+        else:
+            new_key = f"{base_model_prefix}.{key}"
+        new_state_dict[new_key] = val
+
+    model_prefix = f"{base_model_prefix}.transformer"
+    qkv_bias = config.qkv_bias if isinstance(config, Molmo2TextConfig) else config.text_config.qkv_bias
+    use_qk_norm = config.use_qk_norm if isinstance(config, Molmo2TextConfig) else config.text_config.use_qk_norm
+
+    for param in list(new_state_dict.keys()):
+        if param.startswith(f"model.connectors.0"):
+            new_state_dict[param.replace("model.connectors.0", "model.connector")] = new_state_dict.pop(param)
+
+    for layer_i in range(config.num_hidden_layers):
+        prefix = f"{model_prefix}.blocks.{layer_i}"
+
+        move_to_attn = ["att_proj.weight", "attn_out.weight"]
+        if qkv_bias:
+            move_to_attn.append("att_proj.bias")
+        if use_qk_norm:
+            move_to_attn += ["q_norm.weight", "k_norm.weight"]
+        
+        for k in move_to_attn:
+            assert f"{prefix}.self_attn.{k}" not in new_state_dict
+            new_state_dict[f"{prefix}.self_attn.{k}"] = new_state_dict.pop(f"{prefix}.{k}")
+        
+        move_to_mlp = ["ff_proj.weight", "ff_out.weight"]
+        for k in move_to_mlp:
+            assert f"{prefix}.mlp.{k}" not in new_state_dict
+            new_state_dict[f"{prefix}.mlp.{k}"] = new_state_dict.pop(f"{prefix}.{k}")
+    
+    return new_state_dict
+
+
+def convert_model(
+    checkpoint_dir: str,
+    model_config: OlmoMolmoPointConfig,
+    hf_config: MolmoPointConfig,
+    use_bfloat16: bool,
+) -> MolmoPointForConditionalGeneration:
+    """Convert model to HF-compatible model"""
+    with torch.device("meta"):
+        model = model_config.build_model()
+        hf_model = MolmoPointForConditionalGeneration(hf_config)
+    model.to_empty(device=torch.device("cpu"))
+    hf_model.to_empty(device=torch.device("cpu"))
+
+    load_model_state(checkpoint_dir, model)
+    model.eval()
+    model = model.to(torch.float32)
+    state_dict = model.state_dict()
+
+    new_state_dict = convert_molmo2(state_dict, hf_config, model_config.llm.weight_tying)
+    hf_model.eval()
+    hf_model = hf_model.to(torch.bfloat16 if use_bfloat16 else torch.float32)
+    hf_model.load_state_dict(new_state_dict)
+    return hf_model
+
+
+def save(
+    checkpoint_dir: str,
+    output_dir: str,
+    use_bfloat16: bool,
+    attn_implementation: str,
+    override_max_model_len: Optional[int],
+) -> None:
+    logger.info(f"Loading model config from {checkpoint_dir}")
+    config_path = resource_path(select_checkpoint(checkpoint_dir), "config.yaml")
+    model_config: OlmoMolmoPointConfig = OlmoMolmoPointConfig.load(config_path, key="model", validate_paths=False)
+
+    hf_config = convert_config(model_config, attn_implementation, override_max_model_len)
+
+    logger.info(f"Save HF-compatible model config and checkpoint to {output_dir}")
+    logger.info(f"Save HF-compatible model config and checkpoint to {output_dir}")
+    hf_model = convert_model(checkpoint_dir, model_config, hf_config, use_bfloat16)
+
+    hf_model.save_pretrained(output_dir)
+
+    gc.collect()
+
+    model_file = os.path.join(output_dir, "modeling_molmo_point.py")
+    if not os.path.exists(model_file):
+        logger.warning(f"Copying model file to {model_file} manually")
+        shutil.copyfile(
+            "olmo/hf_model/modeling_molmo_point.py",
+            model_file,
+        )
+
+    with open(os.path.join(output_dir, "config.json")) as f:
+        config = json.load(f)
+
+    auto_map = config.get("auto_map", None)
+    if auto_map is None:
+        auto_map = {}
+    if "AutoModelForImageTextToText" not in auto_map:
+        logger.warning("Add AutoModelForImageTextToText to auto_map")
+        auto_map["AutoModelForImageTextToText"] = "modeling_molmo_point.MolmoPointForConditionalGeneration"
+        with open(os.path.join(output_dir, "config.json"), "w") as f:
+            json.dump(config, f, indent=2)
+
+    tokenizer = model_config.build_tokenizer().tokenizer
+    extra_tokens = [f"<EXTRA_TOKENS_POINT_{k}>" for k in range(model_config.llm.additional_vocab_size-len(EXTRA_TOKENS))]
+    extra_tokens += [f"<POINT_{k}>" for k in range(N_POINT_TOKENS)]
+    num_added = tokenizer.add_tokens(extra_tokens)
+    assert tokenizer.encode(f"<POINT_0>")[0] == model_config.llm.vocab_size + model_config.llm.additional_vocab_size
+    assert num_added == len(extra_tokens), "Failed to add extra tokens"
+
+    if not tokenizer.bos_token:
+        tokenizer.bos_token = tokenizer.eos_token
+        tokenizer.bos_token_id = tokenizer.eos_token_id
+    tokenizer.padding_side = "left"
+
+    tokenizer.chat_template = CHAT_TEMPLATE
+
+    logger.info(f"Save tokenizer and processor to {output_dir}")
+
+    mm_cfg = model_config.mm_preprocessor
+    vit_cfg = model_config.vit
+
+    img_cfg = mm_cfg.image
+
+    assert vit_cfg.resize_mode == "siglip", "Only siglip resize is supported for now"
+    assert vit_cfg.normalize == "siglip", "Only siglip normalization is supported for now"
+    assert img_cfg.crop_mode == "overlap-and-resize-c2", "Only overlap-and-resize-c2 crop mode is supported for now"
+    assert img_cfg.max_crops == img_cfg.max_multi_image_crops, "max_crops and max_multi_image_crops must be the same"
+    assert img_cfg.pooling_w == img_cfg.multi_image_pooling_w, "pooling_w and multi_image_pooling_w must be the same"
+    assert img_cfg.pooling_h == img_cfg.multi_image_pooling_h, "pooling_h and multi_image_pooling_h must be the same"
+
+    image_processor = Molmo2ImageProcessor(
+        size={"height": vit_cfg.image_default_input_size[0], "width": vit_cfg.image_default_input_size[1]},
+        resample=PILImageResampling.BILINEAR,
+        image_mean=IMAGENET_STANDARD_MEAN,
+        image_std=IMAGENET_STANDARD_STD,
+        do_convert_rgb=True,
+        max_crops=img_cfg.max_crops,
+        overlap_margins=img_cfg.overlap_margins,
+        patch_size=vit_cfg.image_patch_size,
+        pooling_size=[img_cfg.pooling_h, img_cfg.pooling_w],
+    )
+
+    image_use_col_tokens = img_cfg.use_col_tokens
+    use_single_crop_col_tokens = img_cfg.use_single_crop_col_tokens
+    use_single_crop_start_token = img_cfg.use_single_crop_start_token
+
+    assert vit_cfg.resize_mode == "siglip", "Only siglip resize is supported for now"
+    assert vit_cfg.normalize == "siglip", "Only siglip normalization is supported for now"
+
+    max_fps = mm_cfg.video.max_fps
+    if isinstance(max_fps, (tuple, list)):
+        assert len(max_fps) == 1, "Only one max_fps is supported for now"
+        max_fps = max_fps[0]
+    video_processor = Molmo2VideoProcessor(
+        size={"height": vit_cfg.image_default_input_size[0], "width": vit_cfg.image_default_input_size[1]},
+        resample=PILImageResampling.BILINEAR,
+        image_mean=IMAGENET_STANDARD_MEAN,
+        image_std=IMAGENET_STANDARD_STD,
+        do_convert_rgb=True,
+        patch_size=vit_cfg.image_patch_size,
+        pooling_size=[mm_cfg.video.pooling_h, mm_cfg.video.pooling_w],
+        frame_sample_mode=mm_cfg.video.frame_sample_mode,
+        num_frames=mm_cfg.video.max_frames,
+        max_fps=max_fps,
+        sampling_fps=2,
+    )
+
+    use_frame_special_tokens = mm_cfg.video.use_frame_special_tokens
+
+    processor = Molmo2Processor(
+        image_processor,
+        video_processor,
+        tokenizer,
+        chat_template=CHAT_TEMPLATE,
+        image_use_col_tokens=image_use_col_tokens,
+        use_single_crop_col_tokens=use_single_crop_col_tokens,
+        use_single_crop_start_token=use_single_crop_start_token,
+        video_use_col_tokens=False,
+        use_frame_special_tokens=use_frame_special_tokens,
+    )
+    processor.audio_tokenizer = None
+    processor.save_pretrained(output_dir)
+
+    logger.info(f"Save generation config to {output_dir}")
+    generation_config = GenerationConfig(
+        bos_token_id=tokenizer.bos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+    )
+    generation_config.save_pretrained(output_dir)
+
+    del hf_model, processor, tokenizer, generation_config
+    gc.collect()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert Molmo checkpoint to HuggingFace format."
+    )
+    parser.add_argument("checkpoint_dir", help="Location of Molmo2 checkpoint.")
+    parser.add_argument("output_dir", help="Location to save the converted checkpoint.", default="./hf-ckpt")
+    parser.add_argument("--use_bfloat16", action="store_true", help="Use bfloat16 weights")
+    parser.add_argument(
+        "--attn_implementation", type=str, default="sdpa", help="Attention type",
+        choices=["eager", "sdpa", "flash_attention_2"],
+    )
+    parser.add_argument(
+        "--override_max_model_len",
+        type=int,
+        default=None,
+        help="Override the max model length",
+    )
+    args = parser.parse_args()
+    prepare_cli_environment()
+
+    save(
+        args.checkpoint_dir,
+        args.output_dir,
+        args.use_bfloat16,
+        args.attn_implementation,
+        args.override_max_model_len,
+    )
+
+
+if __name__ == "__main__":
+    main()

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "bos_token_id": 151645,
+  "eos_token_id": 151645,
+  "pad_token_id": 151643,
+  "transformers_version": "4.57.6"
+}

image_processing_molmo2.py

@@ -0,0 +1,528 @@
+"""Image processor class for Molmo2"""
+from typing import Optional, Union
+import numpy as np
+import einops
+import torch
+import torchvision.transforms
+
+from transformers.image_utils import (
+    IMAGENET_STANDARD_MEAN,
+    IMAGENET_STANDARD_STD,
+    ImageInput,
+    PILImageResampling,
+    make_flat_list_of_images,
+    valid_images,
+    to_numpy_array,
+)
+from transformers.image_transforms import convert_to_rgb
+from transformers.processing_utils import ImagesKwargs
+from transformers.image_processing_utils import BaseImageProcessor, get_size_dict
+from transformers.utils import logging
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.utils import TensorType, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def normalize_image(
+    image: np.ndarray,
+    image_mean: list[float],
+    image_std: list[float],
+) -> np.ndarray:
+    image -= np.array(image_mean, dtype=np.float32)[None, None, :]
+    image /= np.array(image_std, dtype=np.float32)[None, None, :]
+    return image
+
+
+def resize_image(
+    image: np.ndarray,
+    desired_output_size: list[int],
+    resample: PILImageResampling,
+) -> np.ndarray:
+    image = torch.permute(torch.from_numpy(image), [2, 0, 1])
+    dtype = image.dtype
+    if torch.is_floating_point(image):
+        in_min = 0.0
+        in_max = 1.0
+        resized = torchvision.transforms.Resize(
+            desired_output_size,
+            resample,
+            antialias=False,
+        )(image)
+        resized = torch.clip(resized, 0.0, 1.0).to(dtype)
+    else:
+        assert image.dtype == torch.uint8, "SigLIP expects float images or uint8 images, but got {}".format(image.dtype)
+        in_min = 0.0
+        in_max = 255.0
+        resized = torchvision.transforms.Resize(
+            desired_output_size,
+            resample,
+            antialias=False,
+        )(image)
+        resized = torch.clip(resized, 0, 255).to(dtype)
+
+    resized = resized.to(torch.float32)
+    resized = (resized - in_min) / (in_max - in_min)
+
+    resized = torch.permute(resized, [1, 2, 0]).numpy()
+
+    return resized
+
+
+def select_tiling(h, w, patch_size, max_num_crops):
+    """Divide in image of size [w, h] in up to max_num_patches of size patch_size"""
+    original_size = np.stack([h, w])  # [1, 2]
+    original_res = h * w
+    tilings = []
+    for i in range(1, max_num_crops + 1):
+        for j in range(1, max_num_crops + 1):
+            if i*j <= max_num_crops:
+                tilings.append((i, j))
+    # sort so argmin and argmax favour smaller tilings in the event of a tie
+    tilings.sort(key=lambda x: (x[0]*x[1], x[0]))
+    candidate_tilings = np.array(tilings, dtype=np.int32)  # [n_resolutions, 2]
+    candidate_resolutions = candidate_tilings * patch_size  # [n_resolutions, 2]
+
+    # How much we would need to scale the image to fit exactly in each tiling
+    original_size = np.stack([h, w], dtype=np.float32)  # [1, 2]
+
+    # The original size can be zero in rare cases if the image is smaller than the margin
+    # In those cases letting the scale become infinite means the tiling is based on the
+    # other side, or falls back to the smallest tiling
+    with np.errstate(divide='ignore'):
+        required_scale_d = candidate_resolutions.astype(np.float32) / original_size,
+    required_scale = np.min(required_scale_d, axis=-1, keepdims=True)  # [n_resolutions, 1]
+    if np.all(required_scale < 1):
+        # We are forced to downscale, so try to minimize the amount of downscaling
+        ix = np.argmax(required_scale)
+    else:
+        # Pick the resolution that required the least upscaling so that it most closely fits the image
+        required_scale = np.where(required_scale < 1.0, 10e9, required_scale)
+        ix = np.argmin(required_scale)
+    return candidate_tilings[ix]
+
+
+def build_resized_image(
+    image: np.ndarray,
+    base_image_input_size: list[int],
+    resample: PILImageResampling,
+    image_mean: list[float],
+    image_std: list[float],
+    image_patch_size: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    resized = resize_image(
+        image, base_image_input_size, resample,
+    )
+    resized = normalize_image(resized, image_mean, image_std)
+    if len(resized.shape) == 3:
+        resized = np.expand_dims(resized, 0)
+    crop_patch_w = base_image_input_size[1] // image_patch_size
+    crop_patch_h = base_image_input_size[0] // image_patch_size
+    resize_idx = np.arange(crop_patch_w*crop_patch_h).reshape([crop_patch_h, crop_patch_w])
+    return resized, resize_idx
+
+
+def build_overlapping_crops(
+    image: np.ndarray,
+    max_crops: int,
+    overlap_margins: list[int],
+    base_image_input_size: list[int],
+    resample: PILImageResampling,
+    image_mean: list[float],
+    image_std: list[float],
+    image_patch_size: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Decompose an image into a set of overlapping crops
+
+    :return crop_arr: [n_crops, h, w, 3] The crops
+    :return patch_idx: [overlap_patch_h, overlap_patch_w] For each patch in the resized image
+                        the crops were extracted from, what patch in `crop_arr` it corresponds to
+    """
+    original_image_h, original_image_w = image.shape[:2]
+    crop_size = base_image_input_size[0]
+    assert base_image_input_size[0] == base_image_input_size[1]
+
+    left_margin, right_margin = overlap_margins
+    total_margin_pixels = image_patch_size * (right_margin + left_margin)  # pixels removed per dim
+    crop_patches = base_image_input_size[0] // image_patch_size  # patches per crop dim
+    crop_window_patches = crop_patches - (right_margin + left_margin)  # usable patches
+    crop_window_size = crop_window_patches * image_patch_size
+    crop_patch_w = base_image_input_size[1] // image_patch_size
+    crop_patch_h = base_image_input_size[0] // image_patch_size
+    original_image_h, original_image_w = image.shape[:2]
+    crop_size = base_image_input_size[0]
+
+    # Decide how to tile the image, to account for the overlap margins we compute the tiling
+    # as if we had an image without the margins and were using a crop size without the margins
+    tiling = select_tiling(
+        original_image_h - total_margin_pixels,
+        original_image_w - total_margin_pixels,
+        crop_window_size,
+        max_crops,
+    )
+
+    src = resize_image(
+        image,
+        [tiling[0]*crop_window_size+total_margin_pixels, tiling[1]*crop_window_size+total_margin_pixels],
+        resample,
+    )
+    src = normalize_image(src, image_mean, image_std)
+
+    # Now we have to split the image into crops, and track what patches came from
+    # where in `patch_idx_arr`
+    n_crops = tiling[0] * tiling[1]
+    crop_arr = np.zeros([n_crops, crop_size, crop_size, 3], dtype=src.dtype)
+    patch_idx_arr = np.zeros([n_crops, crop_patch_h, crop_patch_w], dtype=np.int32)
+    on_crop = 0
+    for i in range(tiling[0]):
+        # Slide over `src` by `crop_window_size` steps, but extract crops of size `crops_size`
+        # which results in overlapping crop windows
+        y0 = i*crop_window_size
+        for j in range(tiling[1]):
+            x0 = j*crop_window_size
+            crop_arr[on_crop] = src[y0:y0+crop_size, x0:x0+crop_size]
+            patch_idx = np.arange(crop_patch_w*crop_patch_h).reshape(crop_patch_h, crop_patch_w)
+            patch_idx += on_crop * crop_patch_h * crop_patch_w
+
+            # Mask out idx that are in the overlap region
+            if i != 0:
+                patch_idx[:left_margin, :] = -1
+            if j != 0:
+                patch_idx[:, :left_margin] = -1
+            if i != tiling[0]-1:
+                patch_idx[-right_margin:, :] = -1
+            if j != tiling[1]-1:
+                patch_idx[:, -right_margin:] = -1
+            patch_idx_arr[on_crop] = patch_idx
+            on_crop += 1
+
+    # `patch_idx_arr` is ordered crop-by-crop, here we transpose `patch_idx_arr`
+    # so it is ordered left-to-right order
+    patch_idx_arr = np.reshape(
+        patch_idx_arr,
+        [tiling[0], tiling[1], crop_patch_h, crop_patch_w]
+    )
+    patch_idx_arr = np.transpose(patch_idx_arr, [0, 2, 1, 3])
+    patch_idx_arr = np.reshape(patch_idx_arr, [-1])
+
+    # Now get the parts not in the overlap region, so it should map each patch in `src`
+    # to the correct patch it should come from in `crop_arr`
+    patch_idx_arr = patch_idx_arr[patch_idx_arr >= 0].reshape(
+        src.shape[0]//image_patch_size,
+        src.shape[1]//image_patch_size,
+    )
+    return crop_arr, patch_idx_arr
+
+
+def batch_pixels_to_patches(array: np.ndarray, patch_size: int) -> np.ndarray:
+    """Reshape images of [n_images, h, w, 3] -> [n_images, n_patches, pixels_per_patch]"""
+    if len(array.shape) == 3:
+        n_crops, h, w = array.shape
+        h_patches = h//patch_size
+        w_patches = w//patch_size
+        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size])
+        array = np.transpose(array, [0, 1, 3, 2, 4])
+        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size])
+        return array
+    else:
+        n_crops, h, w, c = array.shape
+        h_patches = h//patch_size
+        w_patches = w//patch_size
+        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size, c])
+        array = np.transpose(array, [0, 1, 3, 2, 4, 5])
+        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size*c])
+        return array
+
+
+def arange_for_pooling(
+    idx_arr: np.ndarray,
+    pool_h: int,
+    pool_w: int,
+) -> np.ndarray:
+    h_pad = pool_h * ((idx_arr.shape[0] + pool_h - 1) // pool_h) - idx_arr.shape[0]
+    w_pad = pool_w * ((idx_arr.shape[1] + pool_w - 1) // pool_w) - idx_arr.shape[1]
+    idx_arr = np.pad(idx_arr, [[h_pad//2, (h_pad+1)//2], [w_pad//2, (w_pad+1)//2]],
+                     mode='constant',constant_values=-1)
+    return einops.rearrange(
+        idx_arr, "(h dh) (w dw) -> h w (dh dw)", dh=pool_h, dw=pool_w)
+
+
+def image_to_patches_and_grids(
+    image: np.ndarray,
+    max_crops: int,
+    overlap_margins: list[int],
+    base_image_input_size: list[int],
+    resample: PILImageResampling,
+    image_mean: list[float],
+    image_std: list[float],
+    image_patch_size: int,
+    image_pooling_w: int,
+    image_pooling_h: int,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """
+    :return image_grids, the shape of each (low-res, high-res) image after pooling
+    :return crops, the image crops to processes with the ViT
+    :return pooled_patch_idx, for each patch_id tokens in `image_tokens`, the indices of the
+                                patches in `crops` to pool for that token, masked with -1
+    :rturn patch_idx_arr, map patch coordiantes to patch ids
+    """
+    if isinstance(base_image_input_size, int):
+        base_image_input_size = (base_image_input_size, base_image_input_size)
+    
+    base_image_input_d = image_patch_size
+    pooling_w = image_pooling_w
+    pooling_h = image_pooling_h
+    crop_patch_w = base_image_input_size[1] // base_image_input_d
+    crop_patch_h = base_image_input_size[0] // base_image_input_d
+
+    crop_arr, patch_idx_arr = build_overlapping_crops(
+        image,
+        max_crops,
+        overlap_margins,
+        base_image_input_size,
+        resample,
+        image_mean,
+        image_std,
+        image_patch_size,
+    )
+    pooling_idx = arange_for_pooling(patch_idx_arr, pooling_h, pooling_w)
+    h, w = pooling_idx.shape[:2]
+    pooling_idx = pooling_idx.reshape([-1, pooling_h*pooling_w])
+    
+    # Finally do the same for the global image
+    resized, resize_idx = build_resized_image(
+        image,
+        base_image_input_size,
+        resample,
+        image_mean,
+        image_std,
+        image_patch_size,
+    )
+    patch_idx_arr += crop_patch_h*crop_patch_w
+    crop_arr = np.concatenate([resized, crop_arr], 0)
+
+    resize_idx = arange_for_pooling(resize_idx, pooling_h, pooling_w)
+    resized_h, resized_w = resize_idx.shape[:2]
+    resize_idx = resize_idx.reshape([-1, pooling_h*pooling_w])
+
+    # Global image goes first, so the order of patches in previous crops gets increased
+    pooling_idx = np.where(
+        pooling_idx >= 0,
+        pooling_idx + crop_patch_h*crop_patch_w,
+        -1
+    )
+    pooling_idx = np.concatenate([resize_idx, pooling_idx])
+    image_grid = [np.array([resized_h, resized_w, h, w])]
+
+    return (
+        np.stack(image_grid, 0),
+        batch_pixels_to_patches(crop_arr, image_patch_size),
+        pooling_idx,
+        patch_idx_arr
+    )
+
+
+class Molmo2ImagesKwargs(ImagesKwargs, total=False):
+    max_crops: Optional[int]
+    overlap_margins: Optional[list[int]]
+    patch_size: Optional[int]
+    pooling_size: Optional[list[int]]
+
+
+class Molmo2ImageProcessor(BaseImageProcessor):
+    r"""
+    Constructs a Molmo2 image processor that preprocesses images for the model.
+
+    Args:
+        size (`dict[str, int]` *optional*, defaults to `{"height": 378, "width": 378}`):
+            Size of the image after resizing.
+        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
+            Resampling filter to use when resizing the image.
+        image_mean (`float` or `list[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
+            Mean to use if normalizing the image. This is a float or list of floats for each channel in the image.
+        image_std (`float` or `list[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
+            Standard deviation to use if normalizing the image. This is a float or list of floats for each channel in the image.
+        do_convert_rgb (`bool`, *optional*, defaults to `True`):
+            Whether to convert the image to RGB.
+        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 to use.
+        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.
+    """
+
+    model_input_names = ["pixel_values", "image_token_pooling", "image_grids", "image_num_crops"]
+
+    def __init__(
+        self,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = True,
+        max_crops: int = 8,
+        overlap_margins: list[int] = [4, 4],
+        patch_size: int = 14,
+        pooling_size: list[int] = [2, 2],
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 378, "width": 378}
+        size = get_size_dict(size, default_to_square=True)
+        self.size = size
+
+        self.resample = resample
+        self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
+        self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
+        self.do_convert_rgb = do_convert_rgb
+
+        self.max_crops = max_crops
+        self.overlap_margins = overlap_margins
+        self.patch_size = patch_size
+        self.pooling_size = pooling_size
+    
+    def preprocess(
+        self,
+        images: ImageInput,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        max_crops: Optional[int] = None,
+        overlap_margins: Optional[list[int]] = None,
+        patch_size: Optional[int] = None,
+        pooling_size: Optional[list[int]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_pointing_metadata: bool = False,
+        **kwargs,
+    ) -> BatchFeature:
+        """
+        Args:
+            images (`ImageInput`):
+                Image to preprocess.
+            size (`dict[str, int]`, *optional*, defaults to `self.size`):
+                Size of the image after resizing.
+            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
+                Resampling filter to use when resizing the image. This can be one of the enum `PILImageResampling`. Only
+                has an effect if `do_resize` is set to `True`.
+            image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):
+                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):
+                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
+                `True`.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+            max_crops (`int`, *optional*, defaults to `self.max_crops`):
+                Maximum number of crops to use per image.
+            overlap_margins (`list[int]`, *optional*, defaults to `self.overlap_margins`):
+                Overlap margins to use.
+            patch_size (`int`, *optional*, defaults to `self.patch_size`):
+                The spatial patch size of the vision encoder.
+            pooling_size (`list[int]`, *optional*, defaults to `self.pooling_size`):
+                The pooling size of the vision adapter.
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                - Unset: Return a list of `np.ndarray`.
+                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
+                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
+            return_patch_mappings (bool, optional):
+                Whether to return patch mappings used for decoding MolmoPoint points
+
+        Returns:
+            A `BatchFeature` containing the following keys:
+                - `pixel_values`: The preprocessed images.
+                - `image_token_pooling`: The indices of the patches in `crops` to pool for each token in `image_tokens`.
+                - `image_grids`: The image grids.
+                - `image_num_crops`: The number of crops for each image.
+        """
+        if size is not None:
+            if "height" not in size or "width" not in size:
+                raise ValueError("size must contain 'height' and 'width' keys.")
+        else:
+            size = {**self.size}
+        
+        base_image_input_size = [size["height"], size["width"]]
+        
+        resample = resample or self.resample
+        image_mean = image_mean or self.image_mean
+        image_std = image_std or self.image_std
+        do_convert_rgb = do_convert_rgb or self.do_convert_rgb
+
+        max_crops = max_crops or self.max_crops
+        overlap_margins = overlap_margins or self.overlap_margins
+        patch_size = patch_size or self.patch_size
+        pooling_size = pooling_size or self.pooling_size
+
+        image_pooling_h, image_pooling_w = pooling_size
+
+        if images is not None:
+            images = self.fetch_images(images)
+            images = make_flat_list_of_images(images)
+        
+        if images is not None and not valid_images(images):
+            raise ValueError(
+                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
+                "torch.Tensor, tf.Tensor or jax.ndarray."
+            )
+
+        if do_convert_rgb:
+            images = [convert_to_rgb(image) for image in images]
+
+        # All transformations expect numpy arrays.
+        images = [to_numpy_array(image) for image in images]
+
+        data = {}
+        patch_mappings = []
+        if images is not None:
+            batch_grids = []
+            batch_crops = []
+            batch_pooled_patches_idx = []
+            batch_num_crops = []
+
+            for image in images:
+                image_grid, crops, pooled_idx, patch_mapping = image_to_patches_and_grids(
+                    image,
+                    max_crops,
+                    overlap_margins,
+                    base_image_input_size,
+                    resample,
+                    image_mean,
+                    image_std,
+                    patch_size,
+                    image_pooling_w,
+                    image_pooling_h,
+                )
+                batch_grids.append(image_grid)
+                batch_crops.append(crops)
+                batch_pooled_patches_idx.append(pooled_idx)
+                batch_num_crops.append(crops.shape[0])
+                patch_mappings.append(patch_mapping)
+            
+            pixel_values = np.concatenate(batch_crops, 0)
+            image_token_pooling = np.concatenate(batch_pooled_patches_idx, 0)
+            image_grids = np.concatenate(batch_grids, 0)
+            image_num_crops = np.array(batch_num_crops)
+
+            data.update(
+                pixel_values=pixel_values,
+                image_token_pooling=image_token_pooling,
+                image_grids=image_grids,
+                image_num_crops=image_num_crops,
+            )
+
+        data = BatchFeature(data, tensor_type=return_tensors)
+        if return_pointing_metadata:
+            data["image_token_pooling_np"] = image_token_pooling if len(images) else None
+            data["subpatch_mapping"] = patch_mappings
+            data["image_sizes"] = [x.shape[:2][::-1] for x in images]
+        return data
+
+
+Molmo2ImageProcessor.register_for_auto_class()

model-00001-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f3e72190ef32e8a730642be569f673e3db747342d1233337b8839f805c02ce8
+size 4982833608

model-00002-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:396c71733dec4db53150621ba67aea1c8444e3cbdc06410427ee52c9b34090ff
+size 4798510440

model-00003-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a8d305edb19cd8bae64920bc893e636612e4722e986558338b097246b84efd6
+size 4630720272

model-00004-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e09b72e931382e8a0d7566ce791059f0d1e307358ff445aeab940083ae5b8ba2
+size 4630720320

model-00005-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d583e6ccfd62c1a710bd6405ce81ae965612fe814019004fd9357a4f2810f91
+size 4630720320

model-00006-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:957fd289c1ce380f606ac7c66fcf88fd071e59359e1824b34c56e040f118a302
+size 4630720320

model-00007-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f14ba921030350d6277d794066529298e9db125b763f109c17037611c5f379cb
+size 4997804128

model-00008-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3727a21ae90c0f0fc3e59f9a42354db0edafa6bc8577baed9b7c362446b15dd4
+size 1409480580

model.safetensors.index.json

@@ -0,0 +1,727 @@
+{
+  "metadata": {
+    "total_parameters": 8677855065,
+    "total_size": 34711420260
+  },
+  "weight_map": {
+    "model.add_no_point_class_embed.vector": "model-00008-of-00008.safetensors",
+    "model.build_vit_embedding.bias": "model-00008-of-00008.safetensors",
+    "model.build_vit_embedding.weight": "model-00008-of-00008.safetensors",
+    "model.connector.image_pooling_2d.wk.bias": "model-00008-of-00008.safetensors",
+    "model.connector.image_pooling_2d.wk.weight": "model-00008-of-00008.safetensors",
+    "model.connector.image_pooling_2d.wq.bias": "model-00008-of-00008.safetensors",
+    "model.connector.image_pooling_2d.wq.weight": "model-00008-of-00008.safetensors",
+    "model.connector.image_pooling_2d.wv.bias": "model-00008-of-00008.safetensors",
+    "model.connector.image_pooling_2d.wv.weight": "model-00008-of-00008.safetensors",
+    "model.connector.image_projector.w1.weight": "model-00008-of-00008.safetensors",
+    "model.connector.image_projector.w2.weight": "model-00008-of-00008.safetensors",
+    "model.connector.image_projector.w3.weight": "model-00008-of-00008.safetensors",
+    "model.patch_k.bias": "model-00008-of-00008.safetensors",
+    "model.patch_k.weight": "model-00008-of-00008.safetensors",
+    "model.patch_q.bias": "model-00008-of-00008.safetensors",
+    "model.patch_q.weight": "model-00008-of-00008.safetensors",
+    "model.subpatch_k.bias": "model-00008-of-00008.safetensors",
+    "model.subpatch_k.weight": "model-00008-of-00008.safetensors",
+    "model.subpatch_loc_k.bias": "model-00008-of-00008.safetensors",
+    "model.subpatch_loc_k.weight": "model-00008-of-00008.safetensors",
+    "model.subpatch_q.bias": "model-00008-of-00008.safetensors",
+    "model.subpatch_q.weight": "model-00008-of-00008.safetensors",
+    "model.transformer.blocks.0.attn_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.0.ff_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.0.mlp.ff_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.0.mlp.ff_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.0.self_attn.att_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.0.self_attn.attn_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.0.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.0.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.1.attn_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.1.ff_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.1.mlp.ff_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.1.mlp.ff_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.1.self_attn.att_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.1.self_attn.attn_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.1.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.1.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.10.attn_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.10.ff_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.10.mlp.ff_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.10.mlp.ff_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.10.self_attn.att_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.10.self_attn.attn_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.10.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.10.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.11.attn_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.11.ff_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.11.mlp.ff_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.11.mlp.ff_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.11.self_attn.att_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.11.self_attn.attn_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.11.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.11.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.12.attn_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.12.ff_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.12.mlp.ff_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.12.mlp.ff_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.12.self_attn.att_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.12.self_attn.attn_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.12.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.12.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.13.attn_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.13.ff_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.13.mlp.ff_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.13.mlp.ff_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.13.self_attn.att_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.13.self_attn.attn_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.13.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.13.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.14.attn_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.14.ff_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.14.mlp.ff_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.14.mlp.ff_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.14.self_attn.att_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.14.self_attn.attn_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.14.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.14.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.15.attn_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.15.ff_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.15.mlp.ff_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.15.mlp.ff_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.15.self_attn.att_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.15.self_attn.attn_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.15.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.15.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.16.attn_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.16.ff_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.16.mlp.ff_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.16.mlp.ff_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.16.self_attn.att_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.16.self_attn.attn_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.16.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.16.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.17.attn_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.17.ff_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.17.mlp.ff_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.17.mlp.ff_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.17.self_attn.att_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.17.self_attn.attn_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.17.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.17.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.18.attn_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.18.ff_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.18.mlp.ff_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.18.mlp.ff_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.18.self_attn.att_proj.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.18.self_attn.attn_out.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.18.self_attn.k_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.18.self_attn.q_norm.weight": "model-00004-of-00008.safetensors",
+    "model.transformer.blocks.19.attn_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.19.ff_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.19.mlp.ff_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.19.mlp.ff_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.19.self_attn.att_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.19.self_attn.attn_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.19.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.19.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.2.attn_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.2.ff_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.2.mlp.ff_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.2.mlp.ff_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.2.self_attn.att_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.2.self_attn.attn_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.2.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.2.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.20.attn_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.20.ff_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.20.mlp.ff_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.20.mlp.ff_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.20.self_attn.att_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.20.self_attn.attn_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.20.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.20.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.21.attn_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.21.ff_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.21.mlp.ff_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.21.mlp.ff_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.21.self_attn.att_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.21.self_attn.attn_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.21.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.21.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.22.attn_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.22.ff_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.22.mlp.ff_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.22.mlp.ff_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.22.self_attn.att_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.22.self_attn.attn_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.22.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.22.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.23.attn_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.23.ff_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.23.mlp.ff_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.23.mlp.ff_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.23.self_attn.att_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.23.self_attn.attn_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.23.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.23.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.24.attn_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.24.ff_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.24.mlp.ff_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.24.mlp.ff_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.24.self_attn.att_proj.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.24.self_attn.attn_out.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.24.self_attn.k_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.24.self_attn.q_norm.weight": "model-00005-of-00008.safetensors",
+    "model.transformer.blocks.25.attn_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.25.ff_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.25.mlp.ff_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.25.mlp.ff_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.25.self_attn.att_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.25.self_attn.attn_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.25.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.25.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.26.attn_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.26.ff_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.26.mlp.ff_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.26.mlp.ff_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.26.self_attn.att_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.26.self_attn.attn_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.26.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.26.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.27.attn_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.27.ff_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.27.mlp.ff_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.27.mlp.ff_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.27.self_attn.att_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.27.self_attn.attn_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.27.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.27.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.28.attn_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.28.ff_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.28.mlp.ff_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.28.mlp.ff_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.28.self_attn.att_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.28.self_attn.attn_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.28.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.28.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.29.attn_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.29.ff_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.29.mlp.ff_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.29.mlp.ff_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.29.self_attn.att_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.29.self_attn.attn_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.29.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.29.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.3.attn_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.3.ff_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.3.mlp.ff_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.3.mlp.ff_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.3.self_attn.att_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.3.self_attn.attn_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.3.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.3.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.30.attn_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.30.ff_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.30.mlp.ff_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.30.mlp.ff_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.30.self_attn.att_proj.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.30.self_attn.attn_out.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.30.self_attn.k_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.30.self_attn.q_norm.weight": "model-00006-of-00008.safetensors",
+    "model.transformer.blocks.31.attn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.31.ff_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.31.mlp.ff_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.31.mlp.ff_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.31.self_attn.att_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.31.self_attn.attn_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.31.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.31.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.32.attn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.32.ff_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.32.mlp.ff_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.32.mlp.ff_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.32.self_attn.att_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.32.self_attn.attn_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.32.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.32.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.33.attn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.33.ff_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.33.mlp.ff_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.33.mlp.ff_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.33.self_attn.att_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.33.self_attn.attn_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.33.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.33.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.34.attn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.34.ff_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.34.mlp.ff_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.34.mlp.ff_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.34.self_attn.att_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.34.self_attn.attn_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.34.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.34.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.35.attn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.35.ff_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.35.mlp.ff_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.35.mlp.ff_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.35.self_attn.att_proj.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.35.self_attn.attn_out.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.35.self_attn.k_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.35.self_attn.q_norm.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.blocks.4.attn_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.4.ff_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.4.mlp.ff_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.4.mlp.ff_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.4.self_attn.att_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.4.self_attn.attn_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.4.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.4.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.5.attn_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.5.ff_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.5.mlp.ff_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.5.mlp.ff_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.5.self_attn.att_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.5.self_attn.attn_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.5.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.5.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.6.attn_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.6.ff_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.6.mlp.ff_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.6.mlp.ff_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.6.self_attn.att_proj.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.6.self_attn.attn_out.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.6.self_attn.k_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.6.self_attn.q_norm.weight": "model-00002-of-00008.safetensors",
+    "model.transformer.blocks.7.attn_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.7.ff_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.7.mlp.ff_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.7.mlp.ff_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.7.self_attn.att_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.7.self_attn.attn_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.7.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.7.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.8.attn_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.8.ff_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.8.mlp.ff_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.8.mlp.ff_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.8.self_attn.att_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.8.self_attn.attn_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.8.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.8.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.9.attn_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.9.ff_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.9.mlp.ff_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.9.mlp.ff_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.9.self_attn.att_proj.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.9.self_attn.attn_out.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.9.self_attn.k_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.blocks.9.self_attn.q_norm.weight": "model-00003-of-00008.safetensors",
+    "model.transformer.ln_f.weight": "model-00007-of-00008.safetensors",
+    "model.transformer.wte.embedding": "model-00001-of-00008.safetensors",
+    "model.transformer.wte.new_embedding": "model-00001-of-00008.safetensors",
+    "model.vit.patch_embedding.bias": "model-00007-of-00008.safetensors",
+    "model.vit.patch_embedding.weight": "model-00007-of-00008.safetensors",
+    "model.vit.positional_embedding": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention.wk.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention.wk.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention.wo.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention.wo.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention.wq.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention.wq.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention.wv.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention.wv.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.attention_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.feed_forward.w1.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.feed_forward.w2.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.ffn_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.0.ffn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention.wk.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention.wk.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention.wo.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention.wo.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention.wq.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention.wq.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention.wv.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention.wv.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.attention_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.feed_forward.w1.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.feed_forward.w2.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.ffn_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.1.ffn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.10.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.11.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.12.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.13.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.14.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.15.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.16.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.17.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.18.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.19.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention.wk.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention.wk.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention.wo.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention.wo.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention.wq.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention.wq.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention.wv.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention.wv.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.attention_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.feed_forward.w1.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.feed_forward.w2.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.ffn_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.2.ffn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.20.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.21.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.22.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.23.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.24.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention.wk.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention.wk.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention.wo.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention.wo.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention.wq.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention.wq.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention.wv.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention.wv.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.attention_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.feed_forward.w1.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.feed_forward.w2.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.ffn_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.3.ffn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention.wk.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention.wk.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention.wo.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention.wo.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention.wq.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention.wq.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention.wv.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention.wv.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.attention_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.feed_forward.w1.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.feed_forward.w2.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.ffn_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.4.ffn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention.wk.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention.wk.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention.wo.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention.wo.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention.wq.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention.wq.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention.wv.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention.wv.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.attention_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.feed_forward.w1.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.feed_forward.w2.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.ffn_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.5.ffn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention.wk.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention.wk.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention.wo.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention.wo.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention.wq.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention.wq.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention.wv.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention.wv.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.attention_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.feed_forward.w1.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.feed_forward.w2.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.ffn_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.6.ffn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention.wk.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention.wk.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention.wo.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention.wo.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention.wq.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention.wq.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention.wv.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention.wv.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.attention_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.feed_forward.w1.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.feed_forward.w2.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.ffn_norm.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.7.ffn_norm.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention.wk.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention.wk.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention.wo.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention.wo.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention.wq.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention.wq.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention.wv.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention.wv.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.feed_forward.w1.bias": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.8.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention.wk.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention.wk.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention.wo.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention.wo.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention.wq.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention.wq.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention.wv.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention.wv.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.attention_norm.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.feed_forward.w1.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.feed_forward.w2.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.ffn_norm.bias": "model-00008-of-00008.safetensors",
+    "model.vit.transformer.resblocks.9.ffn_norm.weight": "model-00008-of-00008.safetensors",
+    "model.x_norm.weight": "model-00008-of-00008.safetensors",
+    "new_output_embeddings": "model-00001-of-00008.safetensors",
+    "output_embeddings": "model-00001-of-00008.safetensors"
+  }
+}

modeling_molmo2.py

@@ -0,0 +1,1764 @@
+import math
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Optional, Union, Callable
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from transformers.models.auto import AutoModelForImageTextToText
+from transformers.activations import ACT2FN
+from transformers.configuration_utils import PretrainedConfig
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask, create_masks_for_generate
+from transformers.modeling_flash_attention_utils import (
+    _flash_attention_forward,
+    FlashAttentionKwargs,
+    flash_attn_supports_top_left_mask,
+)
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+)
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    ModelOutput,
+    TransformersKwargs,
+    can_return_tuple,
+    logging,
+)
+
+from .configuration_molmo2 import Molmo2Config, Molmo2VitConfig, Molmo2AdapterConfig, Molmo2TextConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class Molmo2CausalLMOutputWithPast(ModelOutput):
+    """
+    Base class for Molmo2 causal language model (or autoregressive) outputs.
+
+    Args:
+        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: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+    image_hidden_states: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class Molmo2ModelOutputWithPast(BaseModelOutputWithPast):
+    """
+    Base class for Molmo2 outputs, with hidden states and attentions.
+
+    Args:
+        image_hidden_states (`torch.FloatTensor`, *optional*):
+            A `torch.FloatTensor` of size `(batch_num_patches, hidden_size)`.
+            image_hidden_states of the model produced by the vision backbone
+    """
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+    image_hidden_states: Optional[torch.FloatTensor] = None
+
+
+class ViTMLP(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, hidden_act: str, device: Union[str, torch.device] = None):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=True, device=device)
+        self.act = ACT2FN[hidden_act]
+        self.w2 = nn.Linear(hidden_dim, dim, bias=True, device=device)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(self.act(self.w1(x)))
+
+
+class ViTMultiHeadDotProductAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        use_bias: bool = True,
+        input_dim: Optional[int] = None,
+        float32_attention: bool = True,
+        attention_dropout: float = 0.0,
+        residual_dropout: float = 0.0,
+        device: Union[str, torch.device] = None,
+        attn_implementation: str = "eager",
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.attn_implementation = attn_implementation
+        self.is_causal = False
+
+        input_dim = input_dim or hidden_size
+
+        self.wq = nn.Linear(
+            input_dim,
+            self.num_heads * self.head_dim,
+            bias=use_bias,
+            device=device,
+        )
+        self.wk = nn.Linear(
+            input_dim,
+            self.num_key_value_heads * self.head_dim,
+            bias=use_bias,
+            device=device,
+        )
+        self.wv = nn.Linear(
+            input_dim,
+            self.num_key_value_heads * self.head_dim,
+            bias=use_bias,
+            device=device,
+        )
+        self.wo = nn.Linear(
+            self.num_heads * self.head_dim,
+            self.hidden_size,
+        )
+        self.float32_attention = float32_attention
+        self.attention_dropout = attention_dropout
+        self.residual_dropout = nn.Dropout(residual_dropout)
+
+    def _split_heads(self, hidden_states, num_heads) -> torch.Tensor:
+        return hidden_states.reshape(hidden_states.shape[:2] + (num_heads, self.head_dim))
+
+    def _merge_heads(self, hidden_states) -> torch.Tensor:
+        return hidden_states.reshape(hidden_states.shape[:2] + (self.hidden_size,))
+    
+    def forward(
+        self,
+        inputs_q: torch.Tensor,
+        inputs_kv: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        if inputs_kv is not None:
+            inputs_k = inputs_kv
+            inputs_v = inputs_kv
+        else:
+            inputs_k = inputs_q
+            inputs_v = inputs_q
+
+        xq, xk, xv = self.wq(inputs_q), self.wk(inputs_k), self.wv(inputs_v)
+
+        xq = self._split_heads(xq, self.num_heads)
+        xk = self._split_heads(xk, self.num_key_value_heads)
+        xv = self._split_heads(xv, self.num_key_value_heads)
+
+        if self.num_heads != self.num_key_value_heads:
+            xk = xk.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads)
+            xv = xv.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads)
+        
+        og_dtype = xq.dtype
+
+        if self.float32_attention:
+            xq = xq.to(torch.float)
+            xk = xk.to(torch.float)
+        
+        dropout_p = 0.0 if not self.training else self.attention_dropout
+        
+        if self.attn_implementation == "eager":
+            attn_weights = torch.einsum("...qhd,...khd->...hqk", xq / math.sqrt(xq.size(-1)), xk)
+            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(xq.dtype)
+            attn_weights = F.dropout(
+                attn_weights,
+                p=dropout_p,
+                training=self.training
+            )
+            attn_output = torch.einsum("...hqk,...khd->...qhd", attn_weights.to(xv.dtype), xv)
+        
+        elif self.attn_implementation == "sdpa":
+            if not torch.is_autocast_enabled():
+                xv = xv.to(torch.float)
+        
+            attn_output = F.scaled_dot_product_attention(
+                xq.transpose(1, 2).contiguous(),
+                xk.transpose(1, 2).contiguous(),
+                xv.transpose(1, 2).contiguous(),
+                attn_mask=attn_mask,
+                is_causal=False,
+                dropout_p=dropout_p,
+            ).transpose(1, 2)
+        
+        elif self.attn_implementation == "flash_attention_2":
+            if xq.dtype == torch.float32:
+                if torch.is_autocast_enabled():
+                    target_dtype = torch.get_autocast_gpu_dtype()
+                else:
+                    target_dtype = self.wq.weight.dtype
+            attn_output = _flash_attention_forward(
+                xq,
+                xk,
+                xv,
+                attention_mask=attn_mask,
+                query_length=inputs_q.shape[1],
+                is_causal=False,
+                dropout=dropout_p,
+                softmax_scale=xq.shape[-1] ** -0.5,
+                use_top_left_mask=flash_attn_supports_top_left_mask(),
+                target_dtype=target_dtype,
+                implementation=self.attn_implementation,
+            )
+        else:
+            raise ValueError(f"Attention implementation {self.attn_implementation} not supported")
+        
+        attn_output = attn_output.to(og_dtype)
+        attn_output = self._merge_heads(attn_output)
+        attn_output = self.wo(attn_output)
+        attn_output = self.residual_dropout(attn_output)
+
+        return attn_output
+
+
+class Molmo2VisionBlock(nn.Module):
+
+    def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None):
+        super().__init__()
+        self.attention = ViTMultiHeadDotProductAttention(
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_key_value_heads=config.num_key_value_heads,
+            head_dim=config.head_dim,
+            float32_attention=config.float32_attention,
+            attention_dropout=config.attention_dropout,
+            residual_dropout=config.residual_dropout,
+            device=device,
+            attn_implementation=config._attn_implementation,
+        )
+        self.feed_forward = ViTMLP(config.hidden_size, config.intermediate_size, config.hidden_act, device=device)
+        self.attention_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device)
+        self.ffn_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attention(self.attention_norm(x))
+        x = x + self.feed_forward(self.ffn_norm(x))
+        return x
+
+
+class Molmo2VisionBlockCollection(nn.Module):
+    
+    def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None):
+        super().__init__()
+        self.conifg = config
+        self.resblocks = nn.ModuleList([
+            Molmo2VisionBlock(config, device) for _ in range(config.num_hidden_layers)
+        ])
+
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        hidden_states = []
+        for r in self.resblocks:
+            x = r(x)
+            hidden_states.append(x)
+        return hidden_states
+
+
+class Molmo2VisionTransformer(nn.Module):
+
+    def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None):
+        super().__init__()
+        self.config = config
+
+        # positional embeddings
+        self.scale = config.hidden_size ** -0.5
+        self.num_prefix_tokens: int = 0 # no class embeddings
+        self.positional_embedding = nn.Parameter(
+            torch.zeros(config.image_num_pos, config.hidden_size, device=device),
+        )
+
+        image_patch_size = config.image_patch_size
+        self.patch_embedding = nn.Linear(
+            image_patch_size * image_patch_size * 3,
+            config.hidden_size,
+            bias=True,
+            device=device,
+        )
+
+        self.transformer = Molmo2VisionBlockCollection(config, device)
+
+    def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor:
+        pos_emb = self.positional_embedding
+
+        pos_emb = pos_emb.reshape(
+            (int(math.sqrt(pos_emb.shape[0])), int(math.sqrt(pos_emb.shape[0])), pos_emb.shape[1])
+        )
+
+        (patch_num_0, patch_num_1) = patch_num
+
+        if pos_emb.shape[0] != patch_num_0 or pos_emb.shape[1] != patch_num_1:
+            # Dervied from https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+            # antialias: default True in jax.image.resize
+            pos_emb = pos_emb.unsqueeze(0).permute(0, 3, 1, 2)
+            pos_emb = F.interpolate(
+                pos_emb, size=(patch_num_0, patch_num_1), mode="bicubic", align_corners=False, antialias=True,
+            )
+            pos_emb = pos_emb.permute(0, 2, 3, 1).squeeze(0)
+
+        pos_emb = pos_emb.reshape(-1, pos_emb.shape[-1])
+        x = x + pos_emb[None, :, :].to(x.dtype)
+        return x
+
+    def forward(self, x: torch.Tensor, patch_num: int = None) -> list[torch.Tensor]:
+        """
+        : param x: (batch_size, num_patch, n_pixels)
+        """
+        if patch_num is None:
+            patch_num = self.config.image_num_patch
+
+        B, N, D = x.shape
+
+        x = self.patch_embedding(x)
+
+        # class embeddings and positional embeddings
+        x = self.add_pos_emb(x, patch_num)
+
+        hidden_states = self.transformer(x)
+        return hidden_states
+
+
+class ImageProjectorMLP(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        hidden_act: str,
+        device: Union[str, torch.device] = None,
+    ):
+        super().__init__()
+        self.w1 = nn.Linear(input_dim, hidden_dim, bias=False, device=device)
+        self.w2 = nn.Linear(hidden_dim, output_dim, bias=False, device=device)
+        self.w3 = nn.Linear(input_dim, hidden_dim, bias=False, device=device)
+        self.act = ACT2FN[hidden_act]
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(self.act(self.w1(x)) * self.w3(x))
+
+
+class Molmo2VisionBackbone(nn.Module):
+    def __init__(self, vit_config: Molmo2VitConfig, adapter_config: Molmo2AdapterConfig):
+        super().__init__()
+        self.vit_config = vit_config
+        self.adapter_config = adapter_config
+
+        self.vit_layers = []
+        for layer in adapter_config.vit_layers:
+            if layer >= 0:
+                self.vit_layers.append(layer)
+            else:
+                self.vit_layers.append(layer + vit_config.num_hidden_layers)
+        
+        last_layer_needed = max(self.vit_layers) + 1
+        if last_layer_needed < vit_config.num_hidden_layers:
+            new_vit_config = deepcopy(vit_config)
+            new_vit_config.num_hidden_layers = last_layer_needed
+            self.image_vit = Molmo2VisionTransformer(new_vit_config)
+        else:
+            self.image_vit = Molmo2VisionTransformer(vit_config)
+
+        self.num_prefix_tokens: int = self.image_vit.num_prefix_tokens
+
+        pool_dim = vit_config.hidden_size * len(adapter_config.vit_layers)
+        self.image_pooling_2d = ViTMultiHeadDotProductAttention(
+            hidden_size=adapter_config.hidden_size,
+            num_heads=adapter_config.num_attention_heads,
+            num_key_value_heads=adapter_config.num_key_value_heads,
+            head_dim=adapter_config.head_dim,
+            input_dim=pool_dim,
+            float32_attention=adapter_config.float32_attention,
+            attention_dropout=adapter_config.attention_dropout,
+            residual_dropout=adapter_config.residual_dropout,
+            attn_implementation=adapter_config._attn_implementation,
+        )
+        self.image_projector = ImageProjectorMLP(
+            adapter_config.hidden_size,
+            adapter_config.intermediate_size,
+            adapter_config.text_hidden_size,
+            adapter_config.hidden_act,
+        )
+        self.image_feature_dropout = nn.Dropout(adapter_config.image_feature_dropout)
+    
+    def encode_image(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        : param images: (batch_size, num_crops, num_patch, n_pixels)
+        """
+        B, T, N, D = images.shape
+        images = images.view(B * T, N, D)
+        image_features = self.image_vit(images)
+
+        features = []
+        for layer in self.vit_layers:
+            features.append(image_features[layer])
+        image_features = torch.cat(features, dim=-1)
+
+        if self.num_prefix_tokens > 0:
+            image_features = image_features[:, 1:]
+        image_features = image_features.view(B, T, N, -1)
+        return image_features
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.image_vit.patch_embedding.weight.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.image_vit.patch_embedding.weight.device
+    
+    def forward(
+        self,
+        images: torch.Tensor,
+        pooled_patches_idx: torch.Tensor,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+        # image_features: (batch_size, num_crops(=num_image), num_patch, nximage_emb_dim)
+        batch_size, num_image = images.shape[:2]
+        images = images.to(device=self.device, dtype=self.dtype)
+        image_features = self.encode_image(images)
+
+        image_features = self.image_feature_dropout(image_features)
+        dim = image_features.shape[-1]
+        valid = pooled_patches_idx >= 0
+        valid_token = torch.any(valid, -1)
+
+        # Use `pooled_patches_idx` to arange the features for image pooling
+        batch_idx = torch.arange(pooled_patches_idx.shape[0], dtype=torch.long, device=pooled_patches_idx.device)
+        batch_idx = torch.tile(batch_idx.view(batch_size, 1, 1), [1, pooled_patches_idx.shape[1], pooled_patches_idx.shape[2]])
+
+        # Now [batch, num_high_res_features, pool_dim, dim]
+        to_pool = image_features.reshape(batch_size, -1, dim)[batch_idx, torch.clip(pooled_patches_idx, 0)]
+        to_pool = to_pool * valid.to(self.dtype)[:, :, :, None]
+        to_pool = to_pool.reshape([-1, pooled_patches_idx.shape[-1], dim])
+        if self.adapter_config.pooling_attention_mask:
+            attn_mask = valid.reshape([-1, 1, 1, valid.shape[-1]])
+            denom = valid.view(-1, to_pool.shape[-2]).float().sum(-1)
+            denom = torch.where(denom == 0, 1, denom)
+            query = to_pool.sum(-2, keepdim=True) / denom[:, None, None].to(to_pool.dtype)
+        else:
+            attn_mask = None
+            query = to_pool.mean(-2, keepdim=True)
+        pooled_features = self.image_pooling_2d(query, to_pool, attn_mask=attn_mask)
+        pooled_features = pooled_features.reshape([batch_size, -1, pooled_features.shape[-1]])
+
+        # MLP layer to map the feature.
+        pooled_features = self.image_projector(pooled_features)
+        return pooled_features.view(-1, pooled_features.shape[-1])[valid_token.flatten()]
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+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)
+
+
+# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        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
+
+
+class Molmo2RotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(
+        self,
+        config: Molmo2TextConfig,
+        device: Union[str, torch.device] = None,
+        rope_type: Optional[str] = None,
+    ):
+        super().__init__()
+        if rope_type is not None:
+            self.rope_type = rope_type
+        elif hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            # BC: "rope_type" was originally "type"
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+    
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            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)
+
+
+class Molmo2RMSNorm(nn.Module):
+
+    def __init__(
+        self,
+        size: int,
+        eps: float = 1e-6,
+        device: Union[str, torch.device] = None,
+    ):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(size, device=device))
+        self.eps = eps
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        with torch.autocast(enabled=False, device_type=x.device.type):
+            og_dtype = x.dtype
+            x = x.to(torch.float32)
+            variance = x.pow(2).mean(-1, keepdim=True)
+            x = x * torch.rsqrt(variance + self.eps)
+            x = x.to(og_dtype)
+        
+        return self.weight * x
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+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: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    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:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_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
+
+
+class Molmo2Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: Molmo2TextConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.is_causal = True
+
+        self.fused_dims = (
+            config.num_attention_heads * config.head_dim,
+            config.head_dim * config.num_key_value_heads,
+            config.head_dim * config.num_key_value_heads,
+        )
+        self.att_proj = nn.Linear(
+            config.hidden_size,
+            sum(self.fused_dims),
+            bias=config.qkv_bias,
+        )
+
+        # Layer norms.
+        self.k_norm: Optional[Molmo2RMSNorm] = None
+        self.q_norm: Optional[Molmo2RMSNorm] = None
+        self.qk_norm_type: Optional[str] = None
+        if config.use_qk_norm:
+            k_norm_size = (
+                config.head_dim
+                if config.qk_norm_type == "qwen3" else
+                config.num_key_value_heads * config.head_dim
+            )
+            self.k_norm = Molmo2RMSNorm(k_norm_size, eps=config.layer_norm_eps)
+            q_norm_size = (
+                config.head_dim
+                if config.qk_norm_type == "qwen3" else
+                config.num_attention_heads * config.head_dim
+            )
+            self.q_norm = Molmo2RMSNorm(q_norm_size, eps=config.layer_norm_eps)
+            self.qk_norm_type = config.qk_norm_type
+
+        self.attention_dropout = config.attention_dropout
+        
+        self.attn_out = nn.Linear(
+            config.head_dim * config.num_attention_heads,
+            config.hidden_size,
+            bias=False,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        qkv = self.att_proj(hidden_states)
+        query_states, key_states, value_states = qkv.split(self.fused_dims, dim=-1)
+        value_states = value_states.view(hidden_shape)
+
+        # Optionally apply layer norm to keys and queries.
+        if self.q_norm is not None and self.k_norm is not None and self.qk_norm_type != "qwen3":
+            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.q_norm is not None and self.k_norm is not None and 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)
+        value_states = value_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: 
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        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.attn_out(attn_output)
+        return attn_output, attn_weights
+
+
+class LanguageModelMLP(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        intermediate_size: int,
+        hidden_act: str,
+        device: Union[str, torch.device] = None,
+    ):
+        super().__init__()
+        self.ff_proj = nn.Linear(input_dim, intermediate_size * 2, bias=False, device=device)
+        self.ff_out = nn.Linear(intermediate_size, input_dim, bias=False, device=device)
+        self.act = ACT2FN[hidden_act]
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.ff_proj(x)
+        x, gate = x.chunk(2, dim=-1)
+        x = self.act(gate) * x
+        x = self.ff_out(x)
+        return x
+
+
+class Molmo2DecoderLayer(GradientCheckpointingLayer):
+
+    def __init__(
+        self,
+        config: Molmo2TextConfig,
+        layer_idx: Optional[int] = None,
+        device: Union[str, torch.device] = None
+    ):
+        super().__init__()
+        self.config = config
+
+        self.self_attn = Molmo2Attention(config, layer_idx)
+        self.attn_norm = Molmo2RMSNorm(
+            config.hidden_size, eps=config.layer_norm_eps, device=device)
+        self.dropout = nn.Dropout(config.residual_dropout)
+        self.mlp = LanguageModelMLP(
+            config.hidden_size, config.intermediate_size, config.hidden_act, device=device)
+        self.ff_norm = Molmo2RMSNorm(
+            config.hidden_size, eps=config.layer_norm_eps, device=device)
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+
+        residual = hidden_states
+        hidden_states = self.attn_norm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights = 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,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = residual + self.dropout(hidden_states)
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.ff_norm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+
+        hidden_states = residual + self.dropout(hidden_states)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs
+
+
+class Molmo2PostNormDecoderLayer(Molmo2DecoderLayer):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+
+        residual = hidden_states
+
+        # Self Attention
+        hidden_states, self_attn_weights = 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,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+        )
+        hidden_states = self.attn_norm(hidden_states)
+
+        hidden_states = residual + self.dropout(hidden_states)
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.ff_norm(hidden_states)
+
+        hidden_states = residual + self.dropout(hidden_states)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs
+
+
+class Molmo2Embedding(nn.Module):
+    def __init__(
+        self,
+        num_embeddings: int,
+        num_new_embeddings: int,
+        features: int,
+        device: Union[str, torch.device] = None,
+    ):
+        super().__init__()
+        self.embedding = nn.Parameter(
+            torch.zeros(num_embeddings, features, device=device),
+        )
+        self.new_embedding = nn.Parameter(
+            torch.zeros(num_new_embeddings, features, device=device),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.embedding(x, torch.cat([self.embedding, self.new_embedding], dim=0))
+
+
+class Molmo2PreTrainedModel(PreTrainedModel):
+    config: Molmo2Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = [
+        "Molmo2DecoderLayer",
+        "Molmo2PostNormDecoderLayer",
+        "Molmo2VisionBlock",
+        "ViTMultiHeadDotProductAttention",
+    ]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": Molmo2DecoderLayer,
+        "attentions": Molmo2Attention,
+    }
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, (nn.Linear,)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, Molmo2Embedding):
+            module.embedding.data.normal_(mean=0.0, std=std)
+            module.new_embedding.data.normal_(mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, Molmo2RMSNorm):
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+
+class Molmo2TextModel(Molmo2PreTrainedModel):
+    config: Molmo2TextConfig
+    _no_split_modules = ["Molmo2DecoderLayer", "Molmo2PostNormDecoderLayer"]
+
+    def __init__(self, config: Molmo2TextConfig):
+        super().__init__(config)
+        if config.additional_vocab_size is not None:
+            self.wte = Molmo2Embedding(
+                config.vocab_size,
+                config.additional_vocab_size,
+                config.hidden_size,
+            )
+        else:
+            self.wte = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.emb_drop = nn.Dropout(config.embedding_dropout)
+        decoder_layer = Molmo2PostNormDecoderLayer if config.norm_after else Molmo2DecoderLayer
+        self.blocks = nn.ModuleList(
+            [decoder_layer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.ln_f = Molmo2RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+        if config.rope_scaling_layers is not None:
+            self.rotary_embs = nn.ModuleDict(
+                {
+                    "default": Molmo2RotaryEmbedding(config, rope_type="default"),
+                    "scaling": Molmo2RotaryEmbedding(config),
+                }
+            )
+        else:
+            self.rotary_emb = Molmo2RotaryEmbedding(config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.wte
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.wte = value
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+        
+        if inputs_embeds is None:
+            input_ids = input_ids * (input_ids != -1).to(input_ids.dtype)
+            inputs_embeds = self.wte(input_ids)
+
+        # torch.jit.trace() doesn't support cache objects in the output
+        if use_cache and past_key_values is None and not torch.jit.is_tracing():
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+
+            # Create the mask
+            causal_mask_mapping = create_causal_mask(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        if self.config.rope_scaling_layers is not None:
+            position_embeddings_mapping = {
+                "default": self.rotary_embs["default"](hidden_states, position_ids),
+                "scaling": self.rotary_embs["scaling"](hidden_states, position_ids),
+            }
+        else:
+            position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for layer_idx, decoder_block in enumerate(self.blocks[: self.config.num_hidden_layers]):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            
+            if self.config.rope_scaling_layers is not None:
+                position_embeddings_i = (
+                    position_embeddings_mapping["scaling"]
+                    if layer_idx in self.config.rope_scaling_layers
+                    else position_embeddings_mapping["default"]
+                )
+            else:
+                position_embeddings_i = position_embeddings
+
+            layer_outputs = decoder_block(
+                hidden_states,
+                attention_mask=causal_mask_mapping,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings_i,
+                **kwargs,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.ln_f(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+# Adapted from transformers.models.gemma3.modeling_gemma3
+def token_type_ids_mask_function(
+    token_type_ids: Optional[torch.Tensor] = None,
+) -> Optional[Callable]:
+    """
+    This function adds the correct offsets to the `q_idx` and `kv_idx` as the torch API can only accept lengths,
+    not start and end indices.
+    """
+    # Do not return an additional mask in this case
+    if token_type_ids is None:
+        return None
+    
+    def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
+        # If it's 1 for both query and key/value, we are in an image block
+        # NOTE: static cache shape goes beyond input seq length, while token_type_ids.shape[1] == input seq length
+        # Since vmap doesn't support `if statement` we workaround it with `torch.where`
+        safe_idx = torch.where(kv_idx < token_type_ids.shape[1], kv_idx, 0)
+        token_type_ids_at_kv_idx = token_type_ids[batch_idx, safe_idx]
+        token_type_ids_at_kv_idx = torch.where(kv_idx < token_type_ids.shape[1], token_type_ids_at_kv_idx, 0)
+
+        is_image_block = (token_type_ids[batch_idx, q_idx] == 1) & (token_type_ids_at_kv_idx == 1)
+
+        # This is bidirectional attention whenever we are dealing with image tokens
+        return is_image_block & is_image_block
+    
+    return inner_mask
+
+
+class Molmo2Model(Molmo2PreTrainedModel):
+    base_model_prefix = ""
+    _checkpoint_conversion_mapping = {}
+    # Reference: fix gemma3 grad acc #37208
+    accepts_loss_kwargs = False
+    config: Molmo2Config
+
+
+    def __init__(self, config: Molmo2Config):
+        super().__init__(config)
+        self.transformer: Molmo2TextModel = Molmo2TextModel(config.text_config)
+        self.vision_backbone: Optional[Molmo2VisionBackbone] = None
+        if config.vit_config is not None and config.adapter_config is not None:
+            self.vision_backbone = Molmo2VisionBackbone(config.vit_config, config.adapter_config)
+        
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.transformer.wte
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.transformer.wte = value
+    
+    def set_decoder(self, decoder):
+        self.transformer = decoder
+    
+    def get_decoder(self):
+        return self.transformer
+
+    @property
+    def device(self) -> torch.device:
+        return self.transformer.ln_f.weight.device
+    
+    def build_batched_images(
+        self,
+        input_ids: torch.LongTensor,
+        pixel_values: torch.Tensor,
+        image_token_pooling: torch.Tensor,
+        image_grids: torch.Tensor,
+        image_num_crops: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # 1) Count the number of images in each example
+        raw_counts = (input_ids == self.config.image_end_token_id).sum(1)  # [N]
+        # Each image is represented by global view and high-res view
+        # so we divide by 2 to get the number of images
+        counts = raw_counts // 2
+        N = counts.size(0)
+        device = input_ids.device
+
+        # Total number of images in the batch
+        num_images = int(counts.sum().item())
+
+        # Sanity check
+        assert image_grids.size(0) == num_images, \
+            f"Expected {num_images} image grids, but got {image_grids.size(0)}"
+        assert image_num_crops.size(0) == num_images, \
+            f"Expected {num_images} image num crops, but got {image_num_crops.size(0)}"
+
+        # 1-1) Compute per-image pooled patch count from image grids
+        with torch.no_grad():
+            first_prod = image_grids[:, :2].prod(dim=1)    # [num_images]
+            second_prod = image_grids[:, 2:].prod(dim=1)   # [num_images]
+            num_pooled_patches_per_image = (first_prod + second_prod).to(image_num_crops.dtype)  # [num_images]
+        
+        # pixel_values: [n_crops, n_patches, pixels_per_patch]
+        n_crops, n_patches, pixels_per_patch = pixel_values.shape
+        
+        # 2) Map each image index → example index
+        # Example: if counts = [2, 1, 3], then this becomes [0,0,1,2,2,2]
+        example_ids_for_image = torch.arange(N, device=device).repeat_interleave(counts)  # [num_images]
+        assert example_ids_for_image.numel() == num_images
+
+        # 2-1) Compute crops_per_example by summing per-image crop counts
+        crops_per_example = torch.zeros(
+            N, dtype=image_num_crops.dtype, device=image_num_crops.device
+        )
+        crops_per_example.index_add_(0, example_ids_for_image, image_num_crops)  # [N]
+
+        # 2-2) Per-image number of patches = (crops per image) * n_patches
+        patches_per_image = image_num_crops * n_patches  # [num_images]
+
+        # 2-3) Compute per-example per-image patch offsets
+        counts_list = counts.tolist()
+        index_offset_per_example_list = []
+        offset_img = 0
+        for c in counts_list:
+            per_img_patches = patches_per_image[offset_img:offset_img + c]  # [c]
+            # Offsets: [0, img0_total_patches, img0+img1_total_patches, ...]
+            index_offset = [0] + per_img_patches.cumsum(0).tolist()[:-1]
+            index_offset_per_example_list.append(index_offset)
+            offset_img += c
+        
+        # 2-4) Compute num_pooled_patches_per_example
+        num_pooled_patches_per_example = torch.zeros(
+            N, dtype=num_pooled_patches_per_image.dtype, device=num_pooled_patches_per_image.device
+        )
+        num_pooled_patches_per_example.index_add_(
+            0, example_ids_for_image, num_pooled_patches_per_image
+        )
+
+        # Sanity checks
+        total_crops = int(crops_per_example.sum().item())
+        assert total_crops == n_crops, \
+            f"Expected {total_crops} crops, but got {n_crops}"
+
+        total_num_pooled_patches = int(num_pooled_patches_per_example.sum().item())
+        assert total_num_pooled_patches == image_token_pooling.size(0), \
+            f"Expected {total_num_pooled_patches} pooled patches, but got {image_token_pooling.size(0)}"
+
+        # 3) Build images tensor filled with -1
+        M = int(crops_per_example.max().item())
+        images = torch.full(
+            (N, M, n_patches, pixels_per_patch),
+            fill_value=-1,
+            dtype=pixel_values.dtype,
+            device=pixel_values.device,
+        )
+
+        # 4) Fill images with per-example slices from pixel_values
+        offset_crop = 0
+        for i in range(N):
+            num = int(crops_per_example[i].item())
+            cur = pixel_values[offset_crop:offset_crop + num]  # [num, n_patches, pixels_per_patch]
+            images[i, :num] = cur
+            offset_crop += num
+
+        # Sanity check
+        assert offset_crop == n_crops
+
+        # 5) Build new_token_pooling tensor filled with -1
+        P = int(num_pooled_patches_per_example.max().item())
+        _, dim = image_token_pooling.shape
+        new_token_pooling = torch.full(
+            (N, P, dim),
+            fill_value=-1,
+            dtype=image_token_pooling.dtype,
+            device=image_token_pooling.device,
+        )
+
+        # 6) Fill token_pooling with per-example slices, adding per-image patch offsets
+        patch_offset = 0
+        img_offset = 0
+
+        for i, c in enumerate(counts_list):
+            num_patches = int(num_pooled_patches_per_example[i].item())
+
+            # Subsequence of pooled tokens belonging to this example
+            cur = image_token_pooling[patch_offset:patch_offset + num_patches].clone()  # [num_patches, dim]
+
+            index_offset_per_example = index_offset_per_example_list[i]  # length = c
+            per_img_pooled = num_pooled_patches_per_image[img_offset:img_offset + c]   # [c]
+
+            assert len(index_offset_per_example) == per_img_pooled.numel()
+
+            # Apply per-image offsets to the (ragged) subsequence
+            offset = 0
+            for j in range(c):
+                index_offset = int(index_offset_per_example[j])
+                n = int(per_img_pooled[j].item())
+                cur_slice = cur[offset:offset + n]
+
+                # Apply offset across all columns
+                cur[offset:offset + n] = torch.where(
+                    cur_slice >= 0,
+                    cur_slice + index_offset,
+                    cur_slice,
+                )
+                offset += n
+
+            new_token_pooling[i, :num_patches] = cur
+
+            patch_offset += num_patches
+            img_offset += c
+
+        # Final sanity checks
+        assert patch_offset == total_num_pooled_patches
+        assert img_offset == num_images
+
+        return images, new_token_pooling
+    
+    def build_batched_videos(
+        self,
+        input_ids: torch.LongTensor,
+        pixel_values_videos: torch.Tensor,
+        video_token_pooling: torch.Tensor,
+        video_grids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        # 1) Count the number of videos in each example
+        if self.config.use_frame_special_tokens:
+            end_token_id = self.config.frame_end_token_id
+        else:
+            end_token_id = self.config.image_end_token_id
+        counts = (input_ids == end_token_id).any(dim=1).long()  # [N]
+        N = counts.size(0)
+        device = input_ids.device
+
+        # Total number of videos in the batch
+        num_videos = int(counts.sum().item())
+
+        # Sanity check
+        assert video_grids.size(0) == num_videos, \
+            f"Expected {num_videos} videos, but got {video_grids.size(0)}"
+        
+        video_num_frames = video_grids[:, 0]  # [num_videos]
+        num_pooled_patches_per_video = video_grids.prod(dim=1)  # [num_videos]
+
+        # pixel_values_videos: [n_frames, n_patches, pixels_per_patch]
+        n_frames, n_patches, pixels_per_patch = pixel_values_videos.shape
+
+        # 2) Map each video index -> example index
+        # Example: if counts = [2, 1, 3], then this becomes [0,0,1,2,2,2]
+        example_ids_for_video = torch.arange(N, device=device).repeat_interleave(counts)  # [num_videos]
+        assert example_ids_for_video.numel() == num_videos
+
+        # 2-1) Compute frames_per_example by summing per-video frame counts
+        frames_per_example = torch.zeros(
+            N, dtype=video_num_frames.dtype, device=device,
+        )
+        frames_per_example.index_add_(0, example_ids_for_video, video_num_frames)  # [N]
+
+        # 2-2) Compute num_pooled_patches_per_example
+        num_pooled_patches_per_example = torch.zeros(
+            N, dtype=num_pooled_patches_per_video.dtype, device=num_pooled_patches_per_video.device,
+        )
+        num_pooled_patches_per_example.index_add_(
+            0, example_ids_for_video, num_pooled_patches_per_video,
+        )
+
+        # Sanity checks
+        total_frames = int(frames_per_example.sum().item())
+        assert total_frames == n_frames, \
+            f"Expected {total_frames} frames, but got {n_frames}"
+        
+        total_num_pooled_patches = int(num_pooled_patches_per_example.sum().item())
+        assert total_num_pooled_patches == video_token_pooling.size(0), \
+            f"Expected {total_num_pooled_patches} pooled patches, but got {video_token_pooling.size(0)}"
+        
+        # 3) Build videos tensor filled with -1
+        M = int(frames_per_example.max().item())
+        videos = torch.full(
+            (N, M, n_patches, pixels_per_patch),
+            fill_value=-1,
+            dtype=pixel_values_videos.dtype,
+            device=device,
+        )
+
+        # 4) Fill videos with per-examples slices from pixel_values_videos
+        offset_frame = 0
+        for i in range(N):
+            num = int(frames_per_example[i].item())
+            cur = pixel_values_videos[offset_frame:offset_frame + num]  # [num, n_patches, pixels_per_patch]
+            videos[i, :num] = cur
+            offset_frame += num
+        
+        # Sanity check
+        assert offset_frame == n_frames
+
+        # 5) Build new token_pooling tensor filled with -1
+        P = int(num_pooled_patches_per_example.max().item())
+        _, dim = video_token_pooling.shape
+        new_token_pooling = torch.full(
+            (N, P, dim),
+            fill_value=-1,
+            dtype=video_token_pooling.dtype,
+            device=video_token_pooling.device,
+        )
+
+        # 6) Fill new token_pooling with per-examples slices from video_token_pooling
+        patch_offset = 0
+        for i in range(N):
+            num_patches = int(num_pooled_patches_per_example[i].item())
+            cur = video_token_pooling[patch_offset:patch_offset + num_patches]  # [num_patches, dim]
+            new_token_pooling[i, :num_patches] = cur
+            patch_offset += num_patches
+
+        # Final sanity checks
+        assert patch_offset == total_num_pooled_patches
+
+        return videos, new_token_pooling
+    
+    def merge_visual_inputs(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        image_token_pooling: Optional[torch.Tensor] = None,
+        image_grids: Optional[torch.Tensor] = None,
+        image_num_crops: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.Tensor] = None,
+        video_token_pooling: Optional[torch.Tensor] = None,
+        video_grids: Optional[torch.Tensor] = None,
+    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        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")
+        elif pixel_values is not None:
+            assert input_ids is not None
+            images, token_pooling = self.build_batched_images(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                image_token_pooling=image_token_pooling,
+                image_grids=image_grids,
+                image_num_crops=image_num_crops,
+            )
+        elif pixel_values_videos is not None:
+            assert input_ids is not None
+            images, token_pooling = self.build_batched_videos(
+                input_ids=input_ids,
+                pixel_values_videos=pixel_values_videos,
+                video_token_pooling=video_token_pooling,
+                video_grids=video_grids,
+            )
+        else:
+            images, token_pooling = None, None
+        return images, token_pooling
+
+    def build_input_embeddings(
+        self,
+        input_ids: torch.LongTensor,
+        images: Optional[torch.FloatTensor] = None,  # image inputs
+        token_pooling: Optional[torch.LongTensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+        # Get embeddings of input.
+        # shape: (batch_size, seq_len, d_model)
+        input_ids = input_ids * (input_ids != -1).to(input_ids.dtype)
+        x = self.transformer.wte(input_ids)
+
+        image_features: Optional[torch.FloatTensor] = None    
+        if images is not None:
+            image_features = self.vision_backbone(images, token_pooling).to(x.device)
+            is_image_patch = input_ids.view(-1) == self.config.image_patch_id
+            assert is_image_patch.sum() == len(image_features)
+            x.view(-1, x.shape[-1])[is_image_patch] += image_features
+
+        # shape: (batch_size, seq_len, d_model)
+        x = self.transformer.emb_drop(x)  # type: ignore
+
+        return x, image_features
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        image_token_pooling: Optional[torch.Tensor] = None,
+        image_grids: Optional[torch.Tensor] = None,
+        image_num_crops: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.Tensor] = None,
+        video_token_pooling: Optional[torch.Tensor] = None,
+        video_grids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, Molmo2ModelOutputWithPast]:
+        
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        
+        images, token_pooling = self.merge_visual_inputs(
+            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,
+        )
+
+        if images is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both images and inputs_embeds at the same time."
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds, image_features = self.build_input_embeddings(
+                input_ids, images, token_pooling,
+            )
+        
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        # Adapted from transformers.models.gemma3.modeling_gemma3
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            mask_kwargs = {
+                "config": self.config.get_text_config(),
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+
+            # NOTE: this `is_prefill` logic is not flawless, it fails when we're using a cache eagerly initialized
+            # (e.g. compiled prefill) AND `images` are not provided. Determining prefill in that case requires
+            # checking data values, which is not compile-compatible.
+            is_prefill = (
+                not use_cache
+                or past_key_values is None
+                or not past_key_values.is_initialized
+                or images is not None
+            )
+            if token_type_ids is not None and is_prefill:
+                # We need to pass an additional mask function to account for token type ids, and it needs to be an `or`
+                mask_kwargs["or_mask_function"] = token_type_ids_mask_function(
+                    token_type_ids.to(cache_position.device)
+                )
+            
+            # Create the mask
+            causal_mask_mapping = create_causal_mask(**mask_kwargs)
+        
+        outputs = self.transformer(
+            attention_mask=causal_mask_mapping,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            cache_position=cache_position,
+            **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 if images is not None else None,
+        )
+
+
+class Molmo2ForConditionalGeneration(Molmo2PreTrainedModel, GenerationMixin):
+    _checkpoint_conversion_mapping = {}
+    _tied_weights_keys = []  # Weights are not tied
+    # Reference: fix gemma3 grad acc #37208
+    accepts_loss_kwargs = False
+    config: Molmo2Config
+
+    def __init__(self, config: Molmo2Config):
+        super().__init__(config)
+
+        self.model = Molmo2Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.vocab_size = config.vocab_size
+        
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.transformer.wte
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.transformer.wte = value
+    
+    def set_decoder(self, decoder):
+        self.model.set_decoder(decoder)
+    
+    def get_decoder(self):
+        return self.model.get_decoder()
+    
+    # Make modules available throught conditional class for BC
+    @property
+    def language_model(self) -> torch.nn.Module:
+        return self.model.transformer
+
+    @property
+    def vision_backbone(self) -> torch.nn.Module:
+        return self.model.vision_backbone
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        image_token_pooling: Optional[torch.Tensor] = None,
+        image_grids: Optional[torch.Tensor] = None,
+        image_num_crops: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.Tensor] = None,
+        video_token_pooling: Optional[torch.Tensor] = None,
+        video_grids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, Molmo2CausalLMOutputWithPast]:
+        r"""
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, Molmo2ForConditionalGeneration
+
+        >>> model = Molmo2ForConditionalGeneration.from_pretrained("...")
+        >>> processor = AutoProcessor.from_pretrained("...")
+
+        >>> 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,
+            token_type_ids=token_type_ids,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            cache_position=cache_position,
+            **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)
+
+        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,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        image_token_pooling: Optional[torch.Tensor] = None,
+        image_grids: Optional[torch.Tensor] = None,
+        image_num_crops: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.Tensor] = None,
+        video_token_pooling: Optional[torch.Tensor] = None,
+        video_grids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Optional[Union[int, torch.Tensor]] = None,
+        **kwargs,
+    ):
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            logits_to_keep=logits_to_keep,
+            token_type_ids=token_type_ids,
+            **kwargs,
+        )
+
+        if cache_position[0] == 0:
+            model_inputs["pixel_values"] = pixel_values
+            model_inputs["image_token_pooling"] = image_token_pooling
+            model_inputs["image_grids"] = image_grids
+            model_inputs["image_num_crops"] = image_num_crops
+            model_inputs["pixel_values_videos"] = pixel_values_videos
+            model_inputs["video_token_pooling"] = video_token_pooling
+            model_inputs["video_grids"] = video_grids
+
+        return model_inputs
+    
+    # Adapted from transformers.models.gemma3.modeling_gemma3
+    @staticmethod
+    def create_masks_for_generate(
+        config: PretrainedConfig,
+        input_embeds: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        cache_position: torch.Tensor,
+        past_key_values: Optional[Cache],
+        position_ids: Optional[torch.Tensor],
+        token_type_ids: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> dict:
+        # Prepare mask arguments
+        mask_kwargs = {
+            "config": config.get_text_config(),
+            "input_embeds": input_embeds,
+            "attention_mask": attention_mask,
+            "cache_position": cache_position,
+            "past_key_values": past_key_values,
+            "position_ids": position_ids,
+        }
+        # Add the token type ids mask for generate as well
+        if token_type_ids is not None and input_embeds.shape[1] != 1:
+            # We need to pass an additional mask function to account for token type ids, and it needs to be an `or`
+            mask_kwargs["or_mask_function"] = token_type_ids_mask_function(
+                token_type_ids.to(cache_position.device)
+            )
+        
+        return create_masks_for_generate(**mask_kwargs)
+
+
+# Always register for multi-modal features
+AutoModelForImageTextToText.register(Molmo2Config, Molmo2ForConditionalGeneration)

modelling_molmo_point.py

@@ -0,0 +1,1914 @@
+import math
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Optional, Union, Callable
+
+import torch
+from torch import nn
+
+from torch.nn import functional as F
+
+from transformers.models.auto import AutoModelForImageTextToText
+from transformers.activations import ACT2FN
+from transformers.configuration_utils import PretrainedConfig
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask, create_masks_for_generate
+from transformers.modeling_flash_attention_utils import (
+    _flash_attention_forward,
+    FlashAttentionKwargs,
+    flash_attn_supports_top_left_mask,
+)
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+)
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    ModelOutput,
+    TransformersKwargs,
+    can_return_tuple,
+    logging,
+)
+
+from .configuration_molmo2 import Molmo2VitConfig, Molmo2TextConfig, Molmo2AdapterConfig
+from .configuration_molmo_point import MolmoPointConfig, MolmoPointAdapterConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class MolmoPointCausalLMOutputWithPast(ModelOutput):
+    """
+    Base class for MolmoPoint causal language model (or autoregressive) outputs.
+
+    Args:
+        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: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+    image_hidden_states: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class MolmoPointModelOutputWithPast(BaseModelOutputWithPast):
+    """
+    Base class for Molmo2 outputs, with hidden states and attentions.
+
+    Args:
+        image_hidden_states (`torch.FloatTensor`, *optional*):
+            A `torch.FloatTensor` of size `(batch_num_patches, hidden_size)`.
+            image_hidden_states of the model produced by the vision backbone
+    """
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+    image_hidden_states: Optional[torch.FloatTensor] = None
+
+
+
+class ViTMLP(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, hidden_act: str, device: Union[str, torch.device] = None):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=True, device=device)
+        self.act = ACT2FN[hidden_act]
+        self.w2 = nn.Linear(hidden_dim, dim, bias=True, device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(self.act(self.w1(x)))
+
+
+class ViTMultiHeadDotProductAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        use_bias: bool = True,
+        input_dim: Optional[int] = None,
+        float32_attention: bool = True,
+        attention_dropout: float = 0.0,
+        residual_dropout: float = 0.0,
+        device: Union[str, torch.device] = None,
+        attn_implementation: str = "eager",
+        out_layer: bool=True
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.attn_implementation = attn_implementation
+        self.is_causal = False
+
+        input_dim = input_dim or hidden_size
+
+        self.wq = nn.Linear(
+            input_dim,
+            self.num_heads * self.head_dim,
+            bias=use_bias,
+            device=device,
+            )
+        self.wk = nn.Linear(
+            input_dim,
+            self.num_key_value_heads * self.head_dim,
+            bias=use_bias,
+            device=device,
+            )
+        self.wv = nn.Linear(
+            input_dim,
+            self.num_key_value_heads * self.head_dim,
+            bias=use_bias,
+            device=device,
+            )
+        if out_layer:
+            self.wo = nn.Linear(
+                self.num_heads * self.head_dim,
+                self.hidden_size,
+                )
+        else:
+            self.w0 = None
+        self.float32_attention = float32_attention
+        self.attention_dropout = attention_dropout
+        self.residual_dropout = nn.Dropout(residual_dropout)
+
+    def _split_heads(self, hidden_states, num_heads) -> torch.Tensor:
+        return hidden_states.reshape(hidden_states.shape[:2] + (num_heads, self.head_dim))
+
+    def _merge_heads(self, hidden_states) -> torch.Tensor:
+        return hidden_states.reshape(hidden_states.shape[:2] + (self.hidden_size,))
+
+    def forward(
+        self,
+        inputs_q: torch.Tensor,
+        inputs_kv: Optional[torch.Tensor] = None,
+        attn_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        if inputs_kv is not None:
+            inputs_k = inputs_kv
+            inputs_v = inputs_kv
+        else:
+            inputs_k = inputs_q
+            inputs_v = inputs_q
+
+        xq, xk, xv = self.wq(inputs_q), self.wk(inputs_k), self.wv(inputs_v)
+
+        xq = self._split_heads(xq, self.num_heads)
+        xk = self._split_heads(xk, self.num_key_value_heads)
+        xv = self._split_heads(xv, self.num_key_value_heads)
+
+        if self.num_heads != self.num_key_value_heads:
+            xk = xk.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads)
+            xv = xv.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads)
+
+        og_dtype = xq.dtype
+
+        if self.float32_attention:
+            xq = xq.to(torch.float)
+            xk = xk.to(torch.float)
+
+        dropout_p = 0.0 if not self.training else self.attention_dropout
+
+        if self.attn_implementation == "eager":
+            attn_weights = torch.einsum("...qhd,...khd->...hqk", xq / math.sqrt(xq.size(-1)), xk)
+            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(xq.dtype)
+            attn_weights = F.dropout(
+                attn_weights,
+                p=dropout_p,
+                training=self.training
+            )
+            attn_output = torch.einsum("...hqk,...khd->...qhd", attn_weights.to(xv.dtype), xv)
+
+        elif self.attn_implementation == "sdpa":
+            if not torch.is_autocast_enabled():
+                xv = xv.to(torch.float)
+
+            attn_output = F.scaled_dot_product_attention(
+                xq.transpose(1, 2).contiguous(),
+                xk.transpose(1, 2).contiguous(),
+                xv.transpose(1, 2).contiguous(),
+                attn_mask=attn_mask,
+                is_causal=False,
+                dropout_p=dropout_p,
+            ).transpose(1, 2)
+
+        elif self.attn_implementation == "flash_attention_2":
+            if xq.dtype == torch.float32:
+                if torch.is_autocast_enabled():
+                    target_dtype = torch.get_autocast_gpu_dtype()
+                else:
+                    target_dtype = self.wq.weight.dtype
+            attn_output = _flash_attention_forward(
+                xq,
+                xk,
+                xv,
+                attention_mask=attn_mask,
+                query_length=inputs_q.shape[1],
+                is_causal=False,
+                dropout=dropout_p,
+                softmax_scale=xq.shape[-1] ** -0.5,
+                use_top_left_mask=flash_attn_supports_top_left_mask(),
+                target_dtype=target_dtype,
+                implementation=self.attn_implementation,
+            )
+        else:
+            raise ValueError(f"Attention implementation {self.attn_implementation} not supported")
+
+        attn_output = attn_output.to(og_dtype)
+        attn_output = self._merge_heads(attn_output)
+        if self.wo is not None:
+            attn_output = self.wo(attn_output)
+        attn_output = self.residual_dropout(attn_output)
+
+        return attn_output
+
+
+class Molmo2VisionBlock(nn.Module):
+
+    def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None):
+        super().__init__()
+        self.attention = ViTMultiHeadDotProductAttention(
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_key_value_heads=config.num_key_value_heads,
+            head_dim=config.head_dim,
+            float32_attention=config.float32_attention,
+            attention_dropout=config.attention_dropout,
+            residual_dropout=config.residual_dropout,
+            device=device,
+            attn_implementation=config._attn_implementation,
+        )
+        self.feed_forward = ViTMLP(config.hidden_size, config.intermediate_size, config.hidden_act, device=device)
+        self.attention_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device)
+        self.ffn_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attention(self.attention_norm(x))
+        x = x + self.feed_forward(self.ffn_norm(x))
+        return x
+
+
+class Molmo2VisionBlockCollection(nn.Module):
+
+    def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None):
+        super().__init__()
+        self.conifg = config
+        self.resblocks = nn.ModuleList([
+            Molmo2VisionBlock(config, device) for _ in range(config.num_hidden_layers)
+        ])
+
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        hidden_states = []
+        for r in self.resblocks:
+            x = r(x)
+            hidden_states.append(x)
+        return hidden_states
+
+
+class Molmo2VisionTransformer(nn.Module):
+
+    def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None):
+        super().__init__()
+        self.config = config
+
+        # positional embeddings
+        self.scale = config.hidden_size ** -0.5
+        self.num_prefix_tokens: int = 0 # no class embeddings
+        self.positional_embedding = nn.Parameter(
+            torch.zeros(config.image_num_pos, config.hidden_size, device=device),
+        )
+
+        image_patch_size = config.image_patch_size
+        self.patch_embedding = nn.Linear(
+            image_patch_size * image_patch_size * 3,
+            config.hidden_size,
+            bias=True,
+            device=device,
+            )
+
+        self.transformer = Molmo2VisionBlockCollection(config, device)
+
+    def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor:
+        pos_emb = self.positional_embedding
+
+        pos_emb = pos_emb.reshape(
+            (int(math.sqrt(pos_emb.shape[0])), int(math.sqrt(pos_emb.shape[0])), pos_emb.shape[1])
+        )
+
+        (patch_num_0, patch_num_1) = patch_num
+
+        if pos_emb.shape[0] != patch_num_0 or pos_emb.shape[1] != patch_num_1:
+            # Dervied from https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+            # antialias: default True in jax.image.resize
+            pos_emb = pos_emb.unsqueeze(0).permute(0, 3, 1, 2)
+            pos_emb = F.interpolate(
+                pos_emb, size=(patch_num_0, patch_num_1), mode="bicubic", align_corners=False, antialias=True,
+            )
+            pos_emb = pos_emb.permute(0, 2, 3, 1).squeeze(0)
+
+        pos_emb = pos_emb.reshape(-1, pos_emb.shape[-1])
+        x = x + pos_emb[None, :, :].to(x.dtype)
+        return x
+
+    def forward(self, x: torch.Tensor, patch_num: int = None) -> list[torch.Tensor]:
+        """
+        : param x: (batch_size, num_patch, n_pixels)
+        """
+        if patch_num is None:
+            patch_num = self.config.image_num_patch
+
+        B, N, D = x.shape
+
+        x = self.patch_embedding(x)
+
+        # class embeddings and positional embeddings
+        x = self.add_pos_emb(x, patch_num)
+
+        hidden_states = self.transformer(x)
+        return hidden_states
+
+
+class ImageProjectorMLP(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        hidden_act: str,
+        device: Union[str, torch.device] = None,
+    ):
+        super().__init__()
+        self.w1 = nn.Linear(input_dim, hidden_dim, bias=False, device=device)
+        self.w2 = nn.Linear(hidden_dim, output_dim, bias=False, device=device)
+        self.w3 = nn.Linear(input_dim, hidden_dim, bias=False, device=device)
+        self.act = ACT2FN[hidden_act]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(self.act(self.w1(x)) * self.w3(x))
+
+
+class Molmo2VisionBackbone(nn.Module):
+    def __init__(self, vit_config: Molmo2VitConfig, adapter_config: Molmo2AdapterConfig):
+        super().__init__()
+        self.vit_config = vit_config
+        self.adapter_config = adapter_config
+
+        self.vit_layers = []
+        for layer in adapter_config.vit_layers:
+            if layer >= 0:
+                self.vit_layers.append(layer)
+            else:
+                self.vit_layers.append(layer + vit_config.num_hidden_layers)
+
+        last_layer_needed = max(self.vit_layers) + 1
+        if last_layer_needed < vit_config.num_hidden_layers:
+            new_vit_config = deepcopy(vit_config)
+            new_vit_config.num_hidden_layers = last_layer_needed
+            self.image_vit = Molmo2VisionTransformer(new_vit_config)
+        else:
+            self.image_vit = Molmo2VisionTransformer(vit_config)
+
+        self.num_prefix_tokens: int = self.image_vit.num_prefix_tokens
+
+        pool_dim = vit_config.hidden_size * len(adapter_config.vit_layers)
+        self.image_pooling_2d = ViTMultiHeadDotProductAttention(
+            hidden_size=adapter_config.hidden_size,
+            num_heads=adapter_config.num_attention_heads,
+            num_key_value_heads=adapter_config.num_key_value_heads,
+            head_dim=adapter_config.head_dim,
+            input_dim=pool_dim,
+            float32_attention=adapter_config.float32_attention,
+            attention_dropout=adapter_config.attention_dropout,
+            residual_dropout=adapter_config.residual_dropout,
+            attn_implementation=adapter_config._attn_implementation,
+        )
+        self.image_projector = ImageProjectorMLP(
+            adapter_config.hidden_size,
+            adapter_config.intermediate_size,
+            adapter_config.text_hidden_size,
+            adapter_config.hidden_act,
+        )
+        self.image_feature_dropout = nn.Dropout(adapter_config.image_feature_dropout)
+
+    def encode_image(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        : param images: (batch_size, num_crops, num_patch, n_pixels)
+        """
+        B, T, N, D = images.shape
+        images = images.view(B * T, N, D)
+        image_features = self.image_vit(images)
+
+        features = []
+        for layer in self.vit_layers:
+            features.append(image_features[layer])
+        image_features = torch.cat(features, dim=-1)
+
+        if self.num_prefix_tokens > 0:
+            image_features = image_features[:, 1:]
+        image_features = image_features.view(B, T, N, -1)
+        return image_features
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.image_vit.patch_embedding.weight.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.image_vit.patch_embedding.weight.device
+
+    def forward(
+        self,
+        images: torch.Tensor,
+        pooled_patches_idx: torch.Tensor,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+        # image_features: (batch_size, num_crops(=num_image), num_patch, nximage_emb_dim)
+        batch_size, num_image = images.shape[:2]
+        images = images.to(device=self.device, dtype=self.dtype)
+        image_features = self.encode_image(images)
+
+        image_features = self.image_feature_dropout(image_features)
+        dim = image_features.shape[-1]
+        valid = pooled_patches_idx >= 0
+        valid_token = torch.any(valid, -1)
+
+        # Use `pooled_patches_idx` to arange the features for image pooling
+        batch_idx = torch.arange(pooled_patches_idx.shape[0], dtype=torch.long, device=pooled_patches_idx.device)
+        batch_idx = torch.tile(batch_idx.view(batch_size, 1, 1), [1, pooled_patches_idx.shape[1], pooled_patches_idx.shape[2]])
+
+        # Now [batch, num_high_res_features, pool_dim, dim]
+        to_pool = image_features.reshape(batch_size, -1, dim)[batch_idx, torch.clip(pooled_patches_idx, 0)]
+        to_pool = to_pool * valid.to(self.dtype)[:, :, :, None]
+        to_pool = to_pool.reshape([-1, pooled_patches_idx.shape[-1], dim])
+        if self.adapter_config.pooling_attention_mask:
+            attn_mask = valid.reshape([-1, 1, 1, valid.shape[-1]])
+            denom = valid.view(-1, to_pool.shape[-2]).float().sum(-1)
+            denom = torch.where(denom == 0, 1, denom)
+            query = to_pool.sum(-2, keepdim=True) / denom[:, None, None].to(to_pool.dtype)
+        else:
+            attn_mask = None
+            query = to_pool.mean(-2, keepdim=True)
+        pooled_features = self.image_pooling_2d(query, to_pool, attn_mask=attn_mask)
+        pooled_features = pooled_features.reshape([batch_size, -1, pooled_features.shape[-1]])
+
+        # MLP layer to map the feature.
+        pooled_features = self.image_projector(pooled_features)
+        return pooled_features.view(-1, pooled_features.shape[-1])[valid_token.flatten()]
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+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)
+
+
+# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, 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.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        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
+
+
+class Molmo2RotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(
+        self,
+        config: Molmo2TextConfig,
+        device: Union[str, torch.device] = None,
+        rope_type: Optional[str] = None,
+    ):
+        super().__init__()
+        if rope_type is not None:
+            self.rope_type = rope_type
+        elif hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            # BC: "rope_type" was originally "type"
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            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)
+
+
+class Molmo2RMSNorm(nn.Module):
+
+    def __init__(
+        self,
+        size: int,
+        eps: float = 1e-6,
+        device: Union[str, torch.device] = None,
+    ):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(size, device=device))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        with torch.autocast(enabled=False, device_type=x.device.type):
+            og_dtype = x.dtype
+            x = x.to(torch.float32)
+            variance = x.pow(2).mean(-1, keepdim=True)
+            x = x * torch.rsqrt(variance + self.eps)
+            x = x.to(og_dtype)
+
+        return self.weight * x
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+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: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    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:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_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
+
+
+class Molmo2Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: Molmo2TextConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.is_causal = True
+
+        self.fused_dims = (
+            config.num_attention_heads * config.head_dim,
+            config.head_dim * config.num_key_value_heads,
+            config.head_dim * config.num_key_value_heads,
+        )
+        self.att_proj = nn.Linear(
+            config.hidden_size,
+            sum(self.fused_dims),
+            bias=config.qkv_bias,
+        )
+
+        # Layer norms.
+        self.k_norm: Optional[Molmo2RMSNorm] = None
+        self.q_norm: Optional[Molmo2RMSNorm] = None
+        self.qk_norm_type: Optional[str] = None
+        if config.use_qk_norm:
+            k_norm_size = (
+                config.head_dim
+                if config.qk_norm_type == "qwen3" else
+                config.num_key_value_heads * config.head_dim
+            )
+            self.k_norm = Molmo2RMSNorm(k_norm_size, eps=config.layer_norm_eps)
+            q_norm_size = (
+                config.head_dim
+                if config.qk_norm_type == "qwen3" else
+                config.num_attention_heads * config.head_dim
+            )
+            self.q_norm = Molmo2RMSNorm(q_norm_size, eps=config.layer_norm_eps)
+            self.qk_norm_type = config.qk_norm_type
+
+        self.attention_dropout = config.attention_dropout
+
+        self.attn_out = nn.Linear(
+            config.head_dim * config.num_attention_heads,
+            config.hidden_size,
+            bias=False,
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        qkv = self.att_proj(hidden_states)
+        query_states, key_states, value_states = qkv.split(self.fused_dims, dim=-1)
+        value_states = value_states.view(hidden_shape)
+
+        # Optionally apply layer norm to keys and queries.
+        if self.q_norm is not None and self.k_norm is not None and self.qk_norm_type != "qwen3":
+            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.q_norm is not None and self.k_norm is not None and 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)
+        value_states = value_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:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        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.attn_out(attn_output)
+        return attn_output, attn_weights
+
+
+class LanguageModelMLP(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int,
+        intermediate_size: int,
+        hidden_act: str,
+        device: Union[str, torch.device] = None,
+    ):
+        super().__init__()
+        self.ff_proj = nn.Linear(input_dim, intermediate_size * 2, bias=False, device=device)
+        self.ff_out = nn.Linear(intermediate_size, input_dim, bias=False, device=device)
+        self.act = ACT2FN[hidden_act]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.ff_proj(x)
+        x, gate = x.chunk(2, dim=-1)
+        x = self.act(gate) * x
+        x = self.ff_out(x)
+        return x
+
+
+class Molmo2DecoderLayer(GradientCheckpointingLayer):
+
+    def __init__(
+        self,
+        config: Molmo2TextConfig,
+        layer_idx: Optional[int] = None,
+        device: Union[str, torch.device] = None
+    ):
+        super().__init__()
+        self.config = config
+
+        self.self_attn = Molmo2Attention(config, layer_idx)
+        self.attn_norm = Molmo2RMSNorm(
+            config.hidden_size, eps=config.layer_norm_eps, device=device)
+        self.dropout = nn.Dropout(config.residual_dropout)
+        self.mlp = LanguageModelMLP(
+            config.hidden_size, config.intermediate_size, config.hidden_act, device=device)
+        self.ff_norm = Molmo2RMSNorm(
+            config.hidden_size, eps=config.layer_norm_eps, device=device)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+
+        residual = hidden_states
+        hidden_states = self.attn_norm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights = 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,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = residual + self.dropout(hidden_states)
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.ff_norm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+
+        hidden_states = residual + self.dropout(hidden_states)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs
+
+
+class Molmo2PostNormDecoderLayer(Molmo2DecoderLayer):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
+
+        residual = hidden_states
+
+        # Self Attention
+        hidden_states, self_attn_weights = 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,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+        )
+        hidden_states = self.attn_norm(hidden_states)
+
+        hidden_states = residual + self.dropout(hidden_states)
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.ff_norm(hidden_states)
+
+        hidden_states = residual + self.dropout(hidden_states)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs
+
+
+class Molmo2Embedding(nn.Module):
+    def __init__(
+        self,
+        num_embeddings: int,
+        num_new_embeddings: int,
+        features: int,
+        device: Union[str, torch.device] = None,
+    ):
+        super().__init__()
+        self.embedding = nn.Parameter(
+            torch.zeros(num_embeddings, features, device=device),
+        )
+        self.new_embedding = nn.Parameter(
+            torch.zeros(num_new_embeddings, features, device=device),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.embedding(x, torch.cat([self.embedding, self.new_embedding], dim=0))
+
+
+class MolmoPointPreTrainedModel(PreTrainedModel):
+    config: MolmoPointConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = [
+        "Molmo2DecoderLayer",
+        "Molmo2PostNormDecoderLayer",
+        "Molmo2VisionBlock",
+        "ViTMultiHeadDotProductAttention",
+    ]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": Molmo2DecoderLayer,
+        "attentions": Molmo2Attention,
+    }
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, (nn.Linear,)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, Molmo2Embedding):
+            module.embedding.data.normal_(mean=0.0, std=std)
+            module.new_embedding.data.normal_(mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, Molmo2RMSNorm):
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+
+class MolmoPointTextModel(PreTrainedModel):
+    config: Molmo2TextConfig
+    _no_split_modules = ["Molmo2DecoderLayer", "Molmo2PostNormDecoderLayer"]
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": Molmo2DecoderLayer,
+        "attentions": Molmo2Attention,
+    }
+
+    def __init__(self, config: Molmo2TextConfig):
+        super().__init__(config)
+        if config.additional_vocab_size is not None:
+            self.wte = Molmo2Embedding(
+                config.vocab_size,
+                config.additional_vocab_size,
+                config.hidden_size,
+            )
+        else:
+            self.wte = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.emb_drop = nn.Dropout(config.embedding_dropout)
+        decoder_layer = Molmo2PostNormDecoderLayer if config.norm_after else Molmo2DecoderLayer
+        self.blocks = nn.ModuleList(
+            [decoder_layer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.ln_f = Molmo2RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+        if config.rope_scaling_layers is not None:
+            self.rotary_embs = nn.ModuleDict(
+                {
+                    "default": Molmo2RotaryEmbedding(config, rope_type="default"),
+                    "scaling": Molmo2RotaryEmbedding(config),
+                }
+            )
+        else:
+            self.rotary_emb = Molmo2RotaryEmbedding(config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.wte
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.wte = value
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        if inputs_embeds is None:
+            input_ids = input_ids * (input_ids != -1).to(input_ids.dtype)
+            inputs_embeds = self.wte(input_ids)
+
+        # torch.jit.trace() doesn't support cache objects in the output
+        if use_cache and past_key_values is None and not torch.jit.is_tracing():
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+                )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+
+            # Create the mask
+            causal_mask_mapping = create_causal_mask(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        if self.config.rope_scaling_layers is not None:
+            position_embeddings_mapping = {
+                "default": self.rotary_embs["default"](hidden_states, position_ids),
+                "scaling": self.rotary_embs["scaling"](hidden_states, position_ids),
+            }
+        else:
+            position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for layer_idx, decoder_block in enumerate(self.blocks[: self.config.num_hidden_layers]):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.config.rope_scaling_layers is not None:
+                position_embeddings_i = (
+                    position_embeddings_mapping["scaling"]
+                    if layer_idx in self.config.rope_scaling_layers
+                    else position_embeddings_mapping["default"]
+                )
+            else:
+                position_embeddings_i = position_embeddings
+
+            layer_outputs = decoder_block(
+                hidden_states,
+                attention_mask=causal_mask_mapping,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings_i,
+                **kwargs,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.ln_f(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+# Adapted from transformers.models.gemma3.modeling_gemma3
+def token_type_ids_mask_function(
+    token_type_ids: Optional[torch.Tensor] = None,
+) -> Optional[Callable]:
+    """
+    This function adds the correct offsets to the `q_idx` and `kv_idx` as the torch API can only accept lengths,
+    not start and end indices.
+    """
+    # Do not return an additional mask in this case
+    if token_type_ids is None:
+        return None
+
+    def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
+        # If it's 1 for both query and key/value, we are in an image block
+        # NOTE: static cache shape goes beyond input seq length, while token_type_ids.shape[1] == input seq length
+        # Since vmap doesn't support `if statement` we workaround it with `torch.where`
+        safe_idx = torch.where(kv_idx < token_type_ids.shape[1], kv_idx, 0)
+        token_type_ids_at_kv_idx = token_type_ids[batch_idx, safe_idx]
+        token_type_ids_at_kv_idx = torch.where(kv_idx < token_type_ids.shape[1], token_type_ids_at_kv_idx, 0)
+
+        is_image_block = (token_type_ids[batch_idx, q_idx] == 1) & (token_type_ids_at_kv_idx == 1)
+
+        # This is bidirectional attention whenever we are dealing with image tokens
+        return is_image_block & is_image_block
+
+    return inner_mask
+
+
+class MolmoPointPadWithLearnedVector(nn.Module):
+    """Module that pads vector
+
+    Used to add in the no-more-point key value
+    """
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+        self.vector = nn.Parameter(torch.zeros([dim]))
+
+    def reset_parameters(self):
+        torch.nn.init.zeros_(self.vector)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        vector = torch.tile(self.vector[None, :], [x.shape[0], 1])
+        return torch.concatenate([x, vector[:, None, :]], dim=1)
+
+
+class AddPosEmbed(nn.Module):
+
+    def __init__(self, in_features: int, n_pos: int) -> None:
+        super().__init__()
+        self.bias = nn.Parameter(torch.zeros([n_pos, in_features]))
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return input + self.bias[None, :input.shape[-2], :]
+
+
+class MolmoPointConnector(nn.Module):
+    def __init__(self, config: MolmoPointAdapterConfig, vit_config: Molmo2VitConfig):
+        super().__init__()
+        self.config = config
+        self.n_vit_layers = len(config.vit_layers)
+        pool_dim = vit_config.hidden_size * self.n_vit_layers
+        self.norm = None
+        self.image_projector = ImageProjectorMLP(
+            config.hidden_size,
+            config.intermediate_size,
+            config.text_hidden_size,
+            config.hidden_act,
+        )
+        self.act = ACT2FN[config.hidden_act]
+        self.image_pooling_2d = ViTMultiHeadDotProductAttention(
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_key_value_heads=config.num_key_value_heads,
+            head_dim=config.head_dim,
+            input_dim=pool_dim,
+            float32_attention=config.float32_attention,
+            attention_dropout=config.attention_dropout,
+            residual_dropout=config.residual_dropout,
+            attn_implementation=config._attn_implementation,
+            out_layer=False
+        )
+        if self.config.positional_embeddings:
+            self.positional_embeddings = AddPosEmbed(pool_dim, self.config.positional_embeddings)
+        else:
+            self.positional_embeddings = None
+
+    def __call__(self, to_pool, to_pool_mask):
+        """
+        to_pool: [n_to_pool, pooling_dim, vit_dim]
+        to_pool_mask: [n_to_pool, pooling_dim]
+
+        returns:
+        pooled_features: [n_to_pool, llm_dim]
+        """
+        cfg = self.config
+
+        if self.config.positional_embeddings:
+            to_pool = self.positional_embeddings(to_pool)
+
+        if self.config.pooling_attention_mask:
+            attn_mask = to_pool_mask.reshape([-1, 1, 1, to_pool_mask.shape[-1]])
+        else:
+            attn_mask = None
+            to_pool = to_pool * to_pool_mask.float()[:, :, None]
+
+        denom = to_pool_mask.view(-1, to_pool.shape[-2]).float().sum(-1)
+        denom = torch.where(denom == 0, 1, denom)
+        query = to_pool.sum(-2, keepdim=True) / denom[:, None, None]
+
+        pooled_features = self.image_pooling_2d(query, to_pool, attn_mask=attn_mask)
+        pooled_features = self.act(pooled_features)
+        pooled_features = self.image_projector(pooled_features)
+        return pooled_features
+
+
+class MolmoPointModel(MolmoPointPreTrainedModel):
+    base_model_prefix = ""
+    _checkpoint_conversion_mapping = {}
+    # Reference: fix gemma3 grad acc #37208
+    accepts_loss_kwargs = False
+    config: MolmoPointConfig
+
+    def __init__(self, config: MolmoPointConfig):
+        super().__init__(config)
+        self.transformer: MolmoPointTextModel = MolmoPointTextModel(config.text_config)
+
+        vit_config = config.vit_config
+        adapter_config = config.adapter_config
+        self.vit_layers = []
+        for layer in adapter_config.vit_layers:
+            if layer >= 0:
+                self.vit_layers.append(layer)
+            else:
+                self.vit_layers.append(layer + vit_config.num_hidden_layers)
+
+        last_layer_needed = max(self.vit_layers) + 1
+        if last_layer_needed < vit_config.num_hidden_layers:
+            new_vit_config = deepcopy(vit_config)
+            new_vit_config.num_hidden_layers = last_layer_needed
+            self.vit = Molmo2VisionTransformer(new_vit_config)
+        else:
+            self.vit = Molmo2VisionTransformer(vit_config)
+
+        self.connector = MolmoPointConnector(adapter_config, vit_config)
+
+        vit_dim = self.config.vit_config.hidden_size * len(self.config.adapter_config.vit_layers)
+        llm_dim = self.config.text_config.hidden_size
+        self.patch_rotary = None
+        self.patch_q = nn.Linear(llm_dim, config.patch_embed_dim)
+        self.patch_k = nn.Linear(llm_dim, config.patch_embed_dim)
+        self.subpatch_q = nn.Linear(llm_dim, config.patch_embed_dim)
+        self.subpatch_k = nn.Linear(vit_dim, config.patch_embed_dim)
+        self.add_no_point_class_embed = MolmoPointPadWithLearnedVector(config.patch_embed_dim)
+
+        if self.config.embed_selected_vit_patch == "linear":
+            self.build_vit_embedding = nn.Linear(vit_dim, llm_dim, bias=True)
+        else:
+            raise NotImplementedError(f"Embedding {self.config.embed_selected_vit_patch} not implemented")
+
+        if self.config.patch_location == "3x3":
+            self.subpatch_loc_k = nn.Linear(llm_dim, 9)
+        elif self.config.patch_location is None:
+            self.subpatch_loc_k = None
+        else:
+            raise NotImplementedError(f"Patch location {self.config.patch_location} not implemented")
+
+        if self.config.layer_norm_x:
+            self.x_norm = Molmo2RMSNorm(llm_dim, eps=self.config.text_config.layer_norm_eps)
+        else:
+            self.x_norm = None
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.transformer.wte
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.transformer.wte = value
+
+    def set_decoder(self, decoder):
+        self.transformer = decoder
+
+    def get_decoder(self):
+        return self.transformer
+
+    @property
+    def device(self) -> torch.device:
+        return self.transformer.ln_f.weight.device
+
+    def build_batched_images(
+        self,
+        input_ids: torch.LongTensor,
+        pixel_values: torch.Tensor,
+        image_token_pooling: torch.Tensor,
+        image_grids: torch.Tensor,
+        image_num_crops: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # 1) Count the number of images in each example
+        raw_counts = (input_ids == self.config.image_end_token_id).sum(1)  # [N]
+        # Each image is represented by global view and high-res view
+        # so we divide by 2 to get the number of images
+        counts = raw_counts // 2
+        N = counts.size(0)
+        device = input_ids.device
+
+        # Total number of images in the batch
+        num_images = int(counts.sum().item())
+
+        # Sanity check
+        assert image_grids.size(0) == num_images, \
+            f"Expected {num_images} image grids, but got {image_grids.size(0)}"
+        assert image_num_crops.size(0) == num_images, \
+            f"Expected {num_images} image num crops, but got {image_num_crops.size(0)}"
+
+        # 1-1) Compute per-image pooled patch count from image grids
+        with torch.no_grad():
+            first_prod = image_grids[:, :2].prod(dim=1)    # [num_images]
+            second_prod = image_grids[:, 2:].prod(dim=1)   # [num_images]
+            num_pooled_patches_per_image = (first_prod + second_prod).to(image_num_crops.dtype)  # [num_images]
+
+        # pixel_values: [n_crops, n_patches, pixels_per_patch]
+        n_crops, n_patches, pixels_per_patch = pixel_values.shape
+
+        # 2) Map each image index → example index
+        # Example: if counts = [2, 1, 3], then this becomes [0,0,1,2,2,2]
+        example_ids_for_image = torch.arange(N, device=device).repeat_interleave(counts)  # [num_images]
+        assert example_ids_for_image.numel() == num_images
+
+        # 2-1) Compute crops_per_example by summing per-image crop counts
+        crops_per_example = torch.zeros(
+            N, dtype=image_num_crops.dtype, device=image_num_crops.device
+        )
+        crops_per_example.index_add_(0, example_ids_for_image, image_num_crops)  # [N]
+
+        # 2-2) Per-image number of patches = (crops per image) * n_patches
+        patches_per_image = image_num_crops * n_patches  # [num_images]
+
+        # 2-3) Compute per-example per-image patch offsets
+        counts_list = counts.tolist()
+        index_offset_per_example_list = []
+        offset_img = 0
+        for c in counts_list:
+            per_img_patches = patches_per_image[offset_img:offset_img + c]  # [c]
+            # Offsets: [0, img0_total_patches, img0+img1_total_patches, ...]
+            index_offset = [0] + per_img_patches.cumsum(0).tolist()[:-1]
+            index_offset_per_example_list.append(index_offset)
+            offset_img += c
+
+        # 2-4) Compute num_pooled_patches_per_example
+        num_pooled_patches_per_example = torch.zeros(
+            N, dtype=num_pooled_patches_per_image.dtype, device=num_pooled_patches_per_image.device
+        )
+        num_pooled_patches_per_example.index_add_(
+            0, example_ids_for_image, num_pooled_patches_per_image
+        )
+
+        # Sanity checks
+        total_crops = int(crops_per_example.sum().item())
+        assert total_crops == n_crops, \
+            f"Expected {total_crops} crops, but got {n_crops}"
+
+        total_num_pooled_patches = int(num_pooled_patches_per_example.sum().item())
+        assert total_num_pooled_patches == image_token_pooling.size(0), \
+            f"Expected {total_num_pooled_patches} pooled patches, but got {image_token_pooling.size(0)}"
+
+        # 3) Build images tensor filled with -1
+        M = int(crops_per_example.max().item())
+        images = torch.full(
+            (N, M, n_patches, pixels_per_patch),
+            fill_value=-1,
+            dtype=pixel_values.dtype,
+            device=pixel_values.device,
+        )
+
+        # 4) Fill images with per-example slices from pixel_values
+        offset_crop = 0
+        for i in range(N):
+            num = int(crops_per_example[i].item())
+            cur = pixel_values[offset_crop:offset_crop + num]  # [num, n_patches, pixels_per_patch]
+            images[i, :num] = cur
+            offset_crop += num
+
+        # Sanity check
+        assert offset_crop == n_crops
+
+        # 5) Build new_token_pooling tensor filled with -1
+        P = int(num_pooled_patches_per_example.max().item())
+        _, dim = image_token_pooling.shape
+        new_token_pooling = torch.full(
+            (N, P, dim),
+            fill_value=-1,
+            dtype=image_token_pooling.dtype,
+            device=image_token_pooling.device,
+        )
+
+        # 6) Fill token_pooling with per-example slices, adding per-image patch offsets
+        patch_offset = 0
+        img_offset = 0
+
+        for i, c in enumerate(counts_list):
+            num_patches = int(num_pooled_patches_per_example[i].item())
+
+            # Subsequence of pooled tokens belonging to this example
+            cur = image_token_pooling[patch_offset:patch_offset + num_patches].clone()  # [num_patches, dim]
+
+            index_offset_per_example = index_offset_per_example_list[i]  # length = c
+            per_img_pooled = num_pooled_patches_per_image[img_offset:img_offset + c]   # [c]
+
+            assert len(index_offset_per_example) == per_img_pooled.numel()
+
+            # Apply per-image offsets to the (ragged) subsequence
+            offset = 0
+            for j in range(c):
+                index_offset = int(index_offset_per_example[j])
+                n = int(per_img_pooled[j].item())
+                cur_slice = cur[offset:offset + n]
+
+                # Apply offset across all columns
+                cur[offset:offset + n] = torch.where(
+                    cur_slice >= 0,
+                    cur_slice + index_offset,
+                    cur_slice,
+                    )
+                offset += n
+
+            new_token_pooling[i, :num_patches] = cur
+
+            patch_offset += num_patches
+            img_offset += c
+
+        # Final sanity checks
+        assert patch_offset == total_num_pooled_patches
+        assert img_offset == num_images
+
+        return images, new_token_pooling
+
+    def build_batched_videos(
+        self,
+        input_ids: torch.LongTensor,
+        pixel_values_videos: torch.Tensor,
+        video_token_pooling: torch.Tensor,
+        video_grids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        # 1) Count the number of videos in each example
+        if self.config.use_frame_special_tokens:
+            end_token_id = self.config.frame_end_token_id
+        else:
+            end_token_id = self.config.image_end_token_id
+        counts = (input_ids == end_token_id).any(dim=1).long()  # [N]
+        N = counts.size(0)
+        device = input_ids.device
+
+        # Total number of videos in the batch
+        num_videos = int(counts.sum().item())
+
+        # Sanity check
+        assert video_grids.size(0) == num_videos, \
+            f"Expected {num_videos} videos, but got {video_grids.size(0)}"
+
+        video_num_frames = video_grids[:, 0]  # [num_videos]
+        num_pooled_patches_per_video = video_grids.prod(dim=1)  # [num_videos]
+
+        # pixel_values_videos: [n_frames, n_patches, pixels_per_patch]
+        n_frames, n_patches, pixels_per_patch = pixel_values_videos.shape
+
+        # 2) Map each video index -> example index
+        # Example: if counts = [2, 1, 3], then this becomes [0,0,1,2,2,2]
+        example_ids_for_video = torch.arange(N, device=device).repeat_interleave(counts)  # [num_videos]
+        assert example_ids_for_video.numel() == num_videos
+
+        # 2-1) Compute frames_per_example by summing per-video frame counts
+        frames_per_example = torch.zeros(
+            N, dtype=video_num_frames.dtype, device=device,
+        )
+        frames_per_example.index_add_(0, example_ids_for_video, video_num_frames)  # [N]
+
+        # 2-2) Compute num_pooled_patches_per_example
+        num_pooled_patches_per_example = torch.zeros(
+            N, dtype=num_pooled_patches_per_video.dtype, device=num_pooled_patches_per_video.device,
+        )
+        num_pooled_patches_per_example.index_add_(
+            0, example_ids_for_video, num_pooled_patches_per_video,
+        )
+
+        # Sanity checks
+        total_frames = int(frames_per_example.sum().item())
+        assert total_frames == n_frames, \
+            f"Expected {total_frames} frames, but got {n_frames}"
+
+        total_num_pooled_patches = int(num_pooled_patches_per_example.sum().item())
+        assert total_num_pooled_patches == video_token_pooling.size(0), \
+            f"Expected {total_num_pooled_patches} pooled patches, but got {video_token_pooling.size(0)}"
+
+        # 3) Build videos tensor filled with -1
+        M = int(frames_per_example.max().item())
+        videos = torch.full(
+            (N, M, n_patches, pixels_per_patch),
+            fill_value=-1,
+            dtype=pixel_values_videos.dtype,
+            device=device,
+        )
+
+        # 4) Fill videos with per-examples slices from pixel_values_videos
+        offset_frame = 0
+        for i in range(N):
+            num = int(frames_per_example[i].item())
+            cur = pixel_values_videos[offset_frame:offset_frame + num]  # [num, n_patches, pixels_per_patch]
+            videos[i, :num] = cur
+            offset_frame += num
+
+        # Sanity check
+        assert offset_frame == n_frames
+
+        # 5) Build new token_pooling tensor filled with -1
+        P = int(num_pooled_patches_per_example.max().item())
+        _, dim = video_token_pooling.shape
+        new_token_pooling = torch.full(
+            (N, P, dim),
+            fill_value=-1,
+            dtype=video_token_pooling.dtype,
+            device=video_token_pooling.device,
+        )
+
+        # 6) Fill new token_pooling with per-examples slices from video_token_pooling
+        patch_offset = 0
+        for i in range(N):
+            num_patches = int(num_pooled_patches_per_example[i].item())
+            cur = video_token_pooling[patch_offset:patch_offset + num_patches]  # [num_patches, dim]
+            new_token_pooling[i, :num_patches] = cur
+            patch_offset += num_patches
+
+        # Final sanity checks
+        assert patch_offset == total_num_pooled_patches
+
+        return videos, new_token_pooling
+
+    def merge_visual_inputs(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        image_token_pooling: Optional[torch.Tensor] = None,
+        image_grids: Optional[torch.Tensor] = None,
+        image_num_crops: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.Tensor] = None,
+        video_token_pooling: Optional[torch.Tensor] = None,
+        video_grids: Optional[torch.Tensor] = None,
+    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        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")
+        elif pixel_values is not None:
+            assert input_ids is not None
+            images, token_pooling = self.build_batched_images(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                image_token_pooling=image_token_pooling,
+                image_grids=image_grids,
+                image_num_crops=image_num_crops,
+            )
+        elif pixel_values_videos is not None:
+            assert input_ids is not None
+            images, token_pooling = self.build_batched_videos(
+                input_ids=input_ids,
+                pixel_values_videos=pixel_values_videos,
+                video_token_pooling=video_token_pooling,
+                video_grids=video_grids,
+            )
+        else:
+            images, token_pooling = None, None
+        return images, token_pooling
+
+    def build_input_embeddings(
+        self,
+        input_ids: torch.LongTensor,
+        images: Optional[torch.FloatTensor] = None,  # image inputs
+        token_pooling: Optional[torch.LongTensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+        # Get embeddings of input.
+        # shape: (batch_size, seq_len, d_model)
+        input_ids = input_ids * (input_ids != -1).to(input_ids.dtype)
+        x = self.transformer.wte(input_ids)
+
+        image_features: Optional[torch.FloatTensor] = None
+        if images is not None:
+            image_features = self.vision_backbone(images, token_pooling).to(x.device)
+            is_image_patch = input_ids.view(-1) == self.config.image_patch_id
+            assert is_image_patch.sum() == len(image_features)
+            x.view(-1, x.shape[-1])[is_image_patch] += image_features
+
+        # shape: (batch_size, seq_len, d_model)
+        x = self.transformer.emb_drop(x)  # type: ignore
+
+        return x, image_features
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        image_token_pooling: Optional[torch.Tensor] = None,
+        image_grids: Optional[torch.Tensor] = None,
+        image_num_crops: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.Tensor] = None,
+        video_token_pooling: Optional[torch.Tensor] = None,
+        video_grids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, MolmoPointModelOutputWithPast]:
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        images, token_pooling = self.merge_visual_inputs(
+            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,
+        )
+
+        if images is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both images and inputs_embeds at the same time."
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds, image_features = self.build_input_embeddings(
+                input_ids, images, token_pooling,
+            )
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+                )
+
+        # Adapted from transformers.models.gemma3.modeling_gemma3
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            # Prepare mask arguments
+            mask_kwargs = {
+                "config": self.config.get_text_config(),
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+
+            # NOTE: this `is_prefill` logic is not flawless, it fails when we're using a cache eagerly initialized
+            # (e.g. compiled prefill) AND `images` are not provided. Determining prefill in that case requires
+            # checking data values, which is not compile-compatible.
+            is_prefill = (
+                not use_cache
+                or past_key_values is None
+                or not past_key_values.is_initialized
+                or images is not None
+            )
+            if token_type_ids is not None and is_prefill:
+                # We need to pass an additional mask function to account for token type ids, and it needs to be an `or`
+                mask_kwargs["or_mask_function"] = token_type_ids_mask_function(
+                    token_type_ids.to(cache_position.device)
+                )
+
+            # Create the mask
+            causal_mask_mapping = create_causal_mask(**mask_kwargs)
+
+        outputs = self.transformer(
+            attention_mask=causal_mask_mapping,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        return MolmoPointModelOutputWithPast(
+            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 if images is not None else None,
+        )
+
+
+class MolmoPointForConditionalGeneration(MolmoPointPreTrainedModel, GenerationMixin):
+    _checkpoint_conversion_mapping = {}
+    _tied_weights_keys = []  # Weights are not tied
+    # Reference: fix gemma3 grad acc #37208
+    accepts_loss_kwargs = False
+    config: MolmoPointConfig
+
+    def __init__(self, config: MolmoPointConfig):
+        super().__init__(config)
+
+        self.model = MolmoPointModel(config)
+        self.output_embeddings = nn.Parameter(torch.zeros([config.vocab_size, config.hidden_size]))
+        self.new_output_embeddings = nn.Parameter(torch.zeros([128, config.hidden_size]))
+        self.vocab_size = config.vocab_size
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.transformer.wte
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.transformer.wte = value
+
+    def set_decoder(self, decoder):
+        self.model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.model.get_decoder()
+
+    # Make modules available throught conditional class for BC
+    @property
+    def language_model(self) -> torch.nn.Module:
+        return self.model.transformer
+
+    @property
+    def vision_backbone(self) -> torch.nn.Module:
+        return self.model.vision_backbone
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        image_token_pooling: Optional[torch.Tensor] = None,
+        image_grids: Optional[torch.Tensor] = None,
+        image_num_crops: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.Tensor] = None,
+        video_token_pooling: Optional[torch.Tensor] = None,
+        video_grids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, MolmoPointCausalLMOutputWithPast]:
+        r"""
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, MolmoPointForConditionalGeneration
+
+        >>> model = Molmo2ForConditionalGeneration.from_pretrained("...")
+        >>> processor = AutoProcessor.from_pretrained("...")
+
+        >>> 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,
+            token_type_ids=token_type_ids,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            cache_position=cache_position,
+            **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
+        lm_head = torch.concatenate([self.output_embeddings, self.new_output_embeddings], dim=0)
+        logits = F.linear(hidden_states[:, slice_indices, :], lm_head)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.vocab_size)
+
+        return MolmoPointCausalLMOutputWithPast(
+            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,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        image_token_pooling: Optional[torch.Tensor] = None,
+        image_grids: Optional[torch.Tensor] = None,
+        image_num_crops: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.Tensor] = None,
+        video_token_pooling: Optional[torch.Tensor] = None,
+        video_grids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Optional[Union[int, torch.Tensor]] = None,
+        **kwargs,
+    ):
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            logits_to_keep=logits_to_keep,
+            token_type_ids=token_type_ids,
+            **kwargs,
+        )
+
+        if cache_position[0] == 0:
+            model_inputs["pixel_values"] = pixel_values
+            model_inputs["image_token_pooling"] = image_token_pooling
+            model_inputs["image_grids"] = image_grids
+            model_inputs["image_num_crops"] = image_num_crops
+            model_inputs["pixel_values_videos"] = pixel_values_videos
+            model_inputs["video_token_pooling"] = video_token_pooling
+            model_inputs["video_grids"] = video_grids
+
+        return model_inputs
+
+    # Adapted from transformers.models.gemma3.modeling_gemma3
+    @staticmethod
+    def create_masks_for_generate(
+        config: PretrainedConfig,
+        input_embeds: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        cache_position: torch.Tensor,
+        past_key_values: Optional[Cache],
+        position_ids: Optional[torch.Tensor],
+        token_type_ids: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> dict:
+        # Prepare mask arguments
+        mask_kwargs = {
+            "config": config.get_text_config(),
+            "input_embeds": input_embeds,
+            "attention_mask": attention_mask,
+            "cache_position": cache_position,
+            "past_key_values": past_key_values,
+            "position_ids": position_ids,
+        }
+        # Add the token type ids mask for generate as well
+        if token_type_ids is not None and input_embeds.shape[1] != 1:
+            # We need to pass an additional mask function to account for token type ids, and it needs to be an `or`
+            mask_kwargs["or_mask_function"] = token_type_ids_mask_function(
+                token_type_ids.to(cache_position.device)
+            )
+
+        return create_masks_for_generate(**mask_kwargs)
+
+
+# Always register for multi-modal features
+AutoModelForImageTextToText.register(MolmoPointConfig, MolmoPointForConditionalGeneration)

preprocessor_config.json

@@ -0,0 +1,34 @@
+{
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_molmo2.Molmo2ImageProcessor",
+    "AutoProcessor": "processing_molmo2.Molmo2Processor"
+  },
+  "do_convert_rgb": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "Molmo2ImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "max_crops": 8,
+  "overlap_margins": [
+    4,
+    4
+  ],
+  "patch_size": 14,
+  "pooling_size": [
+    2,
+    2
+  ],
+  "processor_class": "Molmo2Processor",
+  "resample": 2,
+  "size": {
+    "height": 378,
+    "width": 378
+  }
+}

processing_molmo2.py

@@ -0,0 +1,423 @@
+"""
+Processor class for Molmo2.
+"""
+from typing import Optional, Union
+import dataclasses
+
+import numpy as np
+
+from transformers.image_utils import ImageInput
+from transformers.video_utils import VideoInput
+from transformers.processing_utils import (
+    Unpack,
+    ProcessingKwargs,
+    ProcessorMixin, AllKwargsForChatTemplate,
+)
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.tokenization_utils_base import TextInput, PreTokenizedInput
+from transformers.utils import logging
+
+from transformers import AutoTokenizer
+from .image_processing_molmo2 import Molmo2ImagesKwargs, Molmo2ImageProcessor
+from .video_processing_molmo2 import Molmo2VideoProcessorKwargs, Molmo2VideoProcessor
+
+
+logger = logging.get_logger(__name__)
+
+
+# Special tokens, these should be present in any tokenizer we use since the preprocessor uses them
+IMAGE_PATCH_TOKEN = f"<im_patch>"  # Where to insert high-res tokens
+IMAGE_LOW_RES_TOKEN = f"<im_low>"  # Where to insert low-res tokens
+IM_START_TOKEN = f"<im_start>"
+LOW_RES_IMAGE_START_TOKEN = f"<low_res_im_start>"
+FRAME_START_TOKEN = f"<frame_start>"
+IM_END_TOKEN = f"<im_end>"
+FRAME_END_TOKEN= f"<frame_end>"
+IM_COL_TOKEN = f"<im_col>"
+IMAGE_PROMPT = "<|image|>"
+VIDEO_PROMPT = "<|video|>"
+
+IMAGE_TOKENS = [
+    IMAGE_PATCH_TOKEN,
+    IM_COL_TOKEN,
+    IM_START_TOKEN,
+    LOW_RES_IMAGE_START_TOKEN,
+    FRAME_START_TOKEN,
+    IM_END_TOKEN,
+    FRAME_END_TOKEN,
+    IMAGE_LOW_RES_TOKEN,
+]
+
+
+class Molmo2ProcessorKwargs(ProcessingKwargs, total=False):
+    """Molmo2 processor kwargs"""
+    images_kwargs: Molmo2ImagesKwargs
+    videos_kwargs: Molmo2VideoProcessorKwargs
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+            "return_mm_token_type_ids": True,
+        },
+        "videos_kwargs": {"return_metadata": True},
+    }
+
+
+class Molmo2Processor(ProcessorMixin):
+    attributes = ["image_processor", "video_processor", "tokenizer"]
+    optional_attributes = [
+        "chat_template",
+        "time_mode",
+        "image_use_col_tokens",
+        "use_single_crop_col_tokens",
+        "use_single_crop_start_token",
+        "video_use_col_tokens",
+        "use_frame_special_tokens",
+    ]
+    image_processor_class = "AutoImageProcessor"
+    video_processor_class = "AutoVideoProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor: Molmo2ImageProcessor = None,
+        video_processor: Molmo2VideoProcessor = None,
+        tokenizer: AutoTokenizer = None,
+        chat_template: Optional[str] = None,
+        image_use_col_tokens: Optional[bool] = True,
+        use_single_crop_col_tokens: Optional[bool] = None,
+        use_single_crop_start_token: Optional[bool] = True,
+        video_use_col_tokens: Optional[bool] = False,
+        use_frame_special_tokens: Optional[bool] = True,
+        **kwargs
+    ) -> None:
+        super().__init__(
+            image_processor,
+            video_processor,
+            tokenizer,
+            chat_template=chat_template,
+            image_use_col_tokens=image_use_col_tokens,
+            use_single_crop_col_tokens=use_single_crop_col_tokens,
+            use_single_crop_start_token=use_single_crop_start_token,
+            video_use_col_tokens=video_use_col_tokens,
+            use_frame_special_tokens=use_frame_special_tokens,
+        )
+        self.image_placeholder_token = IMAGE_PROMPT
+        self.video_placeholder_token = VIDEO_PROMPT
+        self.image_token_ids = [
+            tokenizer.convert_tokens_to_ids(token)
+            for token in IMAGE_TOKENS
+        ]
+        self._patch_metadata = None
+
+    def get_image_tokens(self, image_grid: np.ndarray):
+        resized_h, resized_w, height, width = image_grid
+        per_row = np.full(width, IMAGE_PATCH_TOKEN)
+        if self.image_use_col_tokens:
+            per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)
+        joint = [
+            [IM_START_TOKEN],
+            np.tile(per_row, [height]),
+            [IM_END_TOKEN],
+        ]
+        per_row = np.full(resized_w, IMAGE_PATCH_TOKEN)
+        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 = (
+            LOW_RES_IMAGE_START_TOKEN
+            if self.use_single_crop_start_token
+            else IM_START_TOKEN
+        )
+        if use_single_crop_col_tokens:
+            per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)
+        joint = [
+            [image_start_token],
+            np.tile(per_row, [resized_h]),
+            [IM_END_TOKEN],
+        ] + joint
+
+        return np.concatenate(joint)
+    
+    def get_video_string(
+        self,
+        video_grid: np.ndarray,
+        timestamps: np.ndarray,
+    ):  
+        if self.use_frame_special_tokens:
+            start_token_id = FRAME_START_TOKEN
+            end_token_id = FRAME_END_TOKEN
+        else:
+            start_token_id = IM_START_TOKEN
+            end_token_id = IM_END_TOKEN
+        
+        num_frames, h, w = video_grid
+        video_string: str = ""
+        for frame_idx, frame_time in enumerate(timestamps):
+            # `per-frame-compact` time mode
+            prev_space = " " if frame_idx > 0 else ""
+            frame_prefix = prev_space + f"{frame_time:.1f} " # explicit whitespace before/after image tokens
+
+            video_string += frame_prefix
+            per_row = np.full(w, IMAGE_PATCH_TOKEN)
+            if self.video_use_col_tokens:
+                per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)
+            extra_tokens = np.tile(per_row, [h])
+            video_tokens = [
+                [start_token_id],
+                extra_tokens,
+                [end_token_id],
+            ]
+            video_string += "".join(np.concatenate(video_tokens, 0))
+
+        return video_string
+
+    def insert_bos(
+        self,
+        input_ids: np.ndarray,
+        attention_mask: np.ndarray,
+        bos_token_id: int,
+        pad_token_id: int,
+    ):
+        """
+        Args:
+            input_ids: [B, S] array with left padding
+            attention_mask: [B, S] array (0 for pad, 1 for valid)
+            bos_token_id: int
+            pad_token_id: int
+        Returns:
+            input_ids_out: [B, S] or [B, S+1] array with bos inserted if needed
+            attention_mask_out: same shape as input_ids_out
+        """
+
+        need_to_expand = len(input_ids.shape) == 1
+        if need_to_expand:
+            input_ids = input_ids[None, :]
+            attention_mask = attention_mask[None, :]
+
+        B, S = input_ids.shape
+
+        # Handle zero-length sequence
+        if S == 0:
+            new_input_ids = np.full((B, 1), bos_token_id, dtype=input_ids.dtype)
+            new_attention_mask = np.ones((B, 1), dtype=attention_mask.dtype)
+            if need_to_expand:
+                new_input_ids = new_input_ids[0]
+                new_attention_mask = new_attention_mask[0]
+            return new_input_ids, new_attention_mask
+
+        first_valid_index = (attention_mask == 1).argmax(axis=-1)  # [B]
+        bos_already_present = np.all(input_ids[np.arange(B), first_valid_index] == bos_token_id)
+
+        if bos_already_present:
+            if need_to_expand:
+                input_ids = input_ids[0]
+                attention_mask = attention_mask[0]
+            return input_ids, attention_mask
+        else:
+            new_input_ids = np.full((B, S+1), pad_token_id, dtype=input_ids.dtype)
+            new_attention_mask = np.zeros((B, S+1), dtype=attention_mask.dtype)
+
+            src_idx = np.tile(np.arange(S), (B, 1))  # [B, S]
+            valid_mask = src_idx >= first_valid_index[:, None]  # [B, S]
+            tgt_idx = src_idx + 1  # shit right
+            batch_idx = np.tile(np.arange(B)[:, None], (1, S))  # [B, S]
+
+            # flatten valid_positions
+            flat_vals = input_ids[valid_mask]
+            flat_batch = batch_idx[valid_mask]
+            flat_tgt = tgt_idx[valid_mask]
+
+            new_input_ids[flat_batch, flat_tgt] = flat_vals
+            new_attention_mask[flat_batch, flat_tgt] = 1
+            
+            insert_pos = first_valid_index
+            new_input_ids[np.arange(B), insert_pos] = bos_token_id
+            new_attention_mask[np.arange(B), insert_pos] = 1
+
+            if need_to_expand:
+                new_input_ids = new_input_ids[0]
+                new_attention_mask = new_attention_mask[0]
+
+            return new_input_ids, new_attention_mask
+
+    def __call__(
+        self,
+        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
+        images: ImageInput = None,
+        videos: VideoInput = None,
+        return_pointing_metadata: bool = False,
+        **kwargs: Unpack[Molmo2ProcessorKwargs],
+    ) -> BatchFeature:
+        """
+
+        Args:
+            text (`str`, `list[str]`, `list[list[str]]`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. Both channels-first and channels-last formats are supported.
+            videos (`dict[str, Any]` or `list[dict[str, Any]]`):
+                The video or batch of videos to be prepared. Each video can be a dictionary with the following keys:
+                - `"frames"`: `np.ndarray` of shape (T, H, W, 3)
+                - `"timestamps"`: `np.ndarray` of shape (T,)
+                - `"sampled_fps"`: `float` (optional)
+                - `"sampling_augmentation"`: `str` (optional)
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+        Returns:
+            `BatchFeature`: A [`BatchFeature`] with the following fields:
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+            - **image_token_pooling** -- Indices of the patches in `image_grids` to pool for each token in `image_tokens`.
+              Returned when `images` is not `None`.
+            - **image_grids** -- Grids of images. Returned when `images` is not `None`.
+            - **image_num_crops** -- Number of crops for each image. Returned when `images` is not `None`.
+            - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.
+            - **video_token_pooling** -- Indices of the patches in `video_grids` to pool for each token in `video_tokens`.
+              Returned when `videos` is not `None`.
+            - **video_grids** -- Grids of videos. Returned when `videos` is not `None`.
+        """
+
+        output_kwargs = self._merge_kwargs(
+            Molmo2ProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+        patch_metadata = {}
+        if images is not None:
+            image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"],
+                                                return_pointing_metadata=return_pointing_metadata)
+            if return_pointing_metadata:
+                patch_metadata["token_pooling"] = image_inputs.pop("image_token_pooling_np")
+                patch_metadata["subpatch_mapping"] = image_inputs.pop("subpatch_mapping")
+                patch_metadata["image_sizes"] = image_inputs.pop("image_sizes")
+            image_grids = image_inputs["image_grids"]
+        else:
+            image_inputs = {}
+            image_grids = None
+
+        if videos is not None:
+            videos_inputs = self.video_processor(
+                videos=videos, **output_kwargs["videos_kwargs"],
+                return_pointing_metadata=return_pointing_metadata
+            )
+            if return_pointing_metadata:
+                assert len(videos_inputs['video_metadata']) == 1
+                vd_metadata = videos_inputs['video_metadata'][0]
+                patch_metadata["token_pooling"] = videos_inputs.pop("video_token_pooling_np")
+                patch_metadata["subpatch_mapping"] = videos_inputs.pop("subpatch_mapping")
+                patch_metadata["timestamps"] = vd_metadata.timestamps
+                patch_metadata["video_size"] = (vd_metadata.width, vd_metadata.height)
+
+            video_grids = videos_inputs["video_grids"]
+            # If user has not requested video metadata, pop it
+            if "return_metadata" not in kwargs:
+                video_metadata = videos_inputs.pop("video_metadata")
+            else:
+                video_metadata = videos_inputs["video_metadata"]
+        else:
+            videos_inputs = {}
+            video_grids = None
+
+        if not isinstance(text, list):
+            text = [text]
+        
+        text = text.copy() # below lines change text in-place
+
+        if image_grids is not None:
+            index = 0
+            for i in range(len(text)):
+                num_images = text[i].count(self.image_placeholder_token)
+                image_grids_i = image_grids[index:index+num_images]
+                for image_grid in image_grids_i:
+                    image_tokens = self.get_image_tokens(image_grid)
+                    image_string = "".join(image_tokens)
+                    text[i] = text[i].replace(self.image_placeholder_token, image_string, 1)
+                index += num_images
+        
+        if video_grids is not None:
+            index = 0
+            for i in range(len(text)):
+                num_videos = text[i].count(self.video_placeholder_token)
+                assert num_videos in {0, 1}, "At most one video is supported for now"
+                video_grids_i = video_grids[index:index+num_videos]
+                metadata_i = video_metadata[index:index+num_videos]
+                for video_grid, metadata in zip(video_grids_i, metadata_i):
+                    video_string = self.get_video_string(
+                        video_grid,
+                        metadata.timestamps,
+                    )
+                    text[i] = text[i].replace(self.video_placeholder_token, video_string, 1)
+                index += num_videos
+
+        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+        return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)
+        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+        input_ids = text_inputs["input_ids"]
+        attention_mask = text_inputs["attention_mask"]
+
+        input_ids = np.array(input_ids)
+        attention_mask = np.array(attention_mask)
+        
+        bos = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id
+        input_ids, attention_mask = self.insert_bos(
+            input_ids, attention_mask, bos, self.tokenizer.pad_token_id
+        )
+
+        if return_mm_token_type_ids:
+            image_tokens = np.array(self.image_token_ids).astype(input_ids.dtype)
+            token_type_ids = np.any(input_ids[:, :, None] == image_tokens[None, None, :], axis=-1)
+            text_inputs["token_type_ids"] = token_type_ids.tolist()
+        
+        text_inputs["input_ids"] = input_ids.tolist()
+        text_inputs["attention_mask"] = attention_mask.tolist()
+
+        features = BatchFeature(
+            data={**text_inputs, **image_inputs, **videos_inputs},
+            tensor_type=return_tensors,
+        )
+        if return_pointing_metadata:
+            features["metadata"] = patch_metadata
+        return features
+
+    def post_process_image_text_to_text(
+        self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
+    ):
+        """
+        Post-process the output of the model to decode the text.
+
+        Args:
+            generated_outputs (`torch.Tensor` or `np.ndarray`):
+                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
+                or `(sequence_length,)`.
+            skip_special_tokens (`bool`, *optional*, defaults to `True`):
+                Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
+            clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+                Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
+            **kwargs:
+                Additional arguments to be passed to the tokenizer's `batch_decode method`.
+
+        Returns:
+            `list[str]`: The decoded text.
+        """
+        return self.tokenizer.batch_decode(
+            generated_outputs,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+
+Molmo2Processor.register_for_auto_class()

processing_molmo_point.py

@@ -0,0 +1,410 @@
+"""
+Processor class for Molmo2.
+"""
+from typing import Optional, Union
+import dataclasses
+
+import numpy as np
+
+from transformers.image_utils import ImageInput
+from transformers.video_utils import VideoInput
+from transformers.processing_utils import (
+    Unpack,
+    ProcessingKwargs,
+    ProcessorMixin,
+)
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.tokenization_utils_base import TextInput, PreTokenizedInput
+from transformers.utils import logging
+
+from transformers import AutoTokenizer
+from .image_processing_molmo2 import Molmo2ImagesKwargs, Molmo2ImageProcessor
+from .video_processing_molmo2 import Molmo2VideoProcessorKwargs, Molmo2VideoProcessor
+
+
+logger = logging.get_logger(__name__)
+
+
+# Special tokens, these should be present in any tokenizer we use since the preprocessor uses them
+IMAGE_PATCH_TOKEN = f"<im_patch>"  # Where to insert high-res tokens
+IMAGE_LOW_RES_TOKEN = f"<im_low>"  # Where to insert low-res tokens
+IM_START_TOKEN = f"<im_start>"
+LOW_RES_IMAGE_START_TOKEN = f"<low_res_im_start>"
+FRAME_START_TOKEN = f"<frame_start>"
+IM_END_TOKEN = f"<im_end>"
+FRAME_END_TOKEN= f"<frame_end>"
+IM_COL_TOKEN = f"<im_col>"
+IMAGE_PROMPT = "<|image|>"
+VIDEO_PROMPT = "<|video|>"
+
+IMAGE_TOKENS = [
+    IMAGE_PATCH_TOKEN,
+    IM_COL_TOKEN,
+    IM_START_TOKEN,
+    LOW_RES_IMAGE_START_TOKEN,
+    FRAME_START_TOKEN,
+    IM_END_TOKEN,
+    FRAME_END_TOKEN,
+    IMAGE_LOW_RES_TOKEN,
+]
+
+
+class MolmoPointProcessorKwargs(ProcessingKwargs, total=False):
+    """Molmo2 processor kwargs"""
+    images_kwargs: Molmo2ImagesKwargs
+    videos_kwargs: Molmo2VideoProcessorKwargs
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+            "return_mm_token_type_ids": True,
+        },
+        "videos_kwargs": {"return_metadata": True},
+    }
+
+
+class MolmoPointProcessor(ProcessorMixin):
+    attributes = ["image_processor", "video_processor", "tokenizer"]
+    optional_attributes = [
+        "chat_template",
+        "time_mode",
+        "image_use_col_tokens",
+        "use_single_crop_col_tokens",
+        "use_single_crop_start_token",
+        "video_use_col_tokens",
+        "use_frame_special_tokens",
+    ]
+    image_processor_class = "AutoImageProcessor"
+    video_processor_class = "AutoVideoProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor: Molmo2ImageProcessor = None,
+        video_processor: Molmo2VideoProcessor = None,
+        tokenizer: AutoTokenizer = None,
+        chat_template: Optional[str] = None,
+        image_use_col_tokens: Optional[bool] = True,
+        use_single_crop_col_tokens: Optional[bool] = None,
+        use_single_crop_start_token: Optional[bool] = True,
+        video_use_col_tokens: Optional[bool] = False,
+        use_frame_special_tokens: Optional[bool] = True,
+        **kwargs
+    ) -> None:
+        super().__init__(
+            image_processor,
+            video_processor,
+            tokenizer,
+            chat_template=chat_template,
+            image_use_col_tokens=image_use_col_tokens,
+            use_single_crop_col_tokens=use_single_crop_col_tokens,
+            use_single_crop_start_token=use_single_crop_start_token,
+            video_use_col_tokens=video_use_col_tokens,
+            use_frame_special_tokens=use_frame_special_tokens,
+        )
+
+        self.image_placeholder_token = IMAGE_PROMPT
+        self.video_placeholder_token = VIDEO_PROMPT
+        self.image_token_ids = [
+            tokenizer.convert_tokens_to_ids(token)
+            for token in IMAGE_TOKENS
+        ]
+
+    def get_image_tokens(self, image_grid: np.ndarray):
+        resized_h, resized_w, height, width = image_grid
+        per_row = np.full(width, IMAGE_PATCH_TOKEN)
+        if self.image_use_col_tokens:
+            per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)
+        joint = [
+            [IM_START_TOKEN],
+            np.tile(per_row, [height]),
+            [IM_END_TOKEN],
+        ]
+        per_row = np.full(resized_w, IMAGE_PATCH_TOKEN)
+        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 = (
+            LOW_RES_IMAGE_START_TOKEN
+            if self.use_single_crop_start_token
+            else IM_START_TOKEN
+        )
+        if use_single_crop_col_tokens:
+            per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)
+        joint = [
+            [image_start_token],
+            np.tile(per_row, [resized_h]),
+            [IM_END_TOKEN],
+        ] + joint
+
+        return np.concatenate(joint)
+    
+    def get_video_string(
+        self,
+        video_grid: np.ndarray,
+        timestamps: np.ndarray,
+    ):  
+        if self.use_frame_special_tokens:
+            start_token_id = FRAME_START_TOKEN
+            end_token_id = FRAME_END_TOKEN
+        else:
+            start_token_id = IM_START_TOKEN
+            end_token_id = IM_END_TOKEN
+        
+        num_frames, h, w = video_grid
+        video_string: str = ""
+        for frame_idx, frame_time in enumerate(timestamps):
+            # `per-frame-compact` time mode
+            prev_space = " " if frame_idx > 0 else ""
+            frame_prefix = prev_space + f"{frame_time:.1f} " # explicit whitespace before/after image tokens
+
+            video_string += frame_prefix
+            per_row = np.full(w, IMAGE_PATCH_TOKEN)
+            if self.video_use_col_tokens:
+                per_row = np.concatenate([per_row, [IM_COL_TOKEN]], 0)
+            extra_tokens = np.tile(per_row, [h])
+            video_tokens = [
+                [start_token_id],
+                extra_tokens,
+                [end_token_id],
+            ]
+            video_string += "".join(np.concatenate(video_tokens, 0))
+
+        return video_string
+
+    def insert_bos(
+        self,
+        input_ids: np.ndarray,
+        attention_mask: np.ndarray,
+        bos_token_id: int,
+        pad_token_id: int,
+    ):
+        """
+        Args:
+            input_ids: [B, S] array with left padding
+            attention_mask: [B, S] array (0 for pad, 1 for valid)
+            bos_token_id: int
+            pad_token_id: int
+        Returns:
+            input_ids_out: [B, S] or [B, S+1] array with bos inserted if needed
+            attention_mask_out: same shape as input_ids_out
+        """
+
+        need_to_expand = len(input_ids.shape) == 1
+        if need_to_expand:
+            input_ids = input_ids[None, :]
+            attention_mask = attention_mask[None, :]
+
+        B, S = input_ids.shape
+
+        # Handle zero-length sequence
+        if S == 0:
+            new_input_ids = np.full((B, 1), bos_token_id, dtype=input_ids.dtype)
+            new_attention_mask = np.ones((B, 1), dtype=attention_mask.dtype)
+            if need_to_expand:
+                new_input_ids = new_input_ids[0]
+                new_attention_mask = new_attention_mask[0]
+            return new_input_ids, new_attention_mask
+
+        first_valid_index = (attention_mask == 1).argmax(axis=-1)  # [B]
+        bos_already_present = np.all(input_ids[np.arange(B), first_valid_index] == bos_token_id)
+
+        if bos_already_present:
+            if need_to_expand:
+                input_ids = input_ids[0]
+                attention_mask = attention_mask[0]
+            return input_ids, attention_mask
+        else:
+            new_input_ids = np.full((B, S+1), pad_token_id, dtype=input_ids.dtype)
+            new_attention_mask = np.zeros((B, S+1), dtype=attention_mask.dtype)
+
+            src_idx = np.tile(np.arange(S), (B, 1))  # [B, S]
+            valid_mask = src_idx >= first_valid_index[:, None]  # [B, S]
+            tgt_idx = src_idx + 1  # shit right
+            batch_idx = np.tile(np.arange(B)[:, None], (1, S))  # [B, S]
+
+            # flatten valid_positions
+            flat_vals = input_ids[valid_mask]
+            flat_batch = batch_idx[valid_mask]
+            flat_tgt = tgt_idx[valid_mask]
+
+            new_input_ids[flat_batch, flat_tgt] = flat_vals
+            new_attention_mask[flat_batch, flat_tgt] = 1
+            
+            insert_pos = first_valid_index
+            new_input_ids[np.arange(B), insert_pos] = bos_token_id
+            new_attention_mask[np.arange(B), insert_pos] = 1
+
+            if need_to_expand:
+                new_input_ids = new_input_ids[0]
+                new_attention_mask = new_attention_mask[0]
+
+            return new_input_ids, new_attention_mask
+
+    def __call__(
+        self,
+        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
+        images: ImageInput = None,
+        videos: VideoInput = None,
+        return_subpatch_mapping: bool = False,
+        **kwargs: Unpack[MolmoPointProcessorKwargs],
+    ) -> BatchFeature:
+        """
+
+        Args:
+            text (`str`, `list[str]`, `list[list[str]]`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. Both channels-first and channels-last formats are supported.
+            videos (`dict[str, Any]` or `list[dict[str, Any]]`):
+                The video or batch of videos to be prepared. Each video can be a dictionary with the following keys:
+                - `"frames"`: `np.ndarray` of shape (T, H, W, 3)
+                - `"timestamps"`: `np.ndarray` of shape (T,)
+                - `"sampled_fps"`: `float` (optional)
+                - `"sampling_augmentation"`: `str` (optional)
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+        Returns:
+            `BatchFeature`: A [`BatchFeature`] with the following fields:
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+            - **image_token_pooling** -- Indices of the patches in `image_grids` to pool for each token in `image_tokens`.
+              Returned when `images` is not `None`.
+            - **image_grids** -- Grids of images. Returned when `images` is not `None`.
+            - **image_num_crops** -- Number of crops for each image. Returned when `images` is not `None`.
+            - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.
+            - **video_token_pooling** -- Indices of the patches in `video_grids` to pool for each token in `video_tokens`.
+              Returned when `videos` is not `None`.
+            - **video_grids** -- Grids of videos. Returned when `videos` is not `None`.
+        """
+
+        output_kwargs = self._merge_kwargs(
+            MolmoPointProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+
+        subpatch_mapping = None
+        if images is not None:
+            if return_subpatch_mapping:
+                image_inputs, subpatch_mapping = self.image_processor(images, **output_kwargs["images_kwargs"], return_subpatch_mapping=return_subpatch_mapping)
+            else:
+                image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"])
+            image_grids = image_inputs["image_grids"]
+        else:
+            image_inputs = {}
+            image_grids = None
+
+        if videos is not None:
+            videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])
+            video_grids = videos_inputs["video_grids"]
+            # If user has not requested video metadata, pop it
+            if "return_metadata" not in kwargs:
+                video_metadata = videos_inputs.pop("video_metadata")
+            else:
+                video_metadata = videos_inputs["video_metadata"]
+        else:
+            videos_inputs = {}
+            video_grids = None
+
+        if not isinstance(text, list):
+            text = [text]
+        
+        text = text.copy() # below lines change text in-place
+
+        if image_grids is not None:
+            index = 0
+            for i in range(len(text)):
+                num_images = text[i].count(self.image_placeholder_token)
+                image_grids_i = image_grids[index:index+num_images]
+                for image_grid in image_grids_i:
+                    image_tokens = self.get_image_tokens(image_grid)
+                    image_string = "".join(image_tokens)
+                    text[i] = text[i].replace(self.image_placeholder_token, image_string, 1)
+                index += num_images
+        
+        if video_grids is not None:
+            index = 0
+            for i in range(len(text)):
+                num_videos = text[i].count(self.video_placeholder_token)
+                assert num_videos in {0, 1}, "At most one video is supported for now"
+                video_grids_i = video_grids[index:index+num_videos]
+                metadata_i = video_metadata[index:index+num_videos]
+                for video_grid, metadata in zip(video_grids_i, metadata_i):
+                    video_string = self.get_video_string(
+                        video_grid,
+                        metadata.timestamps,
+                    )
+                    text[i] = text[i].replace(self.video_placeholder_token, video_string, 1)
+                index += num_videos
+
+        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+        return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)
+        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+        input_ids = text_inputs["input_ids"]
+        attention_mask = text_inputs["attention_mask"]
+
+        input_ids = np.array(input_ids)
+        attention_mask = np.array(attention_mask)
+        
+        bos = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id
+        input_ids, attention_mask = self.insert_bos(
+            input_ids, attention_mask, bos, self.tokenizer.pad_token_id
+        )
+
+        if return_mm_token_type_ids:
+            image_tokens = np.array(self.image_token_ids).astype(input_ids.dtype)
+            token_type_ids = np.any(input_ids[:, :, None] == image_tokens[None, None, :], axis=-1)
+            text_inputs["token_type_ids"] = token_type_ids.tolist()
+        
+        text_inputs["input_ids"] = input_ids.tolist()
+        text_inputs["attention_mask"] = attention_mask.tolist()
+        features = BatchFeature(
+            data={**text_inputs, **image_inputs, **videos_inputs},
+            tensor_type=return_tensors,
+        )
+        if return_subpatch_mapping:
+            return features, subpatch_mapping
+        return features
+
+    def post_process_image_text_to_text(
+        self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
+    ):
+        """
+        Post-process the output of the model to decode the text.
+
+        Args:
+            generated_outputs (`torch.Tensor` or `np.ndarray`):
+                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
+                or `(sequence_length,)`.
+            skip_special_tokens (`bool`, *optional*, defaults to `True`):
+                Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
+            clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+                Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
+            **kwargs:
+                Additional arguments to be passed to the tokenizer's `batch_decode method`.
+
+        Returns:
+            `list[str]`: The decoded text.
+        """
+        return self.tokenizer.batch_decode(
+            generated_outputs,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+
+MolmoPointProcessor.register_for_auto_class()

processor_config.json

@@ -0,0 +1,11 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_molmo2.Molmo2Processor"
+  },
+  "image_use_col_tokens": true,
+  "processor_class": "Molmo2Processor",
+  "use_frame_special_tokens": true,
+  "use_single_crop_col_tokens": false,
+  "use_single_crop_start_token": true,
+  "video_use_col_tokens": false
+}

special_tokens_map.json

@@ -0,0 +1,300 @@
+{
+  "additional_special_tokens": [
+    "|<EXTRA_TOKENS_0>|",
+    "|<EXTRA_TOKENS_1>|",
+    "|<EXTRA_TOKENS_2>|",
+    "|<EXTRA_TOKENS_3>|",
+    "|<EXTRA_TOKENS_4>|",
+    "|<EXTRA_TOKENS_5>|",
+    "|<EXTRA_TOKENS_6>|",
+    "|<EXTRA_TOKENS_7>|",
+    "|<EXTRA_TOKENS_8>|",
+    "|<EXTRA_TOKENS_9>|",
+    "|<EXTRA_TOKENS_10>|",
+    "|<EXTRA_TOKENS_11>|",
+    "|<EXTRA_TOKENS_12>|",
+    "|<EXTRA_TOKENS_13>|",
+    "|<EXTRA_TOKENS_14>|",
+    "|<EXTRA_TOKENS_15>|",
+    "|<EXTRA_TOKENS_16>|",
+    "|<EXTRA_TOKENS_17>|",
+    "|<EXTRA_TOKENS_18>|",
+    "|<EXTRA_TOKENS_19>|",
+    "|<EXTRA_TOKENS_20>|",
+    "|<EXTRA_TOKENS_21>|",
+    "|<EXTRA_TOKENS_22>|",
+    "|<EXTRA_TOKENS_23>|",
+    "|<EXTRA_TOKENS_24>|",
+    "|<EXTRA_TOKENS_25>|",
+    "|<EXTRA_TOKENS_26>|",
+    "|<EXTRA_TOKENS_27>|",
+    "|<EXTRA_TOKENS_28>|",
+    "|<EXTRA_TOKENS_29>|",
+    "|<EXTRA_TOKENS_30>|",
+    "|<EXTRA_TOKENS_31>|",
+    "|<EXTRA_TOKENS_32>|",
+    "|<EXTRA_TOKENS_33>|",
+    "|<EXTRA_TOKENS_34>|",
+    "|<EXTRA_TOKENS_35>|",
+    "|<EXTRA_TOKENS_36>|",
+    "|<EXTRA_TOKENS_37>|",
+    "|<EXTRA_TOKENS_38>|",
+    "|<EXTRA_TOKENS_39>|",
+    "|<EXTRA_TOKENS_40>|",
+    "|<EXTRA_TOKENS_41>|",
+    "|<EXTRA_TOKENS_42>|",
+    "|<EXTRA_TOKENS_43>|",
+    "|<EXTRA_TOKENS_44>|",
+    "|<EXTRA_TOKENS_45>|",
+    "|<EXTRA_TOKENS_46>|",
+    "|<EXTRA_TOKENS_47>|",
+    "|<EXTRA_TOKENS_48>|",
+    "|<EXTRA_TOKENS_49>|",
+    "|<EXTRA_TOKENS_50>|",
+    "|<EXTRA_TOKENS_51>|",
+    "|<EXTRA_TOKENS_52>|",
+    "|<EXTRA_TOKENS_53>|",
+    "|<EXTRA_TOKENS_54>|",
+    "|<EXTRA_TOKENS_55>|",
+    "|<EXTRA_TOKENS_56>|",
+    "|<EXTRA_TOKENS_57>|",
+    "|<EXTRA_TOKENS_58>|",
+    "|<EXTRA_TOKENS_59>|",
+    "|<EXTRA_TOKENS_60>|",
+    "|<EXTRA_TOKENS_61>|",
+    "|<EXTRA_TOKENS_62>|",
+    "|<EXTRA_TOKENS_63>|",
+    "|<EXTRA_TOKENS_64>|",
+    "|<EXTRA_TOKENS_65>|",
+    "|<EXTRA_TOKENS_66>|",
+    "|<EXTRA_TOKENS_67>|",
+    "|<EXTRA_TOKENS_68>|",
+    "|<EXTRA_TOKENS_69>|",
+    "|<EXTRA_TOKENS_70>|",
+    "|<EXTRA_TOKENS_71>|",
+    "|<EXTRA_TOKENS_72>|",
+    "|<EXTRA_TOKENS_73>|",
+    "|<EXTRA_TOKENS_74>|",
+    "|<EXTRA_TOKENS_75>|",
+    "|<EXTRA_TOKENS_76>|",
+    "|<EXTRA_TOKENS_77>|",
+    "|<EXTRA_TOKENS_78>|",
+    "|<EXTRA_TOKENS_79>|",
+    "|<EXTRA_TOKENS_80>|",
+    "|<EXTRA_TOKENS_81>|",
+    "|<EXTRA_TOKENS_82>|",
+    "|<EXTRA_TOKENS_83>|",
+    "|<EXTRA_TOKENS_84>|",
+    "|<EXTRA_TOKENS_85>|",
+    "|<EXTRA_TOKENS_86>|",
+    "|<EXTRA_TOKENS_87>|",
+    "|<EXTRA_TOKENS_88>|",
+    "|<EXTRA_TOKENS_89>|",
+    "|<EXTRA_TOKENS_90>|",
+    "|<EXTRA_TOKENS_91>|",
+    "|<EXTRA_TOKENS_92>|",
+    "|<EXTRA_TOKENS_93>|",
+    "|<EXTRA_TOKENS_94>|",
+    "|<EXTRA_TOKENS_95>|",
+    "|<EXTRA_TOKENS_96>|",
+    "|<EXTRA_TOKENS_97>|",
+    "|<EXTRA_TOKENS_98>|",
+    "|<EXTRA_TOKENS_99>|",
+    "|<EXTRA_TOKENS_100>|",
+    "|<EXTRA_TOKENS_101>|",
+    "|<EXTRA_TOKENS_102>|",
+    "|<EXTRA_TOKENS_103>|",
+    "|<EXTRA_TOKENS_104>|",
+    "|<EXTRA_TOKENS_105>|",
+    "|<EXTRA_TOKENS_106>|",
+    "|<EXTRA_TOKENS_107>|",
+    "|<EXTRA_TOKENS_108>|",
+    "|<EXTRA_TOKENS_109>|",
+    "|<EXTRA_TOKENS_110>|",
+    "|<EXTRA_TOKENS_111>|",
+    "|<EXTRA_TOKENS_112>|",
+    "|<EXTRA_TOKENS_113>|",
+    "|<EXTRA_TOKENS_114>|",
+    "|<EXTRA_TOKENS_115>|",
+    "|<EXTRA_TOKENS_116>|",
+    "|<EXTRA_TOKENS_117>|",
+    "|<EXTRA_TOKENS_118>|",
+    "|<EXTRA_TOKENS_119>|",
+    "|<EXTRA_TOKENS_120>|",
+    "|<EXTRA_TOKENS_121>|",
+    "|<EXTRA_TOKENS_122>|",
+    "|<EXTRA_TOKENS_123>|",
+    "|<EXTRA_TOKENS_124>|",
+    "|<EXTRA_TOKENS_125>|",
+    "|<EXTRA_TOKENS_126>|",
+    "|<EXTRA_TOKENS_127>|",
+    "|<EXTRA_TOKENS_128>|",
+    "|<EXTRA_TOKENS_129>|",
+    "|<EXTRA_TOKENS_130>|",
+    "|<EXTRA_TOKENS_131>|",
+    "|<EXTRA_TOKENS_132>|",
+    "|<EXTRA_TOKENS_133>|",
+    "|<EXTRA_TOKENS_134>|",
+    "|<EXTRA_TOKENS_135>|",
+    "|<EXTRA_TOKENS_136>|",
+    "|<EXTRA_TOKENS_137>|",
+    "|<EXTRA_TOKENS_138>|",
+    "|<EXTRA_TOKENS_139>|",
+    "|<EXTRA_TOKENS_140>|",
+    "|<EXTRA_TOKENS_141>|",
+    "|<EXTRA_TOKENS_142>|",
+    "|<EXTRA_TOKENS_143>|",
+    "|<EXTRA_TOKENS_144>|",
+    "|<EXTRA_TOKENS_145>|",
+    "|<EXTRA_TOKENS_146>|",
+    "|<EXTRA_TOKENS_147>|",
+    "|<EXTRA_TOKENS_148>|",
+    "|<EXTRA_TOKENS_149>|",
+    "|<EXTRA_TOKENS_150>|",
+    "|<EXTRA_TOKENS_151>|",
+    "|<EXTRA_TOKENS_152>|",
+    "|<EXTRA_TOKENS_153>|",
+    "|<EXTRA_TOKENS_154>|",
+    "|<EXTRA_TOKENS_155>|",
+    "|<EXTRA_TOKENS_156>|",
+    "|<EXTRA_TOKENS_157>|",
+    "|<EXTRA_TOKENS_158>|",
+    "|<EXTRA_TOKENS_159>|",
+    "|<EXTRA_TOKENS_160>|",
+    "|<EXTRA_TOKENS_161>|",
+    "|<EXTRA_TOKENS_162>|",
+    "|<EXTRA_TOKENS_163>|",
+    "|<EXTRA_TOKENS_164>|",
+    "|<EXTRA_TOKENS_165>|",
+    "|<EXTRA_TOKENS_166>|",
+    "|<EXTRA_TOKENS_167>|",
+    "|<EXTRA_TOKENS_168>|",
+    "|<EXTRA_TOKENS_169>|",
+    "|<EXTRA_TOKENS_170>|",
+    "|<EXTRA_TOKENS_171>|",
+    "|<EXTRA_TOKENS_172>|",
+    "|<EXTRA_TOKENS_173>|",
+    "|<EXTRA_TOKENS_174>|",
+    "|<EXTRA_TOKENS_175>|",
+    "|<EXTRA_TOKENS_176>|",
+    "|<EXTRA_TOKENS_177>|",
+    "|<EXTRA_TOKENS_178>|",
+    "|<EXTRA_TOKENS_179>|",
+    "|<EXTRA_TOKENS_180>|",
+    "|<EXTRA_TOKENS_181>|",
+    "|<EXTRA_TOKENS_182>|",
+    "|<EXTRA_TOKENS_183>|",
+    "|<EXTRA_TOKENS_184>|",
+    "|<EXTRA_TOKENS_185>|",
+    "|<EXTRA_TOKENS_186>|",
+    "|<EXTRA_TOKENS_187>|",
+    "|<EXTRA_TOKENS_188>|",
+    "|<EXTRA_TOKENS_189>|",
+    "|<EXTRA_TOKENS_190>|",
+    "|<EXTRA_TOKENS_191>|",
+    "|<EXTRA_TOKENS_192>|",
+    "|<EXTRA_TOKENS_193>|",
+    "|<EXTRA_TOKENS_194>|",
+    "|<EXTRA_TOKENS_195>|",
+    "|<EXTRA_TOKENS_196>|",
+    "|<EXTRA_TOKENS_197>|",
+    "|<EXTRA_TOKENS_198>|",
+    "|<EXTRA_TOKENS_199>|",
+    "|<EXTRA_TOKENS_200>|",
+    "|<EXTRA_TOKENS_201>|",
+    "|<EXTRA_TOKENS_202>|",
+    "|<EXTRA_TOKENS_203>|",
+    "|<EXTRA_TOKENS_204>|",
+    "|<EXTRA_TOKENS_205>|",
+    "|<EXTRA_TOKENS_206>|",
+    "|<EXTRA_TOKENS_207>|",
+    "|<EXTRA_TOKENS_208>|",
+    "|<EXTRA_TOKENS_209>|",
+    "|<EXTRA_TOKENS_210>|",
+    "|<EXTRA_TOKENS_211>|",
+    "|<EXTRA_TOKENS_212>|",
+    "|<EXTRA_TOKENS_213>|",
+    "|<EXTRA_TOKENS_214>|",
+    "|<EXTRA_TOKENS_215>|",
+    "|<EXTRA_TOKENS_216>|",
+    "|<EXTRA_TOKENS_217>|",
+    "|<EXTRA_TOKENS_218>|",
+    "|<EXTRA_TOKENS_219>|",
+    "|<EXTRA_TOKENS_220>|",
+    "|<EXTRA_TOKENS_221>|",
+    "|<EXTRA_TOKENS_222>|",
+    "|<EXTRA_TOKENS_223>|",
+    "|<EXTRA_TOKENS_224>|",
+    "|<EXTRA_TOKENS_225>|",
+    "|<EXTRA_TOKENS_226>|",
+    "|<EXTRA_TOKENS_227>|",
+    "|<EXTRA_TOKENS_228>|",
+    "|<EXTRA_TOKENS_229>|",
+    "|<EXTRA_TOKENS_230>|",
+    "|<EXTRA_TOKENS_231>|",
+    "|<EXTRA_TOKENS_232>|",
+    "|<EXTRA_TOKENS_233>|",
+    "|<EXTRA_TOKENS_234>|",
+    "|<EXTRA_TOKENS_235>|",
+    "|<EXTRA_TOKENS_236>|",
+    "|<EXTRA_TOKENS_237>|",
+    "|<EXTRA_TOKENS_238>|",
+    "|<EXTRA_TOKENS_239>|",
+    "|<EXTRA_TOKENS_240>|",
+    "|<EXTRA_TOKENS_241>|",
+    "|<EXTRA_TOKENS_242>|",
+    "|<EXTRA_TOKENS_243>|",
+    "|<EXTRA_TOKENS_244>|",
+    "|<EXTRA_TOKENS_245>|",
+    "|<EXTRA_TOKENS_246>|",
+    "|<EXTRA_TOKENS_247>|",
+    "|<EXTRA_TOKENS_248>|",
+    "|<EXTRA_TOKENS_249>|",
+    "|<EXTRA_TOKENS_250>|",
+    "|<EXTRA_TOKENS_251>|",
+    "|<EXTRA_TOKENS_252>|",
+    "|<EXTRA_TOKENS_253>|",
+    "|<EXTRA_TOKENS_254>|",
+    "|<EXTRA_TOKENS_255>|",
+    "|<EXTRA_TOKENS_256>|",
+    "|<EXTRA_TOKENS_257>|",
+    "|<EXTRA_TOKENS_258>|",
+    "|<EXTRA_TOKENS_259>|",
+    "|<EXTRA_TOKENS_260>|",
+    "|<EXTRA_TOKENS_261>|",
+    "|<EXTRA_TOKENS_262>|",
+    "|<EXTRA_TOKENS_263>|",
+    "|<EXTRA_TOKENS_264>|",
+    "|<EXTRA_TOKENS_265>|",
+    "|<EXTRA_TOKENS_266>|",
+    "<im_start>",
+    "<im_end>",
+    "<im_patch>",
+    "<im_col>",
+    "<low_res_im_start>",
+    "<|image|>",
+    "<im_low>",
+    "<frame_start>",
+    "<frame_end>",
+    "<|video|>",
+    "<|points|>",
+    "<|token_index|>",
+    "<|vit_index|>",
+    "<|vit_loc|>"
+  ],
+  "bos_token": "<|im_end|>",
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

test_molmo2.py

@@ -0,0 +1,133 @@
+import argparse
+from PIL import Image
+import requests
+
+import torch
+
+from transformers import AutoProcessor, AutoModelForImageTextToText
+
+
+video_path = "https://storage.googleapis.com/oe-training-public/demo_videos/many_penguins.mp4"
+image1_path = "https://picsum.photos/id/237/536/354"
+image2_path = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Test Molmo2 HF-compatible model."
+    )
+    parser.add_argument("checkpoint_dir", help="Location of Molmo2 checkpoint.")
+    args = parser.parse_args()
+
+    processor = AutoProcessor.from_pretrained(
+        args.checkpoint_dir,
+        trust_remote_code=True,
+        dtype="auto",
+        device_map="auto",
+        padding_side="left",
+    )
+
+    model = AutoModelForImageTextToText.from_pretrained(
+        args.checkpoint_dir,
+        trust_remote_code=True,
+        dtype="auto",
+        device_map="auto",
+    )
+
+    single_image_messages = [
+        {
+            "role": "user",
+            "content": [
+                dict(type="text", text="Describe this image."),
+                dict(type="image", image=Image.open(requests.get(image1_path, stream=True).raw)),
+            ]
+        }
+    ]
+
+    multi_image_messages = [
+        {
+            "role": "user",
+            "content": [
+                dict(type="text", text="Compare these images."),
+                dict(
+                    type="image",
+                    image=Image.open(requests.get(image1_path, stream=True).raw),
+                ),
+                dict(
+                    type="image",
+                    image=Image.open(requests.get(image2_path, stream=True).raw),
+                ),
+            ],
+        }
+    ]
+
+    video_messages = [
+        {
+            "role": "user",
+            "content": [
+                dict(type="text", text="Which animal appears in the video?"),
+                dict(type="video", video=video_path),
+            ]
+        }
+    ]
+
+    single_image_inputs = processor.apply_chat_template(
+        single_image_messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+    )
+
+    single_image_inputs = {k: v.to(model.device) for k, v in single_image_inputs.items()}
+
+    with torch.inference_mode():
+        with torch.autocast("cuda", enabled=True, dtype=torch.bfloat16):
+            single_image_generated_ids = model.generate(**single_image_inputs, max_new_tokens=448)
+    single_image_generated_tokens = single_image_generated_ids[:, single_image_inputs['input_ids'].size(1):]
+    single_image_generated_text = processor.post_process_image_text_to_text(
+        single_image_generated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
+    print(single_image_generated_text)
+
+    multi_image_inputs = processor.apply_chat_template(
+        multi_image_messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+    )
+
+    multi_image_inputs = {k: v.to(model.device) for k, v in multi_image_inputs.items()}
+
+    with torch.inference_mode():
+        with torch.autocast("cuda", enabled=True, dtype=torch.bfloat16):
+            multi_image_generated_ids = model.generate(**multi_image_inputs, max_new_tokens=448)
+    multi_image_generated_tokens = multi_image_generated_ids[:, multi_image_inputs['input_ids'].size(1):]
+    multi_image_generated_text = processor.post_process_image_text_to_text(
+        multi_image_generated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
+    print(multi_image_generated_text)
+
+    video_inputs = processor.apply_chat_template(
+        video_messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+    )
+
+    video_inputs = {k: v.to(model.device) for k, v in video_inputs.items()}
+    
+    with torch.inference_mode():
+        with torch.autocast("cuda", enabled=True, dtype=torch.bfloat16):
+            video_generated_ids = model.generate(**video_inputs, max_new_tokens=2048)
+    video_generated_tokens = video_generated_ids[:, video_inputs['input_ids'].size(1):]
+    video_generated_text = processor.post_process_image_text_to_text(
+        video_generated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
+    print(video_generated_text)
+
+
+if __name__ == "__main__":
+    main()

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f437033cf8ca3315943460f7b7681d01130795107d9a99dc124fd9d6898e932
+size 17417468

unified_demo.py

@@ -0,0 +1,334 @@
+import functools
+import os
+import argparse
+import logging
+from collections import defaultdict
+from PIL import Image, ImageFile, ImageDraw
+import PIL
+
+import numpy as np
+import torch
+from transformers import AutoProcessor, AutoModelForImageTextToText
+
+from olmo.models.video_olmo.video_olmo import VideoOlmoConfig
+from olmo.html_utils import postprocess_prompt
+from olmo.util import (
+    prepare_cli_environment,
+    resource_path,
+)
+
+import gradio as gr
+
+try:
+    from molmo_utils import process_vision_info
+except ImportError:
+    # raise ImportError("molmo_utils not found. Please install it with `pip install molmo-utils`.")
+    pass
+
+
+Image.MAX_IMAGE_PIXELS = None
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+CACHE = "model_cache"
+log = logging.getLogger(__name__)
+ALLOWED_PATH = [CACHE]
+MAX_IMAGE_SIZE = 512
+MAX_VIDEO_HEIGHT = 512
+POINT_SIZE = 0.01
+
+DEVICE = None
+
+# load the model, processor
+MODEL = None
+PROCESSOR = None
+POINT_FORMATTER = None
+
+
+def draw_points(image, points):
+    if isinstance(image, np.ndarray):
+        annotation = PIL.Image.fromarray(image)
+    else:
+        annotation = image.copy()
+    draw = ImageDraw.Draw(annotation)
+    w, h = annotation.size
+    size = max(5, int(max(w, h) * POINT_SIZE))
+    for x, y in points:
+        draw.ellipse((x-size, y-size, x+size, y+size), fill="rgb(240, 82, 156)", outline=None)
+    return annotation
+
+
+def get_message(
+    images: list[Image.Image] | None,
+    video_path: str | None,
+    max_frames: int,
+    frame_sample_mode: str,
+    max_fps: int | None,
+    sampling_fps: int | None,
+    input_text: str,
+    style: str,
+):
+    content = [
+        dict(type="text", text=input_text, stye=style)
+    ]
+    if images:
+        image_content = [
+            dict(type="image", image=image)
+            for image in images
+        ]
+        content.extend(image_content)
+    if video_path:
+        video_kwargs = {
+            "num_frames": max_frames,
+            "frame_sample_mode": frame_sample_mode,
+        }
+        if max_fps is not None:
+            video_kwargs["max_fps"] = max_fps
+        if sampling_fps is not None:
+            video_kwargs["sampling_fps"] = sampling_fps
+        video_content = dict(type="video", video=video_path, **video_kwargs)
+        content.append(video_content)
+    
+    return [
+        {
+            "role": "user",
+            "content": content,
+        }
+    ]
+
+
+def cast_float_dtype(t: torch.Tensor):
+    if torch.is_floating_point(t):
+        t = t.to(torch.bfloat16)
+    return t
+
+
+def run_single_inference(*inputs, annotations=None):
+    video_path, images, input_text, style, frame_sample_mode, max_frames, max_fps, sampling_fps, max_steps = inputs
+    assert images is not None or video_path is not None, "Either images or video_path must be provided"
+    assert images is None or video_path is None, "Both images and video_path cannot be provided at the same time"
+    nimages = 0
+    if images:
+        images = [t[0] for t in images]
+        nimages = len(images)
+        logging.info(f"# of images: {nimages}")
+    
+    messages = get_message(
+        images=images,
+        video_path=video_path,
+        max_frames=max_frames,
+        frame_sample_mode=frame_sample_mode,
+        max_fps=max_fps,
+        sampling_fps=sampling_fps,
+        input_text=input_text,
+        style=style,
+    )
+    images, videos, video_kwargs = process_vision_info(messages)
+    if videos:
+        videos, video_metadatas = zip(*videos)
+        videos, video_metadatas = list(videos), list(video_metadatas)
+        logging.info(
+            f"Videos: {videos[0].shape}, frame_sample_mode: {frame_sample_mode}, "
+            f"max_frames: {max_frames}, max_fps: {max_fps}, sampling_fps: {sampling_fps}"
+        )
+    else:
+        video_metadatas = None
+    logging.info(f"Running inference for prompt: \"{input_text}\", style={style} steps={max_steps}")
+    text = PROCESSOR.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    inputs = PROCESSOR(
+        images=images,
+        videos=videos,
+        video_metadata=video_metadatas,
+        text=text,
+        padding=True,
+        return_tensors="pt",
+        **video_kwargs,
+    )
+
+    if MODEL.config.dtype == torch.bfloat16:
+        inputs = {k: cast_float_dtype(v.to(DEVICE)) for k, v in inputs.items()}
+    else:
+        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+    with torch.inference_mode():
+        if MODEL.config.dtype == torch.bfloat16:
+            output = MODEL.generate(**inputs, max_new_tokens=max_steps)
+        else:
+            with torch.autocast("cuda", enabled=True, dtype=torch.bfloat16):
+                output = MODEL.generate(**inputs, max_new_tokens=max_steps)
+    prompts = output[0, :inputs['input_ids'].size(1)]
+    prompt_text = PROCESSOR.decode(prompts, skip_special_tokens=False)
+    prompt_text = postprocess_prompt(prompt_text)
+    logging.info(f"hf prompt: {prompt_text}")
+    generated_tokens = output[0, inputs['input_ids'].size(1):]
+    generated_text = PROCESSOR.decode(generated_tokens, skip_special_tokens=True)
+    logging.info(f"hf generated_text: {generated_text}")
+    if annotations:
+        if video_path is None and nimages == 1:
+            w, h = images[0].size
+            points = POINT_FORMATTER.extract_points(generated_text, w, h)
+            if points:
+                return generated_text, [draw_points(images[0], points)]
+            else:
+                return generated_text, []
+        elif video_path is None and nimages > 1:
+            w, h = [x.size[0] for x in images], [x.size[1] for x in images]
+            points = POINT_FORMATTER.extract_multi_image_points(generated_text, w, h)
+            if points:
+                group_by_index = defaultdict(list)
+                for ix, x, y in points:
+                    group_by_index[ix].append((x, y))
+                out = []
+                for ix, points in group_by_index.items():
+                    out.append(draw_points(images[ix-1], points))
+                return generated_text, out
+            else:
+                return generated_text, []
+        else:
+            h, w = videos[0].shape[1:3]
+            group_by_time = defaultdict(list)
+            points = POINT_FORMATTER.extract_multi_image_points(generated_text, w, h)
+            if points:
+                for ts, x, y in points:
+                    group_by_time[ts].append((x, y))
+            else:
+                track = POINT_FORMATTER.extract_trajectories(generated_text, w, h, 30)
+                for ex in track:
+                    group_by_time[ex["time"]] = [(x["x"], x["y"]) for x in ex["points"]]
+            grouped_by_frame = defaultdict(list)
+            for ts, points in group_by_time.items():
+                timestamps = video_metadatas[0]["frames_indices"] / video_metadatas[0]["fps"]
+                ix = int(np.argmin(np.abs(timestamps - ts)))
+                grouped_by_frame[ix] += points
+            out = []
+            for ix, points in grouped_by_frame.items():
+                out.append(draw_points(videos[0][ix], points))
+            return generated_text, out
+    return generated_text
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("ckpt_home", type=str)
+    parser.add_argument("--server_name")
+    parser.add_argument("--default_max_tokens", type=int, default=2048)
+    parser.add_argument("--cloudflare_tunnel", action="store_true")
+    parser.add_argument("--original_ckpt_home", type=str, default=None)
+    parser.add_argument("--annotations", action="store_true")
+    parser.add_argument("--no_share", action="store_true")
+    parser.add_argument("--port", type=int, default=7860)
+    args = parser.parse_args()
+
+    prepare_cli_environment()
+
+    global DEVICE, MODEL, PROCESSOR
+    if torch.cuda.is_available():
+        DEVICE = torch.device("cuda")
+    else:
+        logging.warning("No GPU available, using CPU")
+        DEVICE = torch.device("cpu")
+        if MODEL is not None:
+            MODEL.to(DEVICE)
+    
+    MODEL = AutoModelForImageTextToText.from_pretrained(
+        args.ckpt_home,
+        trust_remote_code=True,
+        dtype="auto",
+        device_map="auto",
+    )
+    
+    PROCESSOR = AutoProcessor.from_pretrained(
+        args.ckpt_home,
+        trust_remote_code=True,
+        dtype="auto",
+        device_map="auto",
+        padding_side="left",
+    )
+
+    if args.annotations:
+        assert args.original_ckpt_home is not None, "original_ckpt_home must be provided when annotations are enabled"
+        global POINT_FORMATTER
+        model_cfg_path = resource_path(args.original_ckpt_home, "config.yaml")
+        model_cfg = VideoOlmoConfig.load(model_cfg_path, key="model", validate_paths=False)
+        preprocessor = model_cfg.build_preprocessor(for_inference=True, is_training=False)
+        POINT_FORMATTER = preprocessor.formatter._point_formatter
+
+    CSS = """
+    #input_image image {
+        object-fit: contain !important;
+    }
+    #input_video video {
+        object-fit: contain !important;
+    }
+    """
+
+    frame_sample_mode = PROCESSOR.video_processor.frame_sample_mode
+    max_frames = PROCESSOR.video_processor.num_frames
+    max_fps = PROCESSOR.video_processor.max_fps
+    sampling_fps = PROCESSOR.video_processor.sampling_fps
+
+    with gr.Blocks(css=CSS) as demo:
+        gr.Markdown(
+            f"""
+            ## Molmo2 Demo
+            Provide either a video or images and a prompt for question answering.
+            """
+        )
+        with gr.Row():
+            with gr.Tabs():
+                with gr.TabItem("video"):
+                    video = gr.Video(label="Input Video", elem_id="input_video", height=MAX_VIDEO_HEIGHT)
+                with gr.TabItem("image(s)"):
+                    images = gr.Gallery(label="Input Images", elem_id="input_image", type="pil", height=MAX_IMAGE_SIZE)
+        
+        with gr.Row():
+            input_text = gr.Textbox(placeholder="Enter the prompt", label="Input text")
+        
+        with gr.Row():
+            style = gr.Textbox(value="demo", label="style")
+            frame_sample_mode = gr.Textbox(value=frame_sample_mode, label="frame_sample_mode")
+            max_frames = gr.Number(value=max_frames, label="max_frames")
+            max_fps = gr.Number(value=max_fps, label="max_fps")
+            sampling_fps = gr.Number(value=sampling_fps, label="sampling_fps")
+            max_tok_slider = gr.Slider(label="max_tokens", minimum=1, maximum=4096, step=1, value=args.default_max_tokens)
+        
+        with gr.Row():
+            submit_button = gr.Button("Submit", scale=3)
+            clear_all_button = gr.ClearButton(components=[video, images, input_text], value="Clear All", scale=1)
+        
+        with gr.Row():
+            output_text = gr.Textbox(placeholder="Output text", label="Output text", lines=10)
+        
+        if args.annotations:
+            with gr.Row():
+                output_annotations = gr.Gallery(label="Annotations", height=MAX_IMAGE_SIZE)
+            outputs = [output_text, output_annotations]
+            fn = functools.partial(run_single_inference, annotations="points")
+        else:
+            fn = run_single_inference
+            outputs = [output_text]
+        
+        submit_button.click(
+            fn=fn,
+            inputs=[video, images, input_text, style, frame_sample_mode, max_frames, max_fps, sampling_fps, max_tok_slider],
+            outputs=outputs,
+        )
+
+    if args.cloudflare_tunnel:
+        import cloudflared_tunnel
+        with cloudflared_tunnel.run() as port:
+            demo.queue().launch(
+                share=False, show_error=True, max_threads=os.cpu_count() - 10, server_port=port,
+                allowed_paths=ALLOWED_PATH
+            )
+    else:
+        demo.queue().launch(
+            server_name=args.server_name,
+            share=not args.no_share, show_error=True, max_threads=os.cpu_count() - 10,
+            server_port=args.port,
+            allowed_paths=ALLOWED_PATH
+        )
+
+
+if __name__ == "__main__":
+    main()

video_preprocessor_config.json

@@ -0,0 +1,48 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_molmo2.Molmo2Processor",
+    "AutoVideoProcessor": "video_processing_molmo2.Molmo2VideoProcessor"
+  },
+  "crop_size": null,
+  "data_format": "channels_first",
+  "default_to_square": true,
+  "device": null,
+  "do_center_crop": null,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "do_sample_frames": true,
+  "fps": null,
+  "frame_sample_mode": "uniform_last_frame",
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "input_data_format": null,
+  "max_fps": 2.0,
+  "num_frames": 384,
+  "pad_size": null,
+  "patch_size": 14,
+  "pooling_size": [
+    3,
+    3
+  ],
+  "processor_class": "Molmo2Processor",
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "return_metadata": false,
+  "sampling_fps": 2,
+  "size": {
+    "height": 378,
+    "width": 378
+  },
+  "video_metadata": null,
+  "video_processor_type": "Molmo2VideoProcessor"
+}

video_processing_molmo2.py

@@ -0,0 +1,975 @@
+"""Video processor class for Molmo2"""
+from functools import partial
+import os
+import warnings
+from contextlib import redirect_stdout
+from io import BytesIO
+from urllib.parse import urlparse
+from typing import Optional, Union, Callable
+
+import numpy as np
+import requests
+import einops
+import torch
+import torchvision.transforms
+
+from transformers.image_utils import (
+    IMAGENET_STANDARD_MEAN,
+    IMAGENET_STANDARD_STD,
+    ImageInput,
+    PILImageResampling,
+    SizeDict,
+    validate_kwargs,
+)
+from transformers.video_utils import (
+    VideoInput,
+    is_valid_video,
+    make_batched_videos,
+    make_batched_metadata,
+    VideoMetadata,
+)
+from transformers.processing_utils import Unpack, VideosKwargs
+from transformers.video_processing_utils import BaseVideoProcessor
+from transformers.utils import logging
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.utils import (
+    is_av_available,
+    is_decord_available,
+    is_torchcodec_available,
+    is_yt_dlp_available,
+    TensorType,
+    logging,
+    to_numpy,
+)
+
+
+logger = logging.get_logger(__name__)
+
+MAX_VIDEO_FPS = 8
+
+
+def normalize_image(
+    image: np.ndarray,
+    image_mean: list[float],
+    image_std: list[float],
+) -> np.ndarray:
+    image -= np.array(image_mean, dtype=np.float32)[None, None, :]
+    image /= np.array(image_std, dtype=np.float32)[None, None, :]
+    return image
+
+
+def resize_image(
+    image: np.ndarray,
+    desired_output_size: list[int],
+    resample: PILImageResampling,
+) -> np.ndarray:
+    if len(image.shape) == 3:
+        is_video = False
+        image = torch.permute(torch.from_numpy(image), [2, 0, 1])
+    else:
+        is_video = True
+        image = torch.permute(torch.from_numpy(image), [0, 3, 1, 2])
+    dtype = image.dtype
+    if torch.is_floating_point(image):
+        in_min = 0.0
+        in_max = 1.0
+        resized = torchvision.transforms.Resize(
+            desired_output_size,
+            resample,
+            antialias=False,
+        )(image)
+        resized = torch.clip(resized, 0.0, 1.0).to(dtype)
+    else:
+        assert image.dtype == torch.uint8, "SigLIP expects float images or uint8 images, but got {}".format(image.dtype)
+        in_min = 0.0
+        in_max = 255.0
+        resized = torchvision.transforms.Resize(
+            desired_output_size,
+            resample,
+            antialias=False,
+        )(image)
+        resized = torch.clip(resized, 0, 255).to(dtype)
+
+    resized = resized.to(torch.float32)
+    resized = (resized - in_min) / (in_max - in_min)
+
+    if is_video:
+        resized = torch.permute(resized, [0, 2, 3, 1]).numpy()
+    else:
+        resized = torch.permute(resized, [1, 2, 0]).numpy()
+
+    return resized
+
+
+def build_resized_image(
+    image: np.ndarray,
+    base_image_input_size: list[int],
+    resample: PILImageResampling,
+    image_mean: list[float],
+    image_std: list[float],
+    image_patch_size: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    resized = resize_image(
+        image, base_image_input_size, resample,
+    )
+    resized = normalize_image(resized, image_mean, image_std)
+    if len(resized.shape) == 3:
+        resized = np.expand_dims(resized, 0)
+    crop_patch_w = base_image_input_size[1] // image_patch_size
+    crop_patch_h = base_image_input_size[0] // image_patch_size
+    resize_idx = np.arange(crop_patch_w*crop_patch_h).reshape([crop_patch_h, crop_patch_w])
+    return resized, resize_idx
+
+
+def batch_pixels_to_patches(array: np.ndarray, patch_size: int) -> np.ndarray:
+    """Reshape images of [n_images, h, w, 3] -> [n_images, n_patches, pixels_per_patch]"""
+    if len(array.shape) == 3:
+        n_crops, h, w = array.shape
+        h_patches = h//patch_size
+        w_patches = w//patch_size
+        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size])
+        array = np.transpose(array, [0, 1, 3, 2, 4])
+        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size])
+        return array
+    else:
+        n_crops, h, w, c = array.shape
+        h_patches = h//patch_size
+        w_patches = w//patch_size
+        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size, c])
+        array = np.transpose(array, [0, 1, 3, 2, 4, 5])
+        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size*c])
+        return array
+
+
+def arange_for_pooling(
+    idx_arr: np.ndarray,
+    pool_h: int,
+    pool_w: int,
+) -> np.ndarray:
+    h_pad = pool_h * ((idx_arr.shape[0] + pool_h - 1) // pool_h) - idx_arr.shape[0]
+    w_pad = pool_w * ((idx_arr.shape[1] + pool_w - 1) // pool_w) - idx_arr.shape[1]
+    idx_arr = np.pad(idx_arr, [[h_pad//2, (h_pad+1)//2], [w_pad//2, (w_pad+1)//2]],
+                     mode='constant',constant_values=-1)
+    return einops.rearrange(
+        idx_arr, "(h dh) (w dw) -> h w (dh dw)", dh=pool_h, dw=pool_w)
+
+
+def image_to_patches_and_grids(
+    image: ImageInput,
+    base_image_input_size: list[int],
+    resample: PILImageResampling,
+    image_mean: list[float],
+    image_std: list[float],
+    image_patch_size: int,
+    image_pooling_w: int,
+    image_pooling_h: int,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    :return image_grids, the shape of each image after pooling
+    :return crops, the image crops to processes with the ViT
+    :return pooled_patch_idx, for each patch_id tokens in `image_tokens`, the indices of the
+                                patches in `crops` to pool for that token, masked with -1
+    """
+    if isinstance(base_image_input_size, int):
+        base_image_input_size = (base_image_input_size, base_image_input_size)
+    
+    pooling_w = image_pooling_w
+    pooling_h = image_pooling_h
+
+    resized, resize_idx = build_resized_image(
+        image,
+        base_image_input_size,
+        resample,
+        image_mean,
+        image_std,
+        image_patch_size,
+    )
+    pooling_idx = arange_for_pooling(resize_idx, pooling_h, pooling_w)
+    h, w = pooling_idx.shape[:2]
+    pooling_idx = pooling_idx.reshape([-1, pooling_h*pooling_w])
+    image_grid = [h, w]
+    return (
+        image_grid,
+        batch_pixels_to_patches(resized, image_patch_size),
+        pooling_idx,
+    )
+
+
+def get_candidate_target_fps(
+    video_fps: Union[int, float],
+    sampling_fps: Union[int, float],
+    max_fps: Union[int, float] = MAX_VIDEO_FPS,
+) -> list[float]:
+    """
+    Return the subset of `video_fps` factors that remain multiples of `sampling_fps`.
+
+    Examples:
+        >>> get_candidate_target_fps(video_fps=6, sampling_fps=2)
+        [2, 6]
+        >>> get_candidate_target_fps(video_fps=5, sampling_fps=1)
+        [1, 5]
+        >>> get_candidate_target_fps(video_fps=2, sampling_fps=2)
+        [2]
+        >>> get_candidate_target_fps(video_fps=5, sampling_fps=2)
+        Traceback (most recent call last):
+            ...
+        ValueError: sampling_fps=2 must divide video_fps=5 to produce consistent frame steps.
+    """
+    video_fps = int(video_fps)
+    sampling_fps = int(sampling_fps)
+    max_fps = int(max_fps)
+
+    if sampling_fps is None:
+        raise ValueError("sampling_fps must be provided")
+    if video_fps <= 0 or sampling_fps <= 0:
+        raise ValueError(f"video_fps and sampling_fps must be positive (got {video_fps}, {sampling_fps})")
+    if video_fps % sampling_fps != 0:
+        raise ValueError(f"sampling_fps={sampling_fps} must divide video_fps={video_fps}.")
+
+    candidates = []
+    for candidate in range(sampling_fps, video_fps + 1, sampling_fps):
+        if candidate > max_fps:
+            break
+        if video_fps % candidate == 0:
+            candidates.append(float(candidate))
+    
+    return candidates
+
+
+def read_video_decord(
+    video_path,
+    sample_timestamps_fn: Callable,
+    **kwargs,
+) -> np.ndarray:
+    """
+    Decode a video using the Decord backend.
+
+    Args:
+        video_path (`str`):
+            Path to the video file.
+        sample_timestamps_fn (`Callable`):
+            A callable function that will return timestamps at which the video should be sampled.
+
+    Returns:
+        tuple[`np.array`, `VideoMetadata`]: A tuple containing:
+            - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).
+            - `VideoMetadata` object.
+    """
+    # Lazy import from decord
+    import importlib
+    decord = importlib.import_module("decord")
+
+    vr = decord.VideoReader(uri=video_path, ctx=decord.cpu(0))  # decord has problems with gpu
+    video_fps = vr.get_avg_fps()
+    total_num_frames = len(vr)
+    time_stamps = vr.get_frame_timestamp(list(range(len(vr))))
+    duration = time_stamps[-1][1] - time_stamps[0][0]
+
+    metadata = VideoMetadata(
+        total_num_frames=int(total_num_frames),
+        fps=float(video_fps),
+        duration=float(duration),
+        video_backend="decord",
+    )
+
+    target_timestamps = sample_timestamps_fn(metadata=metadata, **kwargs)
+    target_timestamps = np.array(target_timestamps)
+    offset = time_stamps[0, 0]
+
+    ix = np.searchsorted(time_stamps[:, 1], target_timestamps + offset, side='right')
+    ix = np.minimum(ix, len(time_stamps) - 1)
+
+    video = vr.get_batch(ix).asnumpy()
+    metadata.update(
+        {
+            "frames_indices": target_timestamps * video_fps,
+            "height": video.shape[1],
+            "width": video.shape[2],
+        }
+    )
+    return video, metadata
+
+
+def read_video_torchcodec(
+    video_path,
+    sample_timestamps_fn: Callable,
+    **kwargs,
+) -> np.ndarray:
+    """
+    Decode a video using torchcodec decoder.
+
+    Args:
+        video_path (`str`):
+            Path to the video file.
+        sample_timestamps_fn (`Callable`):
+            A callable function that will return timestamps at which the video should be sampled.
+
+    Returns:
+        tuple[`np.array`, `VideoMetadata`]: A tuple containing:
+            - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).
+            - `VideoMetadata` object.
+    """
+    # Lazy import torchcodec
+    import importlib
+    torchcodec = importlib.import_module("torchcodec")
+
+    decoder = torchcodec.decoders.VideoDecoder(
+        video_path,
+        # Interestingly `exact` mode takes less than approximate when we load the whole video
+        seek_mode="exact",
+        # Allow FFmpeg decide on the number of threads for efficiency
+        num_ffmpeg_threads=0,
+    )
+    # If the first frame starts at > 0, we effectively clip the video starting at that time
+    # since (most) video players would also skip to that time
+    time_offset = decoder.metadata.begin_stream_seconds_from_content
+    # Note this duration does assume we started playing at `time_offset`
+    duration = decoder.metadata.duration_seconds
+
+    metadata = VideoMetadata(
+        total_num_frames=decoder.metadata.num_frames,
+        fps=decoder.metadata.average_fps,
+        duration=duration,
+        video_backend="torchcodec",
+        height=decoder.metadata.height,
+        width=decoder.metadata.width,
+    )
+
+    target_timestamps = sample_timestamps_fn(metadata=metadata, **kwargs)
+
+    # Floating point/rounding issues might cause `target_timestamps` to be very slightly
+    # out-of-bounds, to handle this we sanity check then clip them
+    assert all(x >= 0 for x in target_timestamps)
+    assert all(x < duration+1e-6 for x in target_timestamps)
+    # 1e-6 padding since torchcodec can throw out-of-bounds errors even if you ask for the
+    # exact boundary value, we should still get the first/last frame anyway
+    max_timestamp = decoder.metadata.end_stream_seconds_from_content - 1e-6
+    min_timestamp = decoder.metadata.begin_stream_seconds_from_content + 1e-6
+    # Note we avoid using numpy ops here to reduce floating precision issues
+    timestamps = [x + time_offset for x in target_timestamps]
+    timestamps = [max(min_timestamp, min(max_timestamp, x)) for x in timestamps]
+
+    video = decoder.get_frames_played_at(timestamps).data.numpy().transpose(0, 2, 3, 1)  # Convert to THWC format
+    target_timestamps = np.array(target_timestamps)
+    metadata.frames_indices = target_timestamps * metadata.fps
+
+    return video, metadata
+
+
+def read_video_pyav(
+    video_path,
+    sample_timestamps_fn: Callable,
+    **kwargs,
+) -> np.ndarray:
+    """
+    Decode a video using the PyAV backend.
+
+    Args:
+        video_path (`str`):
+            Path to the video file.
+        sample_timestamps_fn (`Callable`):
+            A callable function that will return timestamps at which the video should be sampled.
+
+    Returns:
+        tuple[`np.array`, `VideoMetadata`]: A tuple containing:
+            - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).
+            - `VideoMetadata` object.
+    """
+    # Lazy import torchcodec
+    import importlib
+    av = importlib.import_module("av")
+
+    with av.open(video_path) as container:
+        video_stream = container.streams.video[0]
+        fps = video_stream.average_rate or video_stream.guessed_rate
+        it = container.decode(video=0)
+        frames = list(it)
+
+        stream = container.streams.video[0]
+        start = frames[0].pts * stream.time_base
+        container_end = stream.duration
+        if container_end is not None:
+            container_end *= stream.time_base
+        if container_end is None or container_end < frames[-1].pts:
+            # Some problem with stream duration, so use the frame PTS directly
+            # and guess the duration of the last frame
+            end = frames[-1].pts * stream.time_base + 1/fps
+        else:
+            end = container_end
+        duration = float(end - start)
+
+        metadata = VideoMetadata(
+            total_num_frames=len(frames),
+            fps=float(fps),
+            duration=float(duration),
+            video_backend="pyav",
+            height=video_stream.height,
+            width=video_stream.width,
+        )
+
+        target_timestamps = sample_timestamps_fn(metadata=metadata, **kwargs)
+        offset = float(start)
+
+        target_timestamps = np.array(target_timestamps)
+        end_time_stamps = np.array([float(frame.pts * stream.time_base) for frame in frames[1:]] + [duration])
+        indices = np.searchsorted(end_time_stamps, target_timestamps + offset, side='right')
+        indices = np.minimum(indices, len(end_time_stamps) - 1)
+
+        video = np.stack(
+            [frames[i].to_ndarray(format="rgb24", channel_last=True) for i in indices],
+            axis=0,
+        )
+
+        metadata.frames_indices = target_timestamps * fps
+
+        return video, metadata
+
+
+VIDEO_DECODERS = {
+    "decord": read_video_decord,
+    "torchcodec": read_video_torchcodec,
+    "pyav": read_video_pyav,
+}
+
+
+def load_video(
+    video: VideoInput,
+    backend: str = "decord",
+    sample_timestamps_fn: Optional[Callable] = None,
+    **kwargs,
+):
+    """
+    Loads `video` to a numpy array.
+
+    Args:
+        video (`VideoInput`):
+            The video to convert to the numpy array format. Can be a link to video or local path.
+        backend (`str`, *optional*, defaults to `"decord"`):
+            The backend to use when loading the video. Can be any of ["decord", "pyav", ""torchcodec"]. Defaults to "decord".
+        sample_timestamps_fn (`Callable`):
+            A callable function that will return timestamps at which the video should be sampled.
+    """
+
+    # Early exit if provided an array or `PIL` frames
+    if not isinstance(video, str):
+        metadata = [None] * len(video)
+        return video, metadata
+
+    if urlparse(video).netloc in ["www.youtube.com", "youtube.com"]:
+        if not is_yt_dlp_available():
+            raise ImportError("To load a video from YouTube url you have  to install `yt_dlp` first.")
+        # Lazy import from yt_dlp
+        import importlib
+        yt_dlp = importlib.import_module("yt_dlp")
+
+        buffer = BytesIO()
+        with redirect_stdout(buffer), yt_dlp.YoutubeDL() as f:
+            f.download([video])
+        bytes_obj = buffer.getvalue()
+        file_obj = BytesIO(bytes_obj)
+    elif video.startswith("http://") or video.startswith("https://"):
+        file_obj = BytesIO(requests.get(video).content)
+    elif os.path.isfile(video):
+        file_obj = video
+    else:
+        raise TypeError("Incorrect format used for video. Should be an url linking to an video or a local path.")
+
+    # can also load with decord, but not cv2/torchvision
+    # both will fail in case of url links
+    video_is_url = video.startswith("http://") or video.startswith("https://")
+    if video_is_url and backend == "opencv":
+        raise ValueError("If you are trying to load a video from URL, you cannot use 'opencv' as backend")
+
+    if (
+        (not is_decord_available() and backend == "decord")
+        or (not is_torchcodec_available() and backend == "torchcodec")
+        or (not is_av_available() and backend == "pyav")
+    ):
+        raise ImportError(
+            f"You chose backend={backend} for loading the video but the required library is not found in your environment "
+            f"Make sure to install {backend} before loading the video."
+        )
+    
+    video_decoder = VIDEO_DECODERS[backend]
+    video, metadata = video_decoder(file_obj, sample_timestamps_fn, **kwargs)
+    return video, metadata
+
+
+def get_target_fps(
+    video_fps: float,
+    max_frames: int,
+    total_frames: int,
+    frame_sample_mode: str,
+    candidate_target_fps: tuple[float],
+) -> float:
+    """
+    Get the target fps that best spans the video and has the most frames sampled
+    """
+    num_frames_sampled = 0
+    selected_target_fps = None
+    for target_fps in candidate_target_fps:
+        step_size = max(int(video_fps / target_fps), 1)
+        num_frames_sampled_at_fps = int(total_frames / step_size)
+        if num_frames_sampled == 0:
+            if "uniform" in frame_sample_mode:
+                if num_frames_sampled_at_fps > max_frames:
+                    break
+            selected_target_fps = target_fps
+            num_frames_sampled = num_frames_sampled_at_fps
+
+        else:
+            # the candidate sampling fps increases so frame count can't decrease
+            assert num_frames_sampled <= num_frames_sampled_at_fps
+            if num_frames_sampled_at_fps > max_frames:
+                # choose the sampling fps that spans the video
+                continue
+
+            elif num_frames_sampled_at_fps > num_frames_sampled:
+                # both are less than max_frames, choose the one with higher density of frames sampled
+                selected_target_fps = target_fps
+                num_frames_sampled = num_frames_sampled_at_fps
+    return selected_target_fps
+
+
+def get_frame_times_and_chosen_fps(
+    selected_target_fps,
+    total_frames,
+    max_frames,
+    video_fps
+):
+    if selected_target_fps is None:
+        frame_indices = np.linspace(0, total_frames, max_frames, endpoint=False, dtype=int)
+    else:
+        step_size = max(int(video_fps / selected_target_fps), 1)
+        frame_indices = np.arange(0, total_frames, step_size)
+    if len(frame_indices) > max_frames:
+        frame_indices = frame_indices[:max_frames]
+    return selected_target_fps, frame_indices
+
+
+class Molmo2VideoProcessorKwargs(VideosKwargs, total=False):
+    patch_size: Optional[int]
+    pooling_size: Optional[list[int]]
+    frame_sample_mode: Optional[str]
+    max_fps: Optional[int]
+    sampling_fps: Optional[int]
+
+
+class Molmo2VideoProcessor(BaseVideoProcessor):
+    resample = PILImageResampling.BILINEAR
+    size = {"height": 378, "width": 378}
+    image_mean = IMAGENET_STANDARD_MEAN
+    image_std = IMAGENET_STANDARD_STD
+    do_resize = True
+    do_rescale = True
+    do_normalize = True
+    do_convert_rgb = True
+    patch_size = 14
+    pooling_size = [3, 3]
+    do_sample_frames = True
+    frame_sample_mode = "uniform_last_frame"
+    max_fps = 2
+    sampling_fps = 2
+    valid_kwargs = Molmo2VideoProcessorKwargs
+    model_input_names = ["pixel_values_videos", "video_token_pooling", "video_grids"]
+
+    def __init__(self, **kwargs: Unpack[Molmo2VideoProcessorKwargs]):
+        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 _further_process_kwargs(
+        self,
+        size: Optional[SizeDict] = None,
+        **kwargs,
+    ) -> dict:
+        """
+        Update kwargs that need further processing before being validated
+        Can be overridden by subclasses to customize the processing of kwargs.
+        """
+        if size is not None and ("height" not in size or "width" not in size):
+            raise ValueError("size must contain 'height' and 'width' keys.")
+
+        return super()._further_process_kwargs(size=size, **kwargs)
+
+    def sample_times(
+        self,
+        metadata: VideoMetadata,
+        frame_sample_mode: str,
+        num_frames: int,
+        max_fps: Optional[int] = None,
+        sampling_fps: Optional[int] = None,
+        **kwargs,
+    ) -> np.ndarray:
+        """
+        Time-based sampling if an array video is passed
+        Args:
+            metadata (`VideoMetadata`):
+                Metadata of the video containing information about total duration, fps and total number of frames.
+            frame_sample_mode (`str`, *optional*):
+                Mode to sample frames. Defaults to `self.frame_sample_mode`.
+            num_frames (`int`, *optional*):
+                Maximum number of frames to sample. Defaults to `self.num_frames`.
+            man_fps (`int`, *optional*):
+                Maximum frames per second to sample.
+            sampling_fps (`int`, *optional*):
+                Sampling frames per second. Defaults to `self.sampling_fps`.
+                Used when `frame_sample_mode` is `"fps"`.
+        """
+        frame_sample_mode = frame_sample_mode or self.frame_sample_mode
+        num_frames = num_frames or self.num_frames
+        sampling_fps = sampling_fps or self.sampling_fps
+
+        duration = metadata.duration or metadata.total_num_frames / metadata.fps
+        if frame_sample_mode == "fps":
+            candidate_target_fps = get_candidate_target_fps(metadata.fps, sampling_fps)
+            # Try larger and larger FPSs until we hit one that can't span the video
+            target_fps = candidate_target_fps[0]
+            for candidate_fps in candidate_target_fps[1:]:
+                if num_frames / candidate_fps < duration:
+                    break
+                target_fps = candidate_fps
+            times = np.arange(0, num_frames) / target_fps
+            times = times[times < duration]
+            return times
+        elif frame_sample_mode == "uniform_last_frame":
+            if max_fps is not None:
+                max_duration = (num_frames-1) / max_fps  # -1 to include the last frame
+                if max_duration < duration:
+                    times = np.linspace(
+                        0, duration, num=num_frames, endpoint=True, dtype=np.float64
+                    )
+                else:
+                    times = np.arange(0.0, stop=duration, step=1/max_fps)
+                    times = np.concatenate([times, [duration]], axis=0)
+                    assert len(times) <= num_frames
+            else:
+                times = np.linspace(
+                    0, duration, num=num_frames, endpoint=True, dtype=np.float64
+                )
+            return times
+        else:
+            raise NotImplementedError(frame_sample_mode)
+
+    def sample_frames(
+        self,
+        metadata: VideoMetadata,
+        frame_sample_mode: Optional[str] = None,
+        num_frames: Optional[int] = None,
+        max_fps: Optional[int] = None,
+        sampling_fps: Optional[int] = None,
+        **kwargs,
+    ) -> np.ndarray:
+        """
+        Frame-based sampling if an array video is passed
+        Args:
+            metadata (`VideoMetadata`):
+                Metadata of the video containing information about total duration, fps and total number of frames.
+            frame_sample_mode (`str`, *optional*):
+                Mode to sample frames. Defaults to `self.frame_sample_mode`.
+            num_frames (`int`, *optional*):
+                Maximum number of frames to sample. Defaults to `self.num_frames`.
+            max_fps (`int`, *optional*):
+                Maximum frames per second to sample.
+            sampling_fps (`int`, *optional*):
+                Sampling frames per second. Defaults to `self.sampling_fps`.
+                Used when `frame_sample_mode` is `"fps"`.
+        """
+        frame_sample_mode = frame_sample_mode or self.frame_sample_mode
+        num_frames = num_frames or self.num_frames
+        sampling_fps = sampling_fps or self.sampling_fps
+
+        total_num_frames = metadata.total_num_frames
+        if frame_sample_mode == "uniform_last_frame" and max_fps is not None:
+            duration = total_num_frames / metadata.fps
+            if total_num_frames <= 2:
+                return np.arange(total_num_frames).astype(int)
+            if duration > (num_frames - 1) / max_fps:  # -1 to include the last frame
+                # uniform fallback
+                indices = np.linspace(
+                    0,
+                    total_num_frames - 1,
+                    num=min(num_frames, total_num_frames),
+                    endpoint=True,
+                ).astype(int)
+                return indices
+            else:
+                float_indices = np.arange(
+                    0.0, stop=total_num_frames - 1, step=float(metadata.fps / max_fps),
+                )
+                if np.round(float_indices[-1]) != total_num_frames - 1:
+                    float_indices = np.concatenate([float_indices, [total_num_frames - 1]], axis=0)
+                indices = np.round(float_indices).astype(int)
+                assert indices[-1] < total_num_frames
+                assert len(float_indices) <= num_frames
+                return indices
+        elif frame_sample_mode == "uniform_last_frame":
+            indices = np.linspace(
+                0, total_num_frames - 1, num=min(num_frames, total_num_frames), endpoint=True,
+            ).astype(int)
+            return indices
+        elif frame_sample_mode == "fps":
+            candidate_target_fps = get_candidate_target_fps(metadata.fps, sampling_fps)
+            selected_target_fps = get_target_fps(
+                metadata.fps,
+                num_frames,
+                total_num_frames,
+                frame_sample_mode,
+                candidate_target_fps,
+            )
+            _, indices = get_frame_times_and_chosen_fps(
+                selected_target_fps,
+                total_num_frames,
+                num_frames,
+                metadata.fps,
+            )
+            return indices
+        else:
+            raise NotImplementedError(frame_sample_mode)
+    
+    def fetch_videos(
+        self,
+        video_url_or_urls: Union[str, list[str], list[list[str]]],
+        sample_timestamps_fn=None
+    ):
+        """
+        Convert a single or a list of urls into the corresponding `np.array` objects.
+
+        If a single url is passed, the return value will be a single object. If a list is passed a list of objects is
+        returned.
+        """
+        if (
+            (not is_decord_available())
+            and (not is_torchcodec_available())
+            and (not is_av_available())
+        ):
+            raise ImportError(
+                "Molmo2VideoProcessor requires `decord`, `torchcodec`, or `av` to be installed."
+            )
+
+        if is_decord_available():
+            backend = "decord"
+        elif is_torchcodec_available():
+            warnings.warn(
+                "`decord` is not installed and cannot be used to decode the video by default. "
+                "Falling back to `torchcodec`."
+            )
+            backend = "torchcodec"
+        else:
+            warnings.warn(
+                "`decord` is not installed and cannot be used to decode the video by default. "
+                "Falling back to `PyAV`."
+            )
+            backend = "pyav"
+
+        if isinstance(video_url_or_urls, list):
+            return list(zip(*[self.fetch_videos(x, sample_timestamps_fn=sample_timestamps_fn) for x in video_url_or_urls]))
+        else:
+            return load_video(video_url_or_urls, backend=backend, sample_timestamps_fn=sample_timestamps_fn)
+
+    def _decode_and_sample_videos(
+        self,
+        videos: VideoInput,
+        video_metadata: Union[VideoMetadata, dict],
+        do_sample_frames: Optional[bool] = None,
+        sample_indices_fn: Optional[Callable] = None,
+        sample_timestamps_fn: Optional[Callable] = None,
+    ):
+        """
+        Decode input videos and sample frames if needed.
+        """
+        videos = make_batched_videos(videos)
+        video_metadata = make_batched_metadata(videos, video_metadata=video_metadata)
+
+        # Framed-based sampling if an array video is passed
+        # Otherwise, time-based sampling with decoding
+        if is_valid_video(videos[0]) and do_sample_frames:
+            assert video_metadata[0].fps is not None, "FPS must be provided for video input"
+            sampled_videos = []
+            sampled_metadata = []
+            for video, metadata in zip(videos, video_metadata):
+                indices = sample_indices_fn(metadata=metadata)
+                metadata.frames_indices = indices
+                sampled_videos.append(video[indices])
+                sampled_metadata.append(metadata)
+            videos = sampled_videos
+            video_metadata = sampled_metadata
+        elif not is_valid_video(videos[0]):
+            if sample_indices_fn is None:
+                logger.warning(
+                    "do_sample_frames is False, but video array is not provided: "
+                    "Will decode the video and sample frames using Molmo2's default sampling mode"
+                )
+            if isinstance(videos[0], list):
+                raise ValueError(
+                    "A list of images is not supported for video input!"
+                )
+            else:
+                videos, video_metadata = self.fetch_videos(videos, sample_timestamps_fn=sample_timestamps_fn)
+        
+        return videos, video_metadata
+    
+    def _prepare_input_videos(
+        self,
+        videos: VideoInput,
+        **kwargs,
+    ) -> list[np.ndarray]:
+        processed_videos = [to_numpy(video) for video in videos]
+        return processed_videos
+    
+    def preprocess(
+        self,
+        videos: VideoInput,
+        **kwargs: Unpack[Molmo2VideoProcessorKwargs],
+    ) -> BatchFeature:
+        validate_kwargs(
+            captured_kwargs=kwargs.keys(),
+            valid_processor_keys=list(self.valid_kwargs.__annotations__.keys()) + ["return_tensors"],
+        )
+
+        # Set default kwargs from self. This ensures that if a kwarg is not provided
+        # by the user, it gets its default value from the instance, or is set to None.
+        for kwarg_name in self.valid_kwargs.__annotations__:
+            kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None))
+        
+        do_sample_frames = kwargs.pop("do_sample_frames")
+        video_metadata = kwargs.pop("video_metadata")
+
+        sample_indices_fn = partial(self.sample_frames, **kwargs) if do_sample_frames else None
+        sample_timestamps_fn = partial(self.sample_times, **kwargs)
+        videos, video_metadata = self._decode_and_sample_videos(
+            videos,
+            video_metadata=video_metadata,
+            do_sample_frames=do_sample_frames,
+            sample_indices_fn=sample_indices_fn,
+            sample_timestamps_fn=sample_timestamps_fn,
+        )
+        videos = self._prepare_input_videos(videos=videos)
+
+        kwargs = self._further_process_kwargs(**kwargs)
+
+        return_metadata = kwargs.pop("return_metadata")
+        preprocessed_videos = self._preprocess(videos=videos, **kwargs)
+        if return_metadata:
+            preprocessed_videos["video_metadata"] = video_metadata
+        return preprocessed_videos
+    
+    def _preprocess(
+        self,
+        videos: list[np.ndarray],
+        size: Optional[SizeDict] = None,
+        resample: Optional[PILImageResampling] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        patch_size: Optional[int] = None,
+        pooling_size: Optional[list[int]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_pointing_metadata: bool = False,
+        **kwargs,
+    ) -> BatchFeature:
+        """
+        Preprocess a video for the model.
+        Args:
+            videos (`VideoInput`):
+                Video to preprocess.
+            size (`SizeDict`, *optional*, defaults to `self.size`):
+                Size of the image after resizing.
+            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
+                Resampling filter to use when resizing the image. This can be one of the enum `PILImageResampling`. Only
+                has an effect if `do_resize` is set to `True`.
+            image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):
+                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):
+                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
+                `True`.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+            patch_size (`int`, *optional*, defaults to `self.patch_size`):
+                The spatial patch size of the vision encoder.
+            pooling_size (`list[int]`, *optional*, defaults to `self.pooling_size`):
+                The pooling size of the vision adapter.
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                - Unset: Return a list of `np.ndarray`.
+                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
+                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
+
+        Returns:
+            A `BatchFeature` containing the following keys:
+                - `pixel_values_videos`: The preprocessed videos.
+                - `video_token_pooling`: The indices of the patches in `crops` to pool for each token in `video_tokens`.
+                - `video_grids`: The video grids.
+        """
+        if size.height is None or size.width is None:
+            raise ValueError("size must contain 'height' and 'width' keys.")
+        
+        base_image_input_size = [size.height, size.width]
+
+        resample = resample or self.resample
+        image_mean = image_mean or self.image_mean
+        image_std = image_std or self.image_std
+        do_convert_rgb = do_convert_rgb or self.do_convert_rgb
+
+        patch_size = patch_size or self.patch_size
+        pooling_size = pooling_size or self.pooling_size
+
+        image_pooling_h, image_pooling_w = pooling_size
+
+        batch_grids = []
+        batch_crops = []
+        batch_pooled_patches_idx = []
+
+        for video in videos:
+            all_crops = []
+            pooled_patches_idx = []
+
+            for frame in video:
+                image_grid, crops, pooled_idx = image_to_patches_and_grids(
+                    frame,
+                    base_image_input_size,
+                    resample,
+                    image_mean,
+                    image_std,
+                    patch_size,
+                    image_pooling_w,
+                    image_pooling_h,
+                )
+                offset = sum(np.prod(x.shape[:2]) for x in all_crops)
+                pooled_idx_with_offset = np.where(pooled_idx >= 0, pooled_idx + offset, pooled_idx)
+                pooled_patches_idx.append(pooled_idx_with_offset)
+                all_crops.append(crops)
+
+            video_grid = np.array([len(video), image_grid[0], image_grid[1]])
+            all_crops = np.concatenate(all_crops, 0)
+            pooled_patches_idx = np.concatenate(pooled_patches_idx, 0)
+
+            batch_grids.append(video_grid)
+            batch_crops.append(all_crops)
+            batch_pooled_patches_idx.append(pooled_patches_idx)
+        
+        video_grids = np.stack(batch_grids, 0)
+        pixel_values_videos = np.concatenate(batch_crops, 0)
+        video_token_pooling = np.concatenate(batch_pooled_patches_idx, 0)
+        
+        data = BatchFeature(dict(
+            pixel_values_videos=pixel_values_videos,
+            video_token_pooling=video_token_pooling,
+            video_grids=video_grids,
+        ), tensor_type=return_tensors)
+        if return_pointing_metadata:
+            t = pixel_values_videos.shape[0]
+            assert base_image_input_size[0] % self.patch_size == 0
+            assert base_image_input_size[1] % self.patch_size == 0
+            crop_w = base_image_input_size[0] // self.patch_size
+            crop_h = base_image_input_size[1] // self.patch_size
+            data["subpatch_mapping"] = np.arange(t*crop_w*crop_h).reshape([t, crop_h, crop_w])
+            data["video_token_pooling_np"] = video_token_pooling
+        return data
+
+
+Molmo2VideoProcessor.register_for_auto_class()