mlx-community/MolmoPoint-8B-5bit
016
1import math2import re3from copy import deepcopy4from dataclasses import dataclass5from typing import Optional, Union, Callable, Any, List, Tuple6 7import numpy as np8import torch9from torch import nn10 11from torch.nn import functional as F12from transformers import LogitsProcessorList, LogitsProcessor, AutoProcessor, ViTConfig13from transformers.image_utils import PILImageResampling14 15from transformers.models.auto import AutoModelForImageTextToText16from transformers.activations import ACT2FN17from transformers.configuration_utils import PretrainedConfig18from transformers.cache_utils import Cache, DynamicCache19from transformers.generation import GenerationMixin20from transformers.masking_utils import create_causal_mask, create_masks_for_generate21from transformers.modeling_flash_attention_utils import (22 _flash_attention_forward,23 FlashAttentionKwargs,24 flash_attn_supports_top_left_mask,25)26from transformers.modeling_layers import GradientCheckpointingLayer27from transformers.modeling_outputs import (28 BaseModelOutputWithPast,29)30from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update31from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel32from transformers.processing_utils import Unpack33from transformers.utils import (34 ModelOutput,35 TransformersKwargs,36 can_return_tuple,37 logging,38)39 40from .configuration_molmo2 import Molmo2VitConfig, Molmo2TextConfig, Molmo2AdapterConfig41from .configuration_molmo_point import MolmoPointConfig, MolmoPointAdapterConfig42from .image_processing_molmo2 import Molmo2ImagesKwargs, image_to_patches_and_grids43from .modeling_molmo2 import ImageProjectorMLP, Molmo2VisionTransformer, Molmo2RMSNorm, \44 Molmo2RotaryEmbedding, Molmo2PostNormDecoderLayer, Molmo2DecoderLayer, Molmo2Attention, \45 Molmo2Embedding46 47# FIXME remove48processor = None49def decode(ids):50 global processor51 if processor is None:52 processor = AutoProcessor.from_pretrained(53 "/weka/oe-training-default/mm-olmo/released-models-molmo2-point-0326/MolmoPoint-8B/hf-step2000", trust_remote_code=True,54 padding_side="left")55 return processor.post_process_image_text_to_text(ids.view(1), skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]56 57 58logger = logging.get_logger(__name__)59NO_POINTS_LABEL = 100000060 61 62EXTRACT_POINT_TRIPLE = re.compile(f"<POINT_(\d+)> ?<POINT_(\d+)> ?<POINT_(\d+)> ?([0-9]+)" )63 64 65def get_subpatch_ids(output_text, pooling, no_more_points_class):66 n_patches, n_subpatches = pooling.shape[-2:]67 if no_more_points_class:68 n_patches += 169 for match in EXTRACT_POINT_TRIPLE.finditer(output_text):70 patch_id, subpatch_num = int(match.group(1)), int(match.group(2))71 subpatch_id = subpatch_num - n_patches72 location_num = int(match.group(3))73 location_id = location_num - n_patches - n_subpatches74 example_id = int(match.group(4))75 vit_patch_id = pooling[patch_id, subpatch_id]76 yield vit_patch_id, location_id, example_id77 78 79@dataclass80class ImageCache:81 """Extra stuff we need to cache when doing autoregressive generation with pointing"""82 83 patch_k: torch.FloatTensor84 """K values of the image tokens"""85 86 patch_k_mask: torch.BoolTensor87 """Mask over image tokens that can be selected"""88 89 subpatch_k: torch.FloatTensor90 """K values of the ViT patches before pooling"""91 92 token_pooling: torch.LongTensor93 """token pooling array mapping image_patch_id -> ViT patches pooled for that patch"""94 95 vit_features: torch.FloatTensor96 """Features before pooling, used for building input embeddings"""97 98 image_pos_ids: Optional[torch.LongTensor] = None99 """Position ids of the image tokens if need for rotary embeddings"""100 101 image_features0: Optional[torch.FloatTensor] = None102 """"Image features, might be needed to embed new patch prediction tokens"""103 104 flat_image_tokens_to_flat_image_features: Optional[torch.LongTensor] = None105 """Cached for indexing uses"""106 107 108@dataclass109class MolmoPointCausalLMOutputWithPast(ModelOutput):110 """111 Base class for MolmoPoint causal language model (or autoregressive) outputs.112 113 Args:114 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):115 Language modeling loss (for next-token prediction).116 logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):117 Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).118 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):119 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).120 121 Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see122 `past_key_values` input) to speed up sequential decoding.123 image_hidden_states (`torch.FloatTensor`, *optional*):124 A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.125 image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.126 """127 128 loss: Optional[torch.FloatTensor] = None129 logits: Optional[torch.FloatTensor] = None130 past_key_values: Optional[Cache] = None131 hidden_states: Optional[tuple[torch.FloatTensor]] = None132 attentions: Optional[tuple[torch.FloatTensor]] = None133 image_hidden_states: Optional[torch.FloatTensor] = None134 image_data: Optional[ImageCache] = None135 patch_logits: Optional[torch.FloatTensor] = None136 subpatch_logits: Optional[torch.FloatTensor] = None137 location_logits: Optional[torch.FloatTensor] = None138 last_predicted_patch_id: Optional[torch.LongTensor] = None139 140 141@dataclass142class MolmoPointModelOutputWithPast(BaseModelOutputWithPast):143 """144 Base class for Molmo2 outputs, with hidden states and attentions.145 146 Args:147 image_hidden_states (`torch.FloatTensor`, *optional*):148 A `torch.FloatTensor` of size `(batch_num_patches, hidden_size)`.149 image_hidden_states of the model produced by the vision backbone150 """151 last_hidden_state: Optional[torch.FloatTensor] = None152 past_key_values: Optional[Cache] = None153 hidden_states: Optional[tuple[torch.FloatTensor]] = None154 attentions: Optional[tuple[torch.FloatTensor]] = None155 image_hidden_states: Optional[torch.FloatTensor] = None156 image_data: Optional[ImageCache] = None157 patch_logits: Optional[torch.FloatTensor] = None158 subpatch_logits: Optional[torch.FloatTensor] = None159 location_logits: Optional[torch.FloatTensor] = None160 input_ids: Optional[torch.LongTensor] = None161 last_predicted_patch_id: Optional[torch.LongTensor] = None162 163 164class MolmoPointPatchRope(nn.Module):165 inv_freq: torch.Tensor # fix linting for `register_buffer`166 167 def __init__(168 self,169 theta: float,170 dim: int,171 device: Union[str, torch.device] = None,172 ):173 super().__init__()174 attention_factor = 1.0 # Unused in this type of RoPE175 inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))176 self.register_buffer("inv_freq", inv_freq, persistent=False)177 178 def rotate_half(self, x: torch.Tensor) -> torch.Tensor:179 B, hs = x.size()180 x = x.view(B, 2, hs // 2)181 x1, x2 = x.unbind(dim=-2)182 return torch.cat((-x2, x1), dim=-1)183 184 @torch.no_grad()185 def forward(self, x, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:186 inv_freq_expanded = self.inv_freq.float().to(x.device)187 position_ids_expanded = position_ids.float()188 189 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"190 with torch.autocast(device_type=device_type, enabled=False): # Force float32191 x = x.float()192 freqs = position_ids_expanded[:, None] * inv_freq_expanded[None, :]193 emb = torch.cat((freqs, freqs), dim=-1)194 cos = emb.cos()195 sin = emb.sin()196 out = ((x * cos) + (self.rotate_half(x) * sin))197 198 return out.to(dtype=x.dtype)199 200 201class ViTMultiHeadDotProductAttention(nn.Module):202 def __init__(203 self,204 hidden_size: int,205 num_heads: int,206 num_key_value_heads: int,207 head_dim: int,208 use_bias: bool = True,209 input_dim: Optional[int] = None,210 float32_attention: bool = True,211 attention_dropout: float = 0.0,212 residual_dropout: float = 0.0,213 device: Union[str, torch.device] = None,214 attn_implementation: str = "eager",215 out_layer: bool=True216 ):217 super().__init__()218 219 self.hidden_size = hidden_size220 self.num_heads = num_heads221 self.head_dim = head_dim222 self.num_key_value_heads = num_key_value_heads223 self.num_key_value_groups = self.num_heads // self.num_key_value_heads224 self.attn_implementation = attn_implementation225 self.is_causal = False226 227 input_dim = input_dim or hidden_size228 229 self.wq = nn.Linear(230 input_dim,231 self.num_heads * self.head_dim,232 bias=use_bias,233 device=device,234 )235 self.wk = nn.Linear(236 input_dim,237 self.num_key_value_heads * self.head_dim,238 bias=use_bias,239 device=device,240 )241 self.wv = nn.Linear(242 input_dim,243 self.num_key_value_heads * self.head_dim,244 bias=use_bias,245 device=device,246 )247 if out_layer:248 self.wo = nn.Linear(249 self.num_heads * self.head_dim,250 self.hidden_size,251 )252 else:253 self.wo = None254 self.float32_attention = float32_attention255 self.attention_dropout = attention_dropout256 self.residual_dropout = nn.Dropout(residual_dropout)257 258 def _split_heads(self, hidden_states, num_heads) -> torch.Tensor:259 return hidden_states.reshape(hidden_states.shape[:2] + (num_heads, self.head_dim))260 261 def _merge_heads(self, hidden_states) -> torch.Tensor:262 return hidden_states.reshape(hidden_states.shape[:2] + (self.hidden_size,))263 264 def forward(265 self,266 inputs_q: torch.Tensor,267 inputs_kv: Optional[torch.Tensor] = None,268 attn_mask: Optional[torch.Tensor] = None,269 ) -> torch.Tensor:270 271 if inputs_kv is not None:272 inputs_k = inputs_kv273 inputs_v = inputs_kv274 else:275 inputs_k = inputs_q276 inputs_v = inputs_q277 278 xq, xk, xv = self.wq(inputs_q), self.wk(inputs_k), self.wv(inputs_v)279 280 xq = self._split_heads(xq, self.num_heads)281 xk = self._split_heads(xk, self.num_key_value_heads)282 xv = self._split_heads(xv, self.num_key_value_heads)283 284 if self.num_heads != self.num_key_value_heads:285 xk = xk.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads)286 xv = xv.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads)287 288 og_dtype = xq.dtype289 290 if self.float32_attention:291 xq = xq.to(torch.float)292 xk = xk.to(torch.float)293 294 dropout_p = 0.0 if not self.training else self.attention_dropout295 296 if self.attn_implementation == "eager":297 attn_weights = torch.einsum("...qhd,...khd->...hqk", xq / math.sqrt(xq.size(-1)), xk)298 attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(xq.dtype)299 attn_weights = F.dropout(300 attn_weights,301 p=dropout_p,302 training=self.training303 )304 attn_output = torch.einsum("...hqk,...khd->...qhd", attn_weights.to(xv.dtype), xv)305 306 elif self.attn_implementation == "sdpa":307 if not torch.is_autocast_enabled():308 xv = xv.to(torch.float)309 310 attn_output = F.scaled_dot_product_attention(311 xq.transpose(1, 2).contiguous(),312 xk.transpose(1, 2).contiguous(),313 xv.transpose(1, 2).contiguous(),314 attn_mask=attn_mask,315 is_causal=False,316 dropout_p=dropout_p,317 ).transpose(1, 2)318 319 elif self.attn_implementation == "flash_attention_2":320 if xq.dtype == torch.float32:321 if torch.is_autocast_enabled():322 target_dtype = torch.get_autocast_gpu_dtype()323 else:324 target_dtype = self.wq.weight.dtype325 attn_output = _flash_attention_forward(326 xq,327 xk,328 xv,329 attention_mask=attn_mask,330 query_length=inputs_q.shape[1],331 is_causal=False,332 dropout=dropout_p,333 softmax_scale=xq.shape[-1] ** -0.5,334 use_top_left_mask=flash_attn_supports_top_left_mask(),335 target_dtype=target_dtype,336 implementation=self.attn_implementation,337 )338 else:339 raise ValueError(f"Attention implementation {self.attn_implementation} not supported")340 341 attn_output = attn_output.to(og_dtype)342 attn_output = self._merge_heads(attn_output)343 if self.wo is not None:344 attn_output = self.wo(attn_output)345 attn_output = self.residual_dropout(attn_output)346 347 return attn_output348 349 350class PointPredictor(nn.Module):351 """Point predictor logic"""352 # We separate this out so accelerate will co-locate all these parameters on the same device353 354 def __init__(self, config):355 super().__init__()356 self.config = config357 llm_dim = config.text_config.hidden_size358 patch_embed_dim = config.patch_embed_dim359 vit_dim = self.config.vit_config.hidden_size * len(self.config.adapter_config.vit_layers)360 if self.config.layer_norm_x:361 self.x_norm = Molmo2RMSNorm(llm_dim, eps=self.config.text_config.layer_norm_eps)362 else:363 self.x_norm = None364 if self.config.token_prediction_rotary == "none":365 self.patch_rotary = None366 else:367 theta = self.config.token_prediction_rotary_theta or self.config.llm.rope_theta368 if self.config.token_prediction_rotary == "one_d":369 self.patch_rotary = MolmoPointPatchRope(theta, self.config.patch_embed_dim)370 else:371 raise NotImplementedError()372 self.patch_q = nn.Linear(llm_dim, patch_embed_dim)373 self.patch_k = nn.Linear(llm_dim, patch_embed_dim)374 self.subpatch_q = nn.Linear(llm_dim, patch_embed_dim)375 self.subpatch_k = nn.Linear(vit_dim, patch_embed_dim)376 self.add_no_point_class_embed = MolmoPointPadWithLearnedVector(patch_embed_dim)377 if self.config.patch_location == "3x3":378 self.subpatch_loc_k = nn.Linear(llm_dim, 9)379 elif self.config.patch_location is None:380 self.subpatch_loc_k = None381 else:382 raise NotImplementedError(f"Patch location {self.config.patch_location} not implemented")383 384 def forward(385 self,386 x,387 token_pooling,388 is_image_token,389 is_patch,390 is_subpatch,391 is_indexable_image_token,392 vit_features,393 vit_features_mask,394 image_features_mask,395 input_patch_ids,396 last_predicted_patch_id,397 image_data: ImageCache398 ):399 dim = self.config.text_config.hidden_size400 batch_size = x.shape[0]401 if self.x_norm is not None:402 x_norm = self.x_norm(x)403 elif self.config.norm_x:404 x_norm = x / math.sqrt(dim)405 else:406 x_norm = x407 408 # Build the keys, or get them from the cache409 if image_data is not None:410 patch_k, subpatch_k = image_data.patch_k, image_data.subpatch_k411 patch_k_mask = image_data.patch_k_mask412 token_pooling = image_data.token_pooling413 vit_features_mask = token_pooling >= 0414 image_pos_ids = image_data.image_pos_ids415 else:416 # Build patch keys, this takes a bit of indexing trickery since we want the keys in417 # shape [batch, n_image_tokens] not [batch, sequence_length]418 n_image_tokens = token_pooling.shape[1]419 patch_k_flat = self.patch_k(x_norm.view(-1, dim)[is_image_token.view(-1)])420 if self.patch_rotary is not None:421 image_token_indices = torch.cumsum(is_indexable_image_token, dim=-1) - 1422 image_pos_ids_flat = image_token_indices.view(-1)[is_image_token.view(-1)]423 patch_k_flat = self.patch_rotary(patch_k_flat, image_pos_ids_flat)424 425 # Computed for use with the query vectors426 image_pos_ids = torch.zeros([batch_size, n_image_tokens], dtype=torch.long,427 device=image_pos_ids_flat.device)428 image_pos_ids.view(-1)[image_features_mask.view(-1)] = image_pos_ids_flat429 else:430 image_pos_ids = None431 432 patch_k = torch.zeros([batch_size, n_image_tokens, patch_k_flat.shape[-1]],433 dtype=x.dtype, device=x.device)434 patch_k.view(-1, patch_k_flat.shape[-1])[image_features_mask.flatten()] = patch_k_flat.to(dtype=x.dtype)435 436 patch_k_mask = image_features_mask.clone()437 patch_k_mask.view(-1)[image_features_mask.view(-1)] = (438 is_indexable_image_token.view(-1)[is_image_token.view(-1)])439 440 if self.config.no_more_points_class:441 patch_k = self.add_no_point_class_embed(patch_k)442 patch_k_mask = F.pad(patch_k_mask, (0, 1), value=True)443 444 subpatch_k = self.subpatch_k(vit_features)445 446 patch_logits, subpatch_logits, location_logits = None, None, None447 if image_data is not None:448 # Predict patch locations, only done after pre-filling449 batch_idx = torch.arange(batch_size, device=x_norm.device)450 image_q = self.patch_q(x_norm)451 if self.patch_rotary is not None and last_predicted_patch_id is not None:452 rotate_by = image_pos_ids[batch_idx, last_predicted_patch_id]453 rotate_by = torch.where(last_predicted_patch_id >= 0, rotate_by, 0)454 rotate_by = rotate_by.squeeze(-1)455 image_q = self.patch_rotary(456 image_q.view(-1, image_q.shape[-1]),457 torch.clamp(rotate_by, min=0),458 ).reshape(batch_size, -1, image_q.shape[-1])459 460 dots = torch.matmul(image_q, patch_k.transpose(1, 2)) # [batch, 1, num_images]461 if self.config.norm_logits:462 dots = dots / math.sqrt(dots.shape[-1])463 464 valid = patch_k_mask[:, None, :]465 patch_logits = torch.where(valid, dots, -100000000)466 467 if torch.any(is_patch):468 if x_norm.shape[1] != 1:469 raise NotImplementedError()470 subpatch_point_q = self.subpatch_q(x_norm.squeeze(1))471 subpatch_k = subpatch_k[batch_idx, input_patch_ids.squeeze(1)]472 subpatch_logits = torch.einsum("pd,pcd->pc", subpatch_point_q, subpatch_k)473 if self.config.norm_logits:474 subpatch_logits = subpatch_logits / math.sqrt(patch_k.shape[-1])475 subpatch_mask = vit_features_mask[batch_idx, input_patch_ids.squeeze(1)]476 subpatch_logits = torch.where(subpatch_mask, subpatch_logits, -100000)477 subpatch_logits = subpatch_logits[:, None, :]478 479 if torch.any(is_subpatch):480 location_logits = self.subpatch_loc_k(x)481 482 if image_data is None:483 image_data = ImageCache(484 patch_k=patch_k,485 subpatch_k=subpatch_k,486 vit_features=vit_features,487 patch_k_mask=patch_k_mask,488 token_pooling=token_pooling,489 image_pos_ids=image_pos_ids,490 )491 return patch_logits, subpatch_logits, location_logits, image_data492 493 494class MolmoPointPreTrainedModel(PreTrainedModel):495 config: MolmoPointConfig496 base_model_prefix = "model"497 supports_gradient_checkpointing = True498 _no_split_modules = [499 "Molmo2DecoderLayer",500 "Molmo2PostNormDecoderLayer",501 "Molmo2VisionBlock",502 "ViTMultiHeadDotProductAttention",503 "PointPredictor"504 ]505 _skip_keys_device_placement = "past_key_values"506 _supports_flash_attn = True507 _supports_sdpa = True508 509 _can_compile_fullgraph = True510 _supports_attention_backend = True511 _can_record_outputs = {512 "hidden_states": Molmo2DecoderLayer,513 "attentions": Molmo2Attention,514 }515 516 def _init_weights(self, module):517 std = self.config.initializer_range518 if isinstance(module, (nn.Linear,)):519 module.weight.data.normal_(mean=0.0, std=std)520 if module.bias is not None:521 module.bias.data.zero_()522 elif isinstance(module, Molmo2Embedding):523 module.embedding.data.normal_(mean=0.0, std=std)524 module.new_embedding.data.normal_(mean=0.0, std=std)525 elif isinstance(module, nn.Embedding):526 module.weight.data.normal_(mean=0.0, std=std)527 if module.padding_idx is not None:528 module.weight.data[module.padding_idx].zero_()529 elif isinstance(module, Molmo2RMSNorm):530 module.weight.data.fill_(1.0)531 elif isinstance(module, nn.LayerNorm):532 module.weight.data.fill_(1.0)533 if module.bias is not None:534 module.bias.data.zero_()535 536 537class GeneratedTokenBounds:538 """Describes what tokens id ranges are patch/subpatch/location tokens"""539 540 def __init__(self, vocab_size, n_patches, n_subpatches, n_locations, no_more_points_class):541 self.n_locations = n_locations542 self.n_patches = n_patches543 self.n_subpatches = n_subpatches544 self.vocab_size = vocab_size545 546 if no_more_points_class:547 self.no_more_points_token_id = vocab_size + n_patches548 else:549 self.no_more_points_token_id = -1550 self.patch_start = vocab_size551 self.patch_end_without_no_more_points = vocab_size + n_patches552 self.patch_end = vocab_size + n_patches + int(no_more_points_class)553 self.subpatch_start = self.patch_end554 self.subpatch_end = self.subpatch_start + n_subpatches555 self.location_start = self.subpatch_end556 self.location_end = self.subpatch_end + n_locations557 558 559class MolmoPointLogitProcessor(LogitsProcessor):560 """Force point-special tokens to be generated in a valid order"""561 562 def __init__(self, bounds: GeneratedTokenBounds,563 prevent_repeats, force_patch_sorted, force_subpatch_sorted):564 self.bounds = bounds565 self.prevent_repeats = prevent_repeats566 self.force_patch_sorted = force_patch_sorted567 self.force_subpatch_sorted = force_subpatch_sorted568 569 def __call__(self, input_ids, scores):570 b = self.bounds571 is_complete_patch = (b.patch_start <= input_ids) & (input_ids < b.patch_end)572 is_complete_subpatch = (b.subpatch_start <= input_ids) & (input_ids < b.subpatch_end)573 574 if b.n_locations:575 is_complete_patch[:, -2:] = False576 is_complete_subpatch[:, -2:] = False577 else:578 is_complete_patch[:, -1] = False579 is_complete_subpatch[:, -1] = False580 581 for batch in range(len(input_ids)):582 batch_input_ids = input_ids[batch]583 last_token = batch_input_ids[-1]584 585 batch_is_patch_token = is_complete_patch[batch]586 last_predicted_patch_token = batch_input_ids[is_complete_patch[batch]]587 if len(last_predicted_patch_token):588 last_predicted_patch_token = last_predicted_patch_token[-1]589 else:590 last_predicted_patch_token = None591 592 last_predicted_subpatch_token = batch_input_ids[is_complete_subpatch[batch]]593 if len(last_predicted_subpatch_token):594 last_predicted_subpatch_token = last_predicted_subpatch_token[-1]595 else:596 last_predicted_subpatch_token = None597 598 no_more_points = torch.any(batch_input_ids == b.no_more_points_token_id)599 600 if no_more_points:601 # Cannot generate any kind of point602 scores[batch, b.patch_start:b.location_end] = -float("inf")603 elif last_token < b.patch_start or last_token >= b.subpatch_end:604 # Cannot generate subpatch/location, but might generate a patch605 scores[batch, b.subpatch_start:b.location_end] = -float("inf")606 607 if self.force_patch_sorted and last_predicted_patch_token is not None:608 # Cannot generate patches that occurs before the previously predicted patch609 scores[batch, b.patch_start:last_predicted_patch_token] = -float("inf")610 611 if (612 self.prevent_repeats and613 self.force_subpatch_sorted and614 last_predicted_subpatch_token is not None and615 last_predicted_subpatch_token == (b.subpatch_end-1)616 ):617 # Generating `last_predicted_patch_token` would force us to generate a repeat618 # since the only subpatch we can predict while keeping sorted order619 # will repeat the previous point620 scores[batch, last_predicted_patch_token] = -float("inf")621 622 elif b.patch_start <= last_token < b.patch_end:623 # Last token was a patch token, must select a subpatch next624 scores[batch, :b.subpatch_start] = -float("inf")625 scores[batch, b.subpatch_end:] = -float("inf")626 if (627 self.force_subpatch_sorted and628 last_predicted_patch_token == last_token629 ):630 assert last_predicted_subpatch_token is not None631 if self.prevent_repeats:632 assert last_predicted_subpatch_token != b.subpatch_end-1633 scores[batch, b.subpatch_start:last_predicted_subpatch_token+1] = -float("inf")634 else:635 scores[batch, b.subpatch_start:last_predicted_subpatch_token] = -float("inf")636 637 elif b.n_locations and b.subpatch_start <= last_token < b.subpatch_end:638 # Last token was a subpatch token, must select a location next639 scores[batch, :b.location_start] = -float("inf")640 scores[batch, b.location_end:] = -float("inf")641 else:642 raise RuntimeError("Unreachable")643 return scores644 645 646@dataclass647class Molmo2TextBaseOutput(BaseModelOutputWithPast):648 pre_ln_hidden_state: Optional[torch.FloatTensor] = None649 650 651class MolmoPointTextModel(PreTrainedModel):652 config: Molmo2TextConfig653 _no_split_modules = ["Molmo2DecoderLayer", "Molmo2PostNormDecoderLayer"]654 base_model_prefix = "model"655 supports_gradient_checkpointing = True656 _skip_keys_device_placement = "past_key_values"657 _supports_flash_attn = True658 _supports_sdpa = True659 660 _can_compile_fullgraph = True661 _supports_attention_backend = True662 _can_record_outputs = {663 "hidden_states": Molmo2DecoderLayer,664 "attentions": Molmo2Attention,665 }666 667 def __init__(self, config: Molmo2TextConfig):668 super().__init__(config)669 if config.additional_vocab_size is not None:670 self.wte = Molmo2Embedding(671 config.vocab_size,672 config.additional_vocab_size,673 config.hidden_size,674 )675 else:676 self.wte = nn.Embedding(config.vocab_size, config.hidden_size)677 self.emb_drop = nn.Dropout(config.embedding_dropout)678 decoder_layer = Molmo2PostNormDecoderLayer if config.norm_after else Molmo2DecoderLayer679 self.blocks = nn.ModuleList(680 [decoder_layer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]681 )682 self.ln_f = Molmo2RMSNorm(config.hidden_size, eps=config.layer_norm_eps)683 if config.rope_scaling_layers is not None:684 self.rotary_embs = nn.ModuleDict(685 {686 "default": Molmo2RotaryEmbedding(config, rope_type="default"),687 "scaling": Molmo2RotaryEmbedding(config),688 }689 )690 else:691 self.rotary_emb = Molmo2RotaryEmbedding(config)692 self.gradient_checkpointing = False693 694 # Initialize weights and apply final processing695 self.post_init()696 697 def get_input_embeddings(self) -> torch.nn.Module:698 return self.wte699 700 def set_input_embeddings(self, value: torch.nn.Module) -> None:701 self.wte = value702 703 @can_return_tuple704 def forward(705 self,706 input_ids: Optional[torch.LongTensor] = None,707 attention_mask: Optional[torch.Tensor] = None,708 position_ids: Optional[torch.LongTensor] = None,709 past_key_values: Optional[Cache] = None,710 inputs_embeds: Optional[torch.FloatTensor] = None,711 use_cache: Optional[bool] = None,712 output_attentions: Optional[bool] = None,713 output_hidden_states: Optional[bool] = None,714 output_pre_ln_state: Optional[bool] = None,715 cache_position: Optional[torch.LongTensor] = None,716 **kwargs: Unpack[TransformersKwargs],717 ) -> Molmo2TextBaseOutput:718 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions719 output_hidden_states = (720 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states721 )722 use_cache = use_cache if use_cache is not None else self.config.use_cache723 724 if (input_ids is None) ^ (inputs_embeds is not None):725 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")726 727 if self.gradient_checkpointing and self.training and use_cache:728 logger.warning_once(729 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."730 )731 use_cache = False732 733 if inputs_embeds is None:734 input_ids = input_ids * (input_ids != -1).to(input_ids.dtype)735 inputs_embeds = self.wte(input_ids)736 737 # torch.jit.trace() doesn't support cache objects in the output738 if use_cache and past_key_values is None and not torch.jit.is_tracing():739 past_key_values = DynamicCache(config=self.config)740 741 if cache_position is None:742 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0743 cache_position = torch.arange(744 past_seen_tokens,745 past_seen_tokens + inputs_embeds.shape[1],746 device=inputs_embeds.device,747 )748 749 if position_ids is None:750 position_ids = cache_position.unsqueeze(0)751 752 # It may already have been prepared by e.g. `generate`753 if not isinstance(causal_mask_mapping := attention_mask, dict):754 # Prepare mask arguments755 mask_kwargs = {756 "config": self.config,757 "input_embeds": inputs_embeds,758 "attention_mask": attention_mask,759 "cache_position": cache_position,760 "past_key_values": past_key_values,761 "position_ids": position_ids,762 }763 764 # Create the mask765 causal_mask_mapping = create_causal_mask(**mask_kwargs)766 767 hidden_states = inputs_embeds768 769 # create position embeddings to be shared across the decoder layers770 if self.config.rope_scaling_layers is not None:771 position_embeddings_mapping = {772 "default": self.rotary_embs["default"](hidden_states, position_ids),773 "scaling": self.rotary_embs["scaling"](hidden_states, position_ids),774 }775 else:776 position_embeddings = self.rotary_emb(hidden_states, position_ids)777 778 # decoder layers779 all_hidden_states = () if output_hidden_states else None780 all_self_attns = () if output_attentions else None781 782 for layer_idx, decoder_block in enumerate(self.blocks[: self.config.num_hidden_layers]):783 if output_hidden_states:784 all_hidden_states += (hidden_states,)785 786 if self.config.rope_scaling_layers is not None:787 position_embeddings_i = (788 position_embeddings_mapping["scaling"]789 if layer_idx in self.config.rope_scaling_layers790 else position_embeddings_mapping["default"]791 )792 else:793 position_embeddings_i = position_embeddings794 795 layer_outputs = decoder_block(796 hidden_states,797 attention_mask=causal_mask_mapping,798 position_ids=position_ids,799 past_key_values=past_key_values,800 output_attentions=output_attentions,801 use_cache=use_cache,802 cache_position=cache_position,803 position_embeddings=position_embeddings_i,804 **kwargs,805 )806 807 hidden_states = layer_outputs[0]808 809 if output_attentions:810 all_self_attns += (layer_outputs[1],)811 812 pre_ln_state = hidden_states813 hidden_states = self.ln_f(hidden_states)814 815 # add hidden states from the last decoder layer816 if output_hidden_states:817 all_hidden_states += (hidden_states,)818 819 return Molmo2TextBaseOutput(820 last_hidden_state=hidden_states,821 past_key_values=past_key_values,822 pre_ln_hidden_state=pre_ln_state,823 hidden_states=hidden_states,824 attentions=all_self_attns,825 )826 827# Adapted from transformers.models.gemma3.modeling_gemma3828def token_type_ids_mask_function(829 token_type_ids: Optional[torch.Tensor] = None,830) -> Optional[Callable]:831 """832 This function adds the correct offsets to the `q_idx` and `kv_idx` as the torch API can only accept lengths,833 not start and end indices.834 """835 # Do not return an additional mask in this case836 if token_type_ids is None:837 return None838 839 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:840 # If it's 1 for both query and key/value, we are in an image block841 # NOTE: static cache shape goes beyond input seq length, while token_type_ids.shape[1] == input seq length842 # Since vmap doesn't support `if statement` we workaround it with `torch.where`843 safe_idx = torch.where(kv_idx < token_type_ids.shape[1], kv_idx, 0)844 token_type_ids_at_kv_idx = token_type_ids[batch_idx, safe_idx]845 token_type_ids_at_kv_idx = torch.where(kv_idx < token_type_ids.shape[1], token_type_ids_at_kv_idx, 0)846 847 is_image_block = (token_type_ids[batch_idx, q_idx] == 1) & (token_type_ids_at_kv_idx == 1)848 849 # This is bidirectional attention whenever we are dealing with image tokens850 return is_image_block & is_image_block851 852 return inner_mask853 854 855class MolmoPointPadWithLearnedVector(nn.Module):856 """Module that pads vector857 858 Used to add in the no-more-point key value859 """860 def __init__(self, dim: int):861 super().__init__()862 self.dim = dim863 self.vector = nn.Parameter(torch.zeros([dim]))864 865 def reset_parameters(self):866 torch.nn.init.zeros_(self.vector)867 868 def forward(self, x: torch.Tensor) -> torch.Tensor:869 vector = torch.tile(self.vector[None, :], [x.shape[0], 1])870 return torch.concatenate([x, vector[:, None, :]], dim=1)871 872 873class AddPosEmbed(nn.Module):874 875 def __init__(self, in_features: int, n_pos: int) -> None:876 super().__init__()877 self.bias = nn.Parameter(torch.zeros([n_pos, in_features]))878 879 def forward(self, input: torch.Tensor) -> torch.Tensor:880 return input + self.bias[None, :input.shape[-2], :]881 882 883class MolmoPointConnector(nn.Module):884 def __init__(self, config: MolmoPointAdapterConfig, vit_config: Molmo2VitConfig):885 super().__init__()886 self.config = config887 self.n_vit_layers = len(config.vit_layers)888 pool_dim = vit_config.hidden_size * self.n_vit_layers889 self.norm = None890 self.image_projector = ImageProjectorMLP(891 config.hidden_size,892 config.intermediate_size,893 config.text_hidden_size,894 config.hidden_act,895 )896 self.act = ACT2FN[config.hidden_act]897 self.image_pooling_2d = ViTMultiHeadDotProductAttention(898 hidden_size=config.hidden_size,899 num_heads=config.num_attention_heads,900 num_key_value_heads=config.num_key_value_heads,901 head_dim=config.head_dim,902 input_dim=pool_dim,903 float32_attention=config.float32_attention,904 attention_dropout=config.attention_dropout,905 residual_dropout=config.residual_dropout,906 attn_implementation=config._attn_implementation,907 out_layer=False908 )909 if self.config.positional_embeddings:910 self.positional_embeddings = AddPosEmbed(pool_dim, self.config.positional_embeddings)911 else:912 self.positional_embeddings = None913 914 def __call__(self, to_pool, to_pool_mask):915 """916 to_pool: [n_to_pool, pooling_dim, vit_dim]917 to_pool_mask: [n_to_pool, pooling_dim]918 919 returns:920 pooled_features: [n_to_pool, llm_dim]921 """922 cfg = self.config923 924 if self.config.positional_embeddings:925 to_pool = self.positional_embeddings(to_pool)926 927 if self.config.pooling_attention_mask:928 attn_mask = to_pool_mask.reshape([-1, 1, 1, to_pool_mask.shape[-1]])929 else:930 attn_mask = None931 to_pool = to_pool * to_pool_mask.float()[:, :, None]932 933 denom = to_pool_mask.view(-1, to_pool.shape[-2]).float().sum(-1)934 denom = torch.where(denom == 0, 1, denom)935 query = to_pool.sum(-2, keepdim=True) / denom[:, None, None]936 937 pooled_features = self.image_pooling_2d(query, to_pool, attn_mask=attn_mask)938 pooled_features = self.image_projector(pooled_features)939 return pooled_features940 941 942def extract_image_points(output_text, pooling, mappings, no_more_points_class, location, image_sizes):943 """Extract points from MolmoPoint image output text944 945 return points: [n_points, 4] array of (object_id, image_num, x, y) points946 """947 if len(mappings) != len(image_sizes):948 raise ValueError("Mapping and image sizes must have the same length")949 extracted_points = []950 for vit_patch_id, location_id, example_id in get_subpatch_ids(output_text, pooling, no_more_points_class):951 for image_ix, (mapping, (w, h)) in enumerate(zip(mappings, image_sizes)):952 patch_coords = np.argwhere(mapping == int(vit_patch_id))953 if len(patch_coords) == 1:954 p_y, p_x = patch_coords[0]955 if location_id is not None:956 loc_x = location_id // 3957 loc_y = location_id % 3958 p_x += (loc_x+0.5)*0.33959 p_y += (loc_y+0.5)*0.33960 else:961 p_x += 0.5962 p_y += 0.5963 extracted_points.append([964 example_id,965 image_ix,966 (p_x / mapping.shape[1]) * w,967 (p_y / mapping.shape[0]) * h,968 ])969 break970 else:971 logger.error("Invalid patch id encountered")972 return extracted_points973 974 975def extract_video_points(output_text, pooling, mapping, timestamps, no_more_points_class,976 location, video_size):977 """978 Extract points from MolmoPoint video output text979 980 return points: [n_points, 4] array of (object_id, timestamp, x, y) points981 """982 extracted_points = []983 for vit_patch_id, location_id, example_id in get_subpatch_ids(output_text, pooling, no_more_points_class):984 patch_coords = np.argwhere(mapping == int(vit_patch_id))985 if len(patch_coords) == 1:986 frame_ix, p_y, p_x = patch_coords[0]987 if location_id is not None:988 loc_x = location_id // 3989 loc_y = location_id % 3990 p_x += (loc_x+0.5)*0.33991 p_y += (loc_y+0.5)*0.33992 else:993 p_x += 0.5994 p_y += 0.5995 ts = timestamps[frame_ix]996 extracted_points.append([997 example_id,998 ts,999 (p_x / mapping.shape[2]) * video_size[0],1000 (p_y / mapping.shape[1]) * video_size[1]1001 ])1002 else:1003 logger.error("Invalid patch id encountered")1004 return extracted_points1005 1006 1007class MolmoPointModel(MolmoPointPreTrainedModel):1008 base_model_prefix = ""1009 _checkpoint_conversion_mapping = {}1010 # Reference: fix gemma3 grad acc #372081011 accepts_loss_kwargs = False1012 config: MolmoPointConfig1013 1014 def __init__(self, config: MolmoPointConfig):1015 super().__init__(config)1016 self.transformer: MolmoPointTextModel = MolmoPointTextModel(config.text_config)1017 self.patch_token_id = self.config.patch_token_id1018 self.subpatch_token_id = self.config.subpatch_token_id1019 self.location_token_id = self.config.location_token_id1020 1021 vit_config = config.vit_config1022 adapter_config = config.adapter_config1023 self.vit_layers = []1024 for layer in adapter_config.vit_layers:1025 if layer >= 0:1026 self.vit_layers.append(layer)1027 else:1028 self.vit_layers.append(layer + vit_config.num_hidden_layers)1029 1030 last_layer_needed = max(self.vit_layers) + 11031 if last_layer_needed < vit_config.num_hidden_layers:1032 new_vit_config = deepcopy(vit_config)1033 new_vit_config.num_hidden_layers = last_layer_needed1034 self.vit = Molmo2VisionTransformer(new_vit_config)1035 else:1036 self.vit = Molmo2VisionTransformer(vit_config)1037 1038 self.connector = MolmoPointConnector(adapter_config, vit_config)1039 if self.config.embed_selected_vit_patch == "linear":1040 llm_dim = config.text_config.hidden_size1041 vit_dim = self.config.vit_config.hidden_size * len(self.config.adapter_config.vit_layers)1042 self.build_vit_embedding = nn.Linear(vit_dim, llm_dim, bias=True)1043 else:1044 raise NotImplementedError(f"Embedding {self.config.embed_selected_vit_patch} not implemented")1045 self.point_predictor = PointPredictor(config)1046 1047 # Initialize weights and apply final processing1048 self.post_init()1049 1050 def build_token_bounds(self, token_pooling):1051 n_patches, n_subpatches = token_pooling.shape[-2:]1052 return GeneratedTokenBounds(1053 vocab_size=self.config.vocab_size + self.config.text_config.additional_vocab_size,1054 n_patches=n_patches,1055 n_subpatches=n_subpatches,1056 n_locations=9 if self.config.patch_location else 0,1057 no_more_points_class=self.config.no_more_points_class,1058 )1059 1060 def get_input_embeddings(self) -> torch.nn.Module:1061 return self.transformer.wte1062 1063 def set_input_embeddings(self, value: torch.nn.Module) -> None:1064 self.transformer.wte = value1065 1066 def set_decoder(self, decoder):1067 self.transformer = decoder1068 1069 def get_decoder(self):1070 return self.transformer1071 1072 @property1073 def device(self) -> torch.device:1074 return self.transformer.ln_f.weight.device1075 1076 def build_batched_images(1077 self,1078 input_ids: torch.LongTensor,1079 pixel_values: torch.Tensor,1080 image_token_pooling: torch.Tensor,1081 image_grids: torch.Tensor,1082 image_num_crops: torch.Tensor,1083 ) -> tuple[torch.Tensor, torch.Tensor]:1084 # 1) Count the number of images in each example1085 raw_counts = (input_ids == self.config.image_end_token_id).sum(1) # [N]1086 # Each image is represented by global view and high-res view1087 # so we divide by 2 to get the number of images1088 counts = raw_counts // 21089 N = counts.size(0)1090 device = input_ids.device1091 1092 # Total number of images in the batch1093 num_images = int(counts.sum().item())1094 1095 # Sanity check1096 assert image_grids.size(0) == num_images, \1097 f"Expected {num_images} image grids, but got {image_grids.size(0)}"1098 assert image_num_crops.size(0) == num_images, \1099 f"Expected {num_images} image num crops, but got {image_num_crops.size(0)}"1100 1101 # 1-1) Compute per-image pooled patch count from image grids1102 with torch.no_grad():1103 first_prod = image_grids[:, :2].prod(dim=1) # [num_images]1104 second_prod = image_grids[:, 2:].prod(dim=1) # [num_images]1105 num_pooled_patches_per_image = (first_prod + second_prod).to(image_num_crops.dtype) # [num_images]1106 1107 # pixel_values: [n_crops, n_patches, pixels_per_patch]1108 n_crops, n_patches, pixels_per_patch = pixel_values.shape1109 1110 # 2) Map each image index → example index1111 # Example: if counts = [2, 1, 3], then this becomes [0,0,1,2,2,2]1112 example_ids_for_image = torch.arange(N, device=device).repeat_interleave(counts) # [num_images]1113 assert example_ids_for_image.numel() == num_images1114 1115 # 2-1) Compute crops_per_example by summing per-image crop counts1116 crops_per_example = torch.zeros(1117 N, dtype=image_num_crops.dtype, device=image_num_crops.device1118 )1119 crops_per_example.index_add_(0, example_ids_for_image, image_num_crops) # [N]1120 1121 # 2-2) Per-image number of patches = (crops per image) * n_patches1122 patches_per_image = image_num_crops * n_patches # [num_images]1123 1124 # 2-3) Compute per-example per-image patch offsets1125 counts_list = counts.tolist()1126 index_offset_per_example_list = []1127 offset_img = 01128 for c in counts_list:1129 per_img_patches = patches_per_image[offset_img:offset_img + c] # [c]1130 # Offsets: [0, img0_total_patches, img0+img1_total_patches, ...]1131 index_offset = [0] + per_img_patches.cumsum(0).tolist()[:-1]1132 index_offset_per_example_list.append(index_offset)1133 offset_img += c1134 1135 # 2-4) Compute num_pooled_patches_per_example1136 num_pooled_patches_per_example = torch.zeros(1137 N, dtype=num_pooled_patches_per_image.dtype, device=num_pooled_patches_per_image.device1138 )1139 num_pooled_patches_per_example.index_add_(1140 0, example_ids_for_image, num_pooled_patches_per_image1141 )1142 1143 # Sanity checks1144 total_crops = int(crops_per_example.sum().item())1145 assert total_crops == n_crops, \1146 f"Expected {total_crops} crops, but got {n_crops}"1147 1148 total_num_pooled_patches = int(num_pooled_patches_per_example.sum().item())1149 assert total_num_pooled_patches == image_token_pooling.size(0), \1150 f"Expected {total_num_pooled_patches} pooled patches, but got {image_token_pooling.size(0)}"1151 1152 # 3) Build images tensor filled with -11153 M = int(crops_per_example.max().item())1154 images = torch.full(1155 (N, M, n_patches, pixels_per_patch),1156 fill_value=-1,1157 dtype=pixel_values.dtype,1158 device=pixel_values.device,1159 )1160 1161 # 4) Fill images with per-example slices from pixel_values1162 offset_crop = 01163 for i in range(N):1164 num = int(crops_per_example[i].item())1165 cur = pixel_values[offset_crop:offset_crop + num] # [num, n_patches, pixels_per_patch]1166 images[i, :num] = cur1167 offset_crop += num1168 1169 # Sanity check1170 assert offset_crop == n_crops1171 1172 # 5) Build new_token_pooling tensor filled with -11173 P = int(num_pooled_patches_per_example.max().item())1174 _, dim = image_token_pooling.shape1175 new_token_pooling = torch.full(1176 (N, P, dim),1177 fill_value=-1,1178 dtype=image_token_pooling.dtype,1179 device=image_token_pooling.device,1180 )1181 1182 # 6) Fill token_pooling with per-example slices, adding per-image patch offsets1183 patch_offset = 01184 img_offset = 01185 1186 for i, c in enumerate(counts_list):1187 num_patches = int(num_pooled_patches_per_example[i].item())1188 1189 # Subsequence of pooled tokens belonging to this example1190 cur = image_token_pooling[patch_offset:patch_offset + num_patches].clone() # [num_patches, dim]1191 1192 index_offset_per_example = index_offset_per_example_list[i] # length = c1193 per_img_pooled = num_pooled_patches_per_image[img_offset:img_offset + c] # [c]1194 1195 assert len(index_offset_per_example) == per_img_pooled.numel()1196 1197 # Apply per-image offsets to the (ragged) subsequence1198 offset = 01199 for j in range(c):1200 index_offset = int(index_offset_per_example[j])