import math import re from copy import deepcopy from dataclasses import dataclass from typing import Optional, Union, Callable, Any, List, Tuple import numpy as np import torch from torch import nn from torch.nn import functional as F from transformers import LogitsProcessorList, LogitsProcessor, AutoProcessor, ViTConfig from transformers.image_utils import PILImageResampling 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 from .image_processing_molmo2 import Molmo2ImagesKwargs, image_to_patches_and_grids from .modeling_molmo2 import ImageProjectorMLP, Molmo2VisionTransformer, Molmo2RMSNorm, \ Molmo2RotaryEmbedding, Molmo2PostNormDecoderLayer, Molmo2DecoderLayer, Molmo2Attention, \ Molmo2Embedding # FIXME remove processor = None def decode(ids): global processor if processor is None: processor = AutoProcessor.from_pretrained( "/weka/oe-training-default/mm-olmo/released-models-molmo2-point-0326/MolmoPoint-8B/hf-step2000", trust_remote_code=True, padding_side="left") return processor.post_process_image_text_to_text(ids.view(1), skip_special_tokens=False, clean_up_tokenization_spaces=False)[0] logger = logging.get_logger(__name__) NO_POINTS_LABEL = 1000000 EXTRACT_POINT_TRIPLE = re.compile(f" ? ? ?([0-9]+)" ) def get_subpatch_ids(output_text, pooling, no_more_points_class): n_patches, n_subpatches = pooling.shape[-2:] if no_more_points_class: n_patches += 1 for match in EXTRACT_POINT_TRIPLE.finditer(output_text): patch_id, subpatch_num = int(match.group(1)), int(match.group(2)) subpatch_id = subpatch_num - n_patches location_num = int(match.group(3)) location_id = location_num - n_patches - n_subpatches example_id = int(match.group(4)) vit_patch_id = pooling[patch_id, subpatch_id] yield vit_patch_id, location_id, example_id @dataclass class ImageCache: """Extra stuff we need to cache when doing autoregressive generation with pointing""" patch_k: torch.FloatTensor """K values of the image tokens""" patch_k_mask: torch.BoolTensor """Mask over image tokens that can be selected""" subpatch_k: torch.FloatTensor """K values of the ViT patches before pooling""" token_pooling: torch.LongTensor """token pooling array mapping image_patch_id -> ViT patches pooled for that patch""" vit_features: torch.FloatTensor """Features before pooling, used for building input embeddings""" image_pos_ids: Optional[torch.LongTensor] = None """Position ids of the image tokens if need for rotary embeddings""" image_features0: Optional[torch.FloatTensor] = None """"Image features, might be needed to embed new patch prediction tokens""" flat_image_tokens_to_flat_image_features: Optional[torch.LongTensor] = None """Cached for indexing uses""" @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 image_data: Optional[ImageCache] = None patch_logits: Optional[torch.FloatTensor] = None subpatch_logits: Optional[torch.FloatTensor] = None location_logits: Optional[torch.FloatTensor] = None last_predicted_patch_id: Optional[torch.LongTensor] = 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 image_data: Optional[ImageCache] = None patch_logits: Optional[torch.FloatTensor] = None subpatch_logits: Optional[torch.FloatTensor] = None location_logits: Optional[torch.FloatTensor] = None input_ids: Optional[torch.LongTensor] = None last_predicted_patch_id: Optional[torch.LongTensor] = None class MolmoPointPatchRope(nn.Module): inv_freq: torch.Tensor # fix linting for `register_buffer` def __init__( self, theta: float, dim: int, device: Union[str, torch.device] = None, ): super().__init__() attention_factor = 1.0 # Unused in this type of RoPE inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) def rotate_half(self, x: torch.Tensor) -> torch.Tensor: B, hs = x.size() x = x.view(B, 2, hs // 2) x1, x2 = x.unbind(dim=-2) return torch.cat((-x2, x1), dim=-1) @torch.no_grad() def forward(self, x, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: inv_freq_expanded = self.inv_freq.float().to(x.device) position_ids_expanded = position_ids.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 x = x.float() freqs = position_ids_expanded[:, None] * inv_freq_expanded[None, :] emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() out = ((x * cos) + (self.rotate_half(x) * sin)) return out.to(dtype=x.dtype) 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.wo = 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 PointPredictor(nn.Module): """Point predictor logic""" # We separate this out so accelerate will co-locate all these parameters on the same device def __init__(self, config): super().__init__() self.config = config llm_dim = config.text_config.hidden_size patch_embed_dim = config.patch_embed_dim vit_dim = self.config.vit_config.hidden_size * len(self.config.adapter_config.vit_layers) 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 if self.config.token_prediction_rotary == "none": self.patch_rotary = None else: theta = self.config.token_prediction_rotary_theta or self.config.llm.rope_theta if self.config.token_prediction_rotary == "one_d": self.patch_rotary = MolmoPointPatchRope(theta, self.config.patch_embed_dim) else: raise NotImplementedError() self.patch_q = nn.Linear(llm_dim, patch_embed_dim) self.patch_k = nn.Linear(llm_dim, patch_embed_dim) self.subpatch_q = nn.Linear(llm_dim, patch_embed_dim) self.subpatch_k = nn.Linear(vit_dim, patch_embed_dim) self.add_no_point_class_embed = MolmoPointPadWithLearnedVector(patch_embed_dim) 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") def forward( self, x, token_pooling, is_image_token, is_patch, is_subpatch, is_indexable_image_token, vit_features, vit_features_mask, image_features_mask, input_patch_ids, last_predicted_patch_id, image_data: ImageCache ): dim = self.config.text_config.hidden_size batch_size = x.shape[0] if self.x_norm is not None: x_norm = self.x_norm(x) elif self.config.norm_x: x_norm = x / math.sqrt(dim) else: x_norm = x # Build the keys, or get them from the cache if image_data is not None: patch_k, subpatch_k = image_data.patch_k, image_data.subpatch_k patch_k_mask = image_data.patch_k_mask token_pooling = image_data.token_pooling vit_features_mask = token_pooling >= 0 image_pos_ids = image_data.image_pos_ids else: # Build patch keys, this takes a bit of indexing trickery since we want the keys in # shape [batch, n_image_tokens] not [batch, sequence_length] n_image_tokens = token_pooling.shape[1] patch_k_flat = self.patch_k(x_norm.view(-1, dim)[is_image_token.view(-1)]) if self.patch_rotary is not None: image_token_indices = torch.cumsum(is_indexable_image_token, dim=-1) - 1 image_pos_ids_flat = image_token_indices.view(-1)[is_image_token.view(-1)] patch_k_flat = self.patch_rotary(patch_k_flat, image_pos_ids_flat) # Computed for use with the query vectors image_pos_ids = torch.zeros([batch_size, n_image_tokens], dtype=torch.long, device=image_pos_ids_flat.device) image_pos_ids.view(-1)[image_features_mask.view(-1)] = image_pos_ids_flat else: image_pos_ids = None patch_k = torch.zeros([batch_size, n_image_tokens, patch_k_flat.shape[-1]], dtype=x.dtype, device=x.device) patch_k.view(-1, patch_k_flat.shape[-1])[image_features_mask.flatten()] = patch_k_flat.to(dtype=x.dtype) patch_k_mask = image_features_mask.clone() patch_k_mask.view(-1)[image_features_mask.view(-1)] = ( is_indexable_image_token.view(-1)[is_image_token.view(-1)]) if self.config.no_more_points_class: patch_k = self.add_no_point_class_embed(patch_k) patch_k_mask = F.pad(patch_k_mask, (0, 1), value=True) subpatch_k = self.subpatch_k(vit_features) patch_logits, subpatch_logits, location_logits = None, None, None if image_data is not None: # Predict patch locations, only done after pre-filling batch_idx = torch.arange(batch_size, device=x_norm.device) image_q = self.patch_q(x_norm) if self.patch_rotary is not None and last_predicted_patch_id is not None: rotate_by = image_pos_ids[batch_idx, last_predicted_patch_id] rotate_by = torch.where(last_predicted_patch_id >= 0, rotate_by, 0) rotate_by = rotate_by.squeeze(-1) image_q = self.patch_rotary( image_q.view(-1, image_q.shape[-1]), torch.clamp(rotate_by, min=0), ).reshape(batch_size, -1, image_q.shape[-1]) dots = torch.matmul(image_q, patch_k.transpose(1, 2)) # [batch, 1, num_images] if self.config.norm_logits: dots = dots / math.sqrt(dots.shape[-1]) valid = patch_k_mask[:, None, :] patch_logits = torch.where(valid, dots, -100000000) if torch.any(is_patch): if x_norm.shape[1] != 1: raise NotImplementedError() subpatch_point_q = self.subpatch_q(x_norm.squeeze(1)) subpatch_k = subpatch_k[batch_idx, input_patch_ids.squeeze(1)] subpatch_logits = torch.einsum("pd,pcd->pc", subpatch_point_q, subpatch_k) if self.config.norm_logits: subpatch_logits = subpatch_logits / math.sqrt(patch_k.shape[-1]) subpatch_mask = vit_features_mask[batch_idx, input_patch_ids.squeeze(1)] subpatch_logits = torch.where(subpatch_mask, subpatch_logits, -100000) subpatch_logits = subpatch_logits[:, None, :] if torch.any(is_subpatch): location_logits = self.subpatch_loc_k(x) if image_data is None: image_data = ImageCache( patch_k=patch_k, subpatch_k=subpatch_k, vit_features=vit_features, patch_k_mask=patch_k_mask, token_pooling=token_pooling, image_pos_ids=image_pos_ids, ) return patch_logits, subpatch_logits, location_logits, image_data class MolmoPointPreTrainedModel(PreTrainedModel): config: MolmoPointConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = [ "Molmo2DecoderLayer", "Molmo2PostNormDecoderLayer", "Molmo2VisionBlock", "ViTMultiHeadDotProductAttention", "PointPredictor" ] _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 GeneratedTokenBounds: """Describes what tokens id ranges are patch/subpatch/location tokens""" def __init__(self, vocab_size, n_patches, n_subpatches, n_locations, no_more_points_class): self.n_locations = n_locations self.n_patches = n_patches self.n_subpatches = n_subpatches self.vocab_size = vocab_size if no_more_points_class: self.no_more_points_token_id = vocab_size + n_patches else: self.no_more_points_token_id = -1 self.patch_start = vocab_size self.patch_end_without_no_more_points = vocab_size + n_patches self.patch_end = vocab_size + n_patches + int(no_more_points_class) self.subpatch_start = self.patch_end self.subpatch_end = self.subpatch_start + n_subpatches self.location_start = self.subpatch_end self.location_end = self.subpatch_end + n_locations class MolmoPointLogitProcessor(LogitsProcessor): """Force point-special tokens to be generated in a valid order""" def __init__(self, bounds: GeneratedTokenBounds, prevent_repeats, force_patch_sorted, force_subpatch_sorted): self.bounds = bounds self.prevent_repeats = prevent_repeats self.force_patch_sorted = force_patch_sorted self.force_subpatch_sorted = force_subpatch_sorted def __call__(self, input_ids, scores): b = self.bounds is_complete_patch = (b.patch_start <= input_ids) & (input_ids < b.patch_end) is_complete_subpatch = (b.subpatch_start <= input_ids) & (input_ids < b.subpatch_end) if b.n_locations: is_complete_patch[:, -2:] = False is_complete_subpatch[:, -2:] = False else: is_complete_patch[:, -1] = False is_complete_subpatch[:, -1] = False for batch in range(len(input_ids)): batch_input_ids = input_ids[batch] last_token = batch_input_ids[-1] batch_is_patch_token = is_complete_patch[batch] last_predicted_patch_token = batch_input_ids[is_complete_patch[batch]] if len(last_predicted_patch_token): last_predicted_patch_token = last_predicted_patch_token[-1] else: last_predicted_patch_token = None last_predicted_subpatch_token = batch_input_ids[is_complete_subpatch[batch]] if len(last_predicted_subpatch_token): last_predicted_subpatch_token = last_predicted_subpatch_token[-1] else: last_predicted_subpatch_token = None no_more_points = torch.any(batch_input_ids == b.no_more_points_token_id) if no_more_points: # Cannot generate any kind of point scores[batch, b.patch_start:b.location_end] = -float("inf") elif last_token < b.patch_start or last_token >= b.subpatch_end: # Cannot generate subpatch/location, but might generate a patch scores[batch, b.subpatch_start:b.location_end] = -float("inf") if self.force_patch_sorted and last_predicted_patch_token is not None: # Cannot generate patches that occurs before the previously predicted patch scores[batch, b.patch_start:last_predicted_patch_token] = -float("inf") if ( self.prevent_repeats and self.force_subpatch_sorted and last_predicted_subpatch_token is not None and last_predicted_subpatch_token == (b.subpatch_end-1) ): # Generating `last_predicted_patch_token` would force us to generate a repeat # since the only subpatch we can predict while keeping sorted order # will repeat the previous point scores[batch, last_predicted_patch_token] = -float("inf") elif b.patch_start <= last_token < b.patch_end: # Last token was a patch token, must select a subpatch next scores[batch, :b.subpatch_start] = -float("inf") scores[batch, b.subpatch_end:] = -float("inf") if ( self.force_subpatch_sorted and last_predicted_patch_token == last_token ): assert last_predicted_subpatch_token is not None if self.prevent_repeats: assert last_predicted_subpatch_token != b.subpatch_end-1 scores[batch, b.subpatch_start:last_predicted_subpatch_token+1] = -float("inf") else: scores[batch, b.subpatch_start:last_predicted_subpatch_token] = -float("inf") elif b.n_locations and b.subpatch_start <= last_token < b.subpatch_end: # Last token was a subpatch token, must select a location next scores[batch, :b.location_start] = -float("inf") scores[batch, b.location_end:] = -float("inf") else: raise RuntimeError("Unreachable") return scores @dataclass class Molmo2TextBaseOutput(BaseModelOutputWithPast): pre_ln_hidden_state: Optional[torch.FloatTensor] = None 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, output_pre_ln_state: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> Molmo2TextBaseOutput: 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],) pre_ln_state = hidden_states 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 Molmo2TextBaseOutput( last_hidden_state=hidden_states, past_key_values=past_key_values, pre_ln_hidden_state=pre_ln_state, hidden_states=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.image_projector(pooled_features) return pooled_features def extract_image_points(output_text, pooling, mappings, no_more_points_class, location, image_sizes): """Extract points from MolmoPoint image output text return points: [n_points, 4] array of (object_id, image_num, x, y) points """ if len(mappings) != len(image_sizes): raise ValueError("Mapping and image sizes must have the same length") extracted_points = [] for vit_patch_id, location_id, example_id in get_subpatch_ids(output_text, pooling, no_more_points_class): for image_ix, (mapping, (w, h)) in enumerate(zip(mappings, image_sizes)): patch_coords = np.argwhere(mapping == int(vit_patch_id)) if len(patch_coords) == 1: p_y, p_x = patch_coords[0] if location_id is not None: loc_x = location_id // 3 loc_y = location_id % 3 p_x += (loc_x+0.5)*0.33 p_y += (loc_y+0.5)*0.33 else: p_x += 0.5 p_y += 0.5 extracted_points.append([ example_id, image_ix, (p_x / mapping.shape[1]) * w, (p_y / mapping.shape[0]) * h, ]) break else: logger.error("Invalid patch id encountered") return extracted_points def extract_video_points(output_text, pooling, mapping, timestamps, no_more_points_class, location, video_size): """ Extract points from MolmoPoint video output text return points: [n_points, 4] array of (object_id, timestamp, x, y) points """ extracted_points = [] for vit_patch_id, location_id, example_id in get_subpatch_ids(output_text, pooling, no_more_points_class): patch_coords = np.argwhere(mapping == int(vit_patch_id)) if len(patch_coords) == 1: frame_ix, p_y, p_x = patch_coords[0] if location_id is not None: loc_x = location_id // 3 loc_y = location_id % 3 p_x += (loc_x+0.5)*0.33 p_y += (loc_y+0.5)*0.33 else: p_x += 0.5 p_y += 0.5 ts = timestamps[frame_ix] extracted_points.append([ example_id, ts, (p_x / mapping.shape[2]) * video_size[0], (p_y / mapping.shape[1]) * video_size[1] ]) else: logger.error("Invalid patch id encountered") return extracted_points 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) self.patch_token_id = self.config.patch_token_id self.subpatch_token_id = self.config.subpatch_token_id self.location_token_id = self.config.location_token_id 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) if self.config.embed_selected_vit_patch == "linear": llm_dim = config.text_config.hidden_size vit_dim = self.config.vit_config.hidden_size * len(self.config.adapter_config.vit_layers) 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") self.point_predictor = PointPredictor(config) # Initialize weights and apply final processing self.post_init() def build_token_bounds(self, token_pooling): n_patches, n_subpatches = token_pooling.shape[-2:] return GeneratedTokenBounds( vocab_size=self.config.vocab_size + self.config.text_config.additional_vocab_size, n_patches=n_patches, n_subpatches=n_subpatches, n_locations=9 if self.config.patch_location else 0, no_more_points_class=self.config.no_more_points_class, ) 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 @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, image_data: Optional[ImageCache] = None, last_predicted_patch_id: Optional[torch.LongTensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> Union[tuple, MolmoPointModelOutputWithPast]: """ last_point_patch_id: The patch id the last generated point pointed to """ 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 inputs_embeds is not None: raise NotImplementedError("Custom inputs_embeds is not implemented yet") input_ids = input_ids * (input_ids != -1).to(input_ids.dtype) if image_data is not None: # Figure out where the patch/subpatch/location are and their values, and then convert # the input_ids back into their original special token values can_point = True bounds = self.build_token_bounds(image_data.token_pooling) expanded_inputs = input_ids is_patch = (input_ids >= bounds.patch_start) & (input_ids < bounds.patch_end_without_no_more_points) is_no_more_points = (input_ids == bounds.no_more_points_token_id) is_subpatch = (input_ids >= bounds.subpatch_start) & (input_ids < bounds.subpatch_end) is_location = (input_ids >= bounds.location_start) & (input_ids < bounds.location_end) input_patch_ids = torch.where(is_patch, input_ids - bounds.patch_start, -1) input_subpatch_ids = torch.where(is_subpatch, input_ids - bounds.subpatch_start, -1) input_ids = torch.where(is_patch | is_no_more_points, self.patch_token_id, input_ids) input_ids = torch.where(is_subpatch, self.subpatch_token_id, input_ids) input_ids = torch.where(is_location, self.location_token_id, input_ids) else: # No patch prediction during pre-filling input_subpatch_ids = None input_patch_ids = None is_patch = None is_subpatch = None can_point = False device = input_ids.device x = self.transformer.wte(input_ids).to(device=device) batch_size, _, dim = x.shape batch_idx = torch.arange(batch_size, device=device) vit_features_flat: Optional[torch.FloatTensor] = None if images is not None: is_indexable_image_token = input_ids == self.config.image_patch_id is_non_indexable_image_token = input_ids == self.config.image_non_indexable_patch_id is_image_token = is_indexable_image_token | is_non_indexable_image_token images = images.to(device=self.device, dtype=self.dtype) B, T, N, D = images.shape images = images.view(B * T, N, D) vit_image_features = self.vit(images) features = [] for layer in self.vit_layers: features.append(vit_image_features[layer]) vit_features = torch.cat(features, dim=-1).to(device=device) vit_feature_dim = vit_features.shape[-1] # Gather the features that should be pooled to build patch embeddings vit_features = vit_features.reshape(batch_size, -1, vit_feature_dim)[batch_idx[:, None, None], torch.clip(token_pooling, 0)] vit_features = vit_features * (token_pooling >= 0).float()[:, :, :, None] vit_features_mask = token_pooling >= 0 # Build the sparse version which will be passed to the connector # Now shape [num_image_tokens_in_batch, pooling_dim, dim] image_features_mask = torch.any(vit_features_mask, -1) vit_features_flat = vit_features.reshape([-1, token_pooling.shape[-1], vit_features.shape[-1]]) vit_features_flat = vit_features_flat[image_features_mask.view(-1)] vit_features_to_flat_mask = vit_features_mask.view(-1, token_pooling.shape[-1])[image_features_mask.view(-1)] # Finally, apply the connector and add to input embeddings image_features = self.connector(vit_features_flat, vit_features_to_flat_mask).to(device=device) x = x.clone() x.view(-1, dim)[is_image_token.view(-1)] += image_features.view(-1, dim) else: is_image_token = None is_indexable_image_token = None if image_data is not None: # Get the features/masks from the cache token_pooling = image_data.token_pooling.to(device=device) vit_features_mask = token_pooling >= 0 image_features_mask = torch.any(vit_features_mask, -1) vit_features = image_data.vit_features.to(device=device) else: vit_features = None vit_features_mask = None image_features_mask = None # Embed the points if can_point: image_token_offset = image_data.flat_image_tokens_to_flat_image_features should_embed = (input_patch_ids >= 0) and (input_patch_ids < (bounds.patch_end-1)) input_patch_ids_flat = (input_patch_ids + image_token_offset).view(-1)[should_embed.view(-1)] x.view(-1, dim)[is_patch.view(-1)] += image_data.image_features0.view(-1, dim)[input_patch_ids_flat] if torch.any(is_subpatch): vit_features_flat = vit_features.reshape([-1, token_pooling.shape[-1], vit_features.shape[-1]]) vit_features_flat = vit_features_flat[image_features_mask.view(-1)] assert last_predicted_patch_id is not None, "Patch should always be generated before a subpatch" for_patches = (last_predicted_patch_id.view(batch_size) + image_token_offset)[input_subpatch_ids.view(batch_size) >= 0] vit_features_to_embed = vit_features_flat[for_patches, input_subpatch_ids] x.view(-1, dim)[is_subpatch.view(-1)] = self.build_vit_embedding(vit_features_to_embed).to(device=device, dtype=x.dtype) # shape: (batch_size, seq_len, d_model) x = self.transformer.emb_drop(x) # type: ignore 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, ) # 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 ) # 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": x, "attention_mask": attention_mask, "cache_position": cache_position, "past_key_values": past_key_values, "position_ids": position_ids, } 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=x, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, cache_position=cache_position, output_pre_ln_state=True, **kwargs, ) x = outputs.pre_ln_hidden_state patch_logits = None subpatch_logits = None location_logits = None if images is not None or image_data is not None: patch_logits, subpatch_logits, location_logits, image_data = self.point_predictor( x, token_pooling, is_image_token, is_patch, is_subpatch, is_indexable_image_token, vit_features, vit_features_mask, image_features_mask, input_patch_ids, last_predicted_patch_id, image_data ) if images is not None: # Also cache stuff we need to building the patch/subpatch token embeddings image_data.image_features0 = image_features num_image_tokens = is_image_token.sum(-1) image_token_offset = torch.cumsum(num_image_tokens[:-1], 0) image_token_offset = F.pad(image_token_offset, [1, 0]) image_data.flat_image_tokens_to_flat_image_features = image_token_offset if last_predicted_patch_id is not None: last_predicted_patch_id = torch.where(input_patch_ids == -1, last_predicted_patch_id, input_patch_ids) else: last_predicted_patch_id = input_patch_ids 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, image_data=image_data, patch_logits=patch_logits, subpatch_logits=subpatch_logits, location_logits=location_logits, last_predicted_patch_id=last_predicted_patch_id, ) class ExtendedLmHead(nn.Module): def __init__(self, config): super().__init__() 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])) def __call__(self, hidden_states, slice_indices=None): lm_head = torch.concatenate([self.output_embeddings, self.new_output_embeddings], dim=0) return F.linear(hidden_states[:, slice_indices, :], lm_head) 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.lm_head = ExtendedLmHead(config) self.vocab_size = config.vocab_size # Initialize weights and apply final processing self.post_init() def build_logit_processor_from_inputs(self, inputs) -> LogitsProcessorList: if inputs.get("image_token_pooling") is not None: pooling = inputs["image_token_pooling"] elif inputs.get("video_token_pooling") is not None: pooling = inputs["video_token_pooling"] else: return [] return [self.build_logit_processor(pooling)] def build_logit_processor(self, token_pooling): return MolmoPointLogitProcessor( bounds=self.model.build_token_bounds(token_pooling), prevent_repeats=self.config.mask_repeats in ["all", "inference"], force_patch_sorted=self.config.mask_patches in ["always", "inference"], force_subpatch_sorted=self.config.mask_subpatches in ["always", "inference"], ) def extract_image_points(self, output_text, pooling, subpatch_mapping, image_sizes): return extract_image_points( output_text, pooling, subpatch_mapping, self.config.no_more_points_class, self.config.patch_location, image_sizes) def extract_video_points(self, output_text, pooling, subpatch_mapping, timestamps, video_size): return extract_video_points( output_text, pooling, subpatch_mapping, timestamps, self.config.no_more_points_class, self.config.patch_location, video_size) 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, image_data: Optional[ImageCache] = None, last_predicted_patch_id: Optional[torch.LongTensor] = None, **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: MolmoPointModelOutputWithPast = 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, image_data=image_data, last_predicted_patch_id=last_predicted_patch_id, **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=slice_indices) loss = None if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.vocab_size) bs, seq, _ = logits.shape if image_data is not None: token_pooling = image_data.token_pooling else: token_pooling = video_token_pooling if video_token_pooling is not None else image_token_pooling n_patches, n_subpatches = token_pooling.shape[-2:] if self.config.no_more_points_class: n_patches += 1 small_val = -100000 # The patch token is a bit tricky since we train the model to first select whether to # generate a patch token or not, and then to select the patch, but this two-stage # process is hard to emulate in generation frameworks # Our hack here is to assume that, if we generate a TOKEN, we always select the argmax # patch. Then we can use PATCH_TOKEN scores as the argmax's patch scores device = logits.device predicted_tokens = torch.argmax(logits[:, -1], dim=-1) patch_token_logits = torch.clone(logits[:, :, self.config.patch_token_id]) logits[:, :, self.config.patch_token_id] = small_val predicted_patch = predicted_tokens == self.config.patch_token_id argmax_patch_logits = torch.full([bs, seq, n_patches], small_val, dtype=logits.dtype, device=device) if outputs.patch_logits is not None: selected_patches = torch.argmax(outputs.patch_logits, -1).to(device=device) bs, seq, n_patches = outputs.patch_logits.shape batch_idx = torch.arange(outputs.patch_logits.shape[0], device=device) seq_ix = torch.arange(outputs.patch_logits.shape[1], device=device) argmax_patch_logits[batch_idx.view(-1, 1, 1), seq_ix.view(1, -1, 1), selected_patches] = patch_token_logits logits[:, :, self.config.subpatch_token_id] = small_val if outputs.subpatch_logits is not None: subpatch_logits = outputs.subpatch_logits else: subpatch_logits = torch.full([bs, seq, n_subpatches], small_val, dtype=logits.dtype, device=device) logits[:, :, self.config.location_token_id] = small_val if outputs.location_logits is not None: location_logits = outputs.location_logits else: location_logits = torch.full([bs, seq, 9], small_val, dtype=logits.dtype, device=device) logits = torch.concatenate([ logits, argmax_patch_logits, subpatch_logits.to(device=device), location_logits.to(device=device) ], -1) 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, image_data=outputs.image_data, patch_logits=outputs.patch_logits, subpatch_logits=outputs.subpatch_logits, location_logits=outputs.location_logits, last_predicted_patch_id=outputs.last_predicted_patch_id, ) 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, image_data: Optional[ImageCache] = 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, image_data=image_data, **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 def _update_model_kwargs_for_generation( self, outputs: MolmoPointModelOutputWithPast, model_kwargs: dict[str, Any], is_encoder_decoder: bool = False, num_new_tokens: int = 1, ) -> dict[str, Any]: args = super()._update_model_kwargs_for_generation(outputs, model_kwargs, is_encoder_decoder, num_new_tokens) if outputs.image_data is not None: args["image_data"] = outputs.image_data args["last_predicted_patch_id"] = outputs.last_predicted_patch_id return args # 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)