diff --git "a/modeling_molmo_point.py" "b/modeling_molmo_point.py" new file mode 100644--- /dev/null +++ "b/modeling_molmo_point.py" @@ -0,0 +1,2356 @@ +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 +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 + +# 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 + + +def extract_image_points_from_pooling(output_text, pooling, mappings, no_more_points_class, location, image_sizes): + 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_from_pooling(output_text, pooling, mapping, timestamps, no_more_points_class, + location, video_size): + 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 + + +@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 ViTMLP(nn.Module): + def __init__(self, dim: int, hidden_dim: int, hidden_act: str, device: Union[str, torch.device] = None): + super().__init__() + self.w1 = nn.Linear(dim, hidden_dim, bias=True, device=device) + self.act = ACT2FN[hidden_act] + self.w2 = nn.Linear(hidden_dim, dim, bias=True, device=device) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w2(self.act(self.w1(x))) + + +class ViTMultiHeadDotProductAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_key_value_heads: int, + head_dim: int, + use_bias: bool = True, + input_dim: Optional[int] = None, + float32_attention: bool = True, + attention_dropout: float = 0.0, + residual_dropout: float = 0.0, + device: Union[str, torch.device] = None, + attn_implementation: str = "eager", + out_layer: bool=True + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = head_dim + self.num_key_value_heads = num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.attn_implementation = attn_implementation + self.is_causal = False + + input_dim = input_dim or hidden_size + + self.wq = nn.Linear( + input_dim, + self.num_heads * self.head_dim, + bias=use_bias, + device=device, + ) + self.wk = nn.Linear( + input_dim, + self.num_key_value_heads * self.head_dim, + bias=use_bias, + device=device, + ) + self.wv = nn.Linear( + input_dim, + self.num_key_value_heads * self.head_dim, + bias=use_bias, + device=device, + ) + if out_layer: + self.wo = nn.Linear( + self.num_heads * self.head_dim, + self.hidden_size, + ) + else: + self.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 Molmo2VisionBlock(nn.Module): + + def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None): + super().__init__() + self.attention = ViTMultiHeadDotProductAttention( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + head_dim=config.head_dim, + float32_attention=config.float32_attention, + attention_dropout=config.attention_dropout, + residual_dropout=config.residual_dropout, + device=device, + attn_implementation=config._attn_implementation, + ) + self.feed_forward = ViTMLP(config.hidden_size, config.intermediate_size, config.hidden_act, device=device) + self.attention_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device) + self.ffn_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.attention(self.attention_norm(x)) + x = x + self.feed_forward(self.ffn_norm(x)) + return x + + +class Molmo2VisionBlockCollection(nn.Module): + + def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None): + super().__init__() + self.conifg = config + self.resblocks = nn.ModuleList([ + Molmo2VisionBlock(config, device) for _ in range(config.num_hidden_layers) + ]) + + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: + hidden_states = [] + for r in self.resblocks: + x = r(x) + hidden_states.append(x) + return hidden_states + + +class Molmo2VisionTransformer(nn.Module): + + def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None): + super().__init__() + self.config = config + + # positional embeddings + self.scale = config.hidden_size ** -0.5 + self.num_prefix_tokens: int = 0 # no class embeddings + self.positional_embedding = nn.Parameter( + torch.zeros(config.image_num_pos, config.hidden_size, device=device), + ) + + image_patch_size = config.image_patch_size + self.patch_embedding = nn.Linear( + image_patch_size * image_patch_size * 3, + config.hidden_size, + bias=True, + device=device, + ) + + self.transformer = Molmo2VisionBlockCollection(config, device) + + def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor: + pos_emb = self.positional_embedding + + pos_emb = pos_emb.reshape( + (int(math.sqrt(pos_emb.shape[0])), int(math.sqrt(pos_emb.shape[0])), pos_emb.shape[1]) + ) + + (patch_num_0, patch_num_1) = patch_num + + if pos_emb.shape[0] != patch_num_0 or pos_emb.shape[1] != patch_num_1: + # Dervied from https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py + # antialias: default True in jax.image.resize + pos_emb = pos_emb.unsqueeze(0).permute(0, 3, 1, 2) + pos_emb = F.interpolate( + pos_emb, size=(patch_num_0, patch_num_1), mode="bicubic", align_corners=False, antialias=True, + ) + pos_emb = pos_emb.permute(0, 2, 3, 1).squeeze(0) + + pos_emb = pos_emb.reshape(-1, pos_emb.shape[-1]) + x = x + pos_emb[None, :, :].to(x.dtype) + return x + + def forward(self, x: torch.Tensor, patch_num: int = None) -> list[torch.Tensor]: + """ + : param x: (batch_size, num_patch, n_pixels) + """ + if patch_num is None: + patch_num = self.config.image_num_patch + + B, N, D = x.shape + + x = self.patch_embedding(x) + + # class embeddings and positional embeddings + x = self.add_pos_emb(x, patch_num) + + hidden_states = self.transformer(x) + return hidden_states + + +class ImageProjectorMLP(nn.Module): + + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + hidden_act: str, + device: Union[str, torch.device] = None, + ): + super().__init__() + self.w1 = nn.Linear(input_dim, hidden_dim, bias=False, device=device) + self.w2 = nn.Linear(hidden_dim, output_dim, bias=False, device=device) + self.w3 = nn.Linear(input_dim, hidden_dim, bias=False, device=device) + self.act = ACT2FN[hidden_act] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w2(self.act(self.w1(x)) * self.w3(x)) + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class Molmo2RotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__( + self, + config: Molmo2TextConfig, + device: Union[str, torch.device] = None, + rope_type: Optional[str] = None, + ): + super().__init__() + if rope_type is not None: + self.rope_type = rope_type + elif hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict): + # BC: "rope_type" was originally "type" + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + inv_freq, self.attention_scaling = rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Molmo2RMSNorm(nn.Module): + + def __init__( + self, + size: int, + eps: float = 1e-6, + device: Union[str, torch.device] = None, + ): + super().__init__() + self.weight = nn.Parameter(torch.ones(size, device=device)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + with torch.autocast(enabled=False, device_type=x.device.type): + og_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + x = x.to(og_dtype) + + return self.weight * x + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.eps}" + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class Molmo2Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Molmo2TextConfig, layer_idx: int) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.num_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.head_dim = config.head_dim + self.scaling = self.head_dim**-0.5 + self.is_causal = True + + self.fused_dims = ( + config.num_attention_heads * config.head_dim, + config.head_dim * config.num_key_value_heads, + config.head_dim * config.num_key_value_heads, + ) + self.att_proj = nn.Linear( + config.hidden_size, + sum(self.fused_dims), + bias=config.qkv_bias, + ) + + # Layer norms. + self.k_norm: Optional[Molmo2RMSNorm] = None + self.q_norm: Optional[Molmo2RMSNorm] = None + self.qk_norm_type: Optional[str] = None + if config.use_qk_norm: + k_norm_size = ( + config.head_dim + if config.qk_norm_type == "qwen3" else + config.num_key_value_heads * config.head_dim + ) + self.k_norm = Molmo2RMSNorm(k_norm_size, eps=config.layer_norm_eps) + q_norm_size = ( + config.head_dim + if config.qk_norm_type == "qwen3" else + config.num_attention_heads * config.head_dim + ) + self.q_norm = Molmo2RMSNorm(q_norm_size, eps=config.layer_norm_eps) + self.qk_norm_type = config.qk_norm_type + + self.attention_dropout = config.attention_dropout + + self.attn_out = nn.Linear( + config.head_dim * config.num_attention_heads, + config.hidden_size, + bias=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + qkv = self.att_proj(hidden_states) + query_states, key_states, value_states = qkv.split(self.fused_dims, dim=-1) + value_states = value_states.view(hidden_shape) + + # Optionally apply layer norm to keys and queries. + if self.q_norm is not None and self.k_norm is not None and self.qk_norm_type != "qwen3": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.view(hidden_shape) + key_states = key_states.view(hidden_shape) + if self.q_norm is not None and self.k_norm is not None and self.qk_norm_type == "qwen3": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.attn_out(attn_output) + return attn_output, attn_weights + + +class LanguageModelMLP(nn.Module): + + def __init__( + self, + input_dim: int, + intermediate_size: int, + hidden_act: str, + device: Union[str, torch.device] = None, + ): + super().__init__() + self.ff_proj = nn.Linear(input_dim, intermediate_size * 2, bias=False, device=device) + self.ff_out = nn.Linear(intermediate_size, input_dim, bias=False, device=device) + self.act = ACT2FN[hidden_act] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.ff_proj(x) + x, gate = x.chunk(2, dim=-1) + x = self.act(gate) * x + x = self.ff_out(x) + return x + + +class Molmo2DecoderLayer(GradientCheckpointingLayer): + + def __init__( + self, + config: Molmo2TextConfig, + layer_idx: Optional[int] = None, + device: Union[str, torch.device] = None + ): + super().__init__() + self.config = config + + self.self_attn = Molmo2Attention(config, layer_idx) + self.attn_norm = Molmo2RMSNorm( + config.hidden_size, eps=config.layer_norm_eps, device=device) + self.dropout = nn.Dropout(config.residual_dropout) + self.mlp = LanguageModelMLP( + config.hidden_size, config.intermediate_size, config.hidden_act, device=device) + self.ff_norm = Molmo2RMSNorm( + config.hidden_size, eps=config.layer_norm_eps, device=device) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: + + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = residual + self.dropout(hidden_states) + + # Fully Connected + residual = hidden_states + hidden_states = self.ff_norm(hidden_states) + hidden_states = self.mlp(hidden_states) + + hidden_states = residual + self.dropout(hidden_states) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +class Molmo2PostNormDecoderLayer(Molmo2DecoderLayer): + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: + + residual = hidden_states + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + ) + hidden_states = self.attn_norm(hidden_states) + + hidden_states = residual + self.dropout(hidden_states) + + # Fully Connected + residual = hidden_states + hidden_states = self.mlp(hidden_states) + hidden_states = self.ff_norm(hidden_states) + + hidden_states = residual + self.dropout(hidden_states) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +class Molmo2Embedding(nn.Module): + def __init__( + self, + num_embeddings: int, + num_new_embeddings: int, + features: int, + device: Union[str, torch.device] = None, + ): + super().__init__() + self.embedding = nn.Parameter( + torch.zeros(num_embeddings, features, device=device), + ) + self.new_embedding = nn.Parameter( + torch.zeros(num_new_embeddings, features, device=device), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.embedding(x, torch.cat([self.embedding, self.new_embedding], dim=0)) + + +class MolmoPointPreTrainedModel(PreTrainedModel): + config: MolmoPointConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = [ + "Molmo2DecoderLayer", + "Molmo2PostNormDecoderLayer", + "Molmo2VisionBlock", + "ViTMultiHeadDotProductAttention", + ] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn = True + _supports_sdpa = True + + _can_compile_fullgraph = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": Molmo2DecoderLayer, + "attentions": Molmo2Attention, + } + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, (nn.Linear,)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, Molmo2Embedding): + module.embedding.data.normal_(mean=0.0, std=std) + module.new_embedding.data.normal_(mean=0.0, std=std) + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, Molmo2RMSNorm): + module.weight.data.fill_(1.0) + elif isinstance(module, nn.LayerNorm): + module.weight.data.fill_(1.0) + if module.bias is not None: + module.bias.data.zero_() + + +class 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): + + 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 + + +class MolmoPointModel(MolmoPointPreTrainedModel): + base_model_prefix = "" + _checkpoint_conversion_mapping = {} + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + config: MolmoPointConfig + + def __init__(self, config: MolmoPointConfig): + super().__init__(config) + self.transformer: MolmoPointTextModel = MolmoPointTextModel(config.text_config) + + vit_config = config.vit_config + adapter_config = config.adapter_config + self.vit_layers = [] + for layer in adapter_config.vit_layers: + if layer >= 0: + self.vit_layers.append(layer) + else: + self.vit_layers.append(layer + vit_config.num_hidden_layers) + + last_layer_needed = max(self.vit_layers) + 1 + if last_layer_needed < vit_config.num_hidden_layers: + new_vit_config = deepcopy(vit_config) + new_vit_config.num_hidden_layers = last_layer_needed + self.vit = Molmo2VisionTransformer(new_vit_config) + else: + self.vit = Molmo2VisionTransformer(vit_config) + + 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.connector = MolmoPointConnector(adapter_config, vit_config) + + vit_dim = self.config.vit_config.hidden_size * len(self.config.adapter_config.vit_layers) + llm_dim = self.config.text_config.hidden_size + self.patch_q = nn.Linear(llm_dim, config.patch_embed_dim) + self.patch_k = nn.Linear(llm_dim, config.patch_embed_dim) + self.subpatch_q = nn.Linear(llm_dim, config.patch_embed_dim) + self.subpatch_k = nn.Linear(vit_dim, config.patch_embed_dim) + self.add_no_point_class_embed = MolmoPointPadWithLearnedVector(config.patch_embed_dim) + + 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 + + if self.config.embed_selected_vit_patch == "linear": + self.build_vit_embedding = nn.Linear(vit_dim, llm_dim, bias=True) + else: + raise NotImplementedError(f"Embedding {self.config.embed_selected_vit_patch} not implemented") + + if self.config.patch_location == "3x3": + self.subpatch_loc_k = nn.Linear(llm_dim, 9) + elif self.config.patch_location is None: + self.subpatch_loc_k = None + else: + raise NotImplementedError(f"Patch location {self.config.patch_location} not implemented") + + if self.config.layer_norm_x: + self.x_norm = Molmo2RMSNorm(llm_dim, eps=self.config.text_config.layer_norm_eps) + else: + self.x_norm = None + + # Initialize weights and apply final processing + self.post_init() + + def 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 generatd 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) + print(f"ON: {input_ids[0, -1]}") + + if image_data is not None: + 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: + input_patch_ids = None + can_point = False + + x = self.transformer.wte(input_ids) + batch_size, _, dim = x.shape + batch_idx = torch.arange(batch_size, device=self.device) + + # TODO update embeddings for patch/subpatch tokens + + vit_features_flat: Optional[torch.FloatTensor] = None + if images is not None: + is_image_token = input_ids == self.config.image_patch_id + 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) + 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) + x = x.clone() + x.view(-1, dim)[is_image_token.view(-1)] += image_features.view(-1, dim) + + # Build position ids for the image features, which we might need for rotary + # embeddings + image_token_indices = torch.cumsum(is_image_token, dim=-1) - 1 + if image_grids is not None: + # Global crop is always the first 196 images tokens and cannot be pointed to + is_indexable_image_token = is_image_token & (image_token_indices >= 196) + is_non_indexable_image_token = is_image_token & (image_token_indices < 196) + image_token_indices = torch.clip(image_token_indices - 196, min=0) + else: + is_indexable_image_token = is_image_token + is_non_indexable_image_token = torch.zeros_like(is_indexable_image_token) + image_pos_ids_flat = image_token_indices.view(-1)[is_image_token.view(-1)] + image_pos_ids = torch.zeros([batch_size, token_pooling.shape[1]], dtype=torch.long, device=self.device) + image_pos_ids.view(-1)[image_features_mask.view(-1)] = image_pos_ids_flat + max_image_pos_id = image_pos_ids_flat.max() + 1 + elif image_data is not None: + token_pooling = image_data.token_pooling + image_pos_ids = image_data.image_pos_ids + vit_features = image_data.vit_features + vit_features_mask = token_pooling >= 0 + 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)] + + # 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): + 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) + + # 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: + if self.x_norm: + 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 + else: + patch_k_flat = self.patch_k(x_norm.view(-1, dim)[is_image_token.view(-1)]) + if self.patch_rotary is not None: + patch_k_flat = self.patch_rotary(patch_k_flat, image_pos_ids_flat) + patch_k = torch.zeros([batch_size, image_features_mask.shape[1], patch_k_flat.shape[-1]], dtype=x.dtype, device=self.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) + + # Predict patch locations + if can_point: + 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 can_point and 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 can_point and torch.any(is_subpatch): + location_logits = self.subpatch_loc_k(x) + + if is_prefill: + 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]) + # Return the cache for image keys/features + 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, + image_features0=image_features, + 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 MolmoPointForConditionalGeneration(MolmoPointPreTrainedModel, GenerationMixin): + _checkpoint_conversion_mapping = {} + _tied_weights_keys = [] # Weights are not tied + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + config: MolmoPointConfig + + def __init__(self, config: MolmoPointConfig): + super().__init__(config) + + self.model = MolmoPointModel(config) + self.output_embeddings = nn.Parameter(torch.zeros([config.vocab_size, config.hidden_size])) + self.new_output_embeddings = nn.Parameter(torch.zeros([128, config.hidden_size])) + self.vocab_size = config.vocab_size + + # Initialize weights and apply final processing + self.post_init() + + def 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 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 + lm_head = torch.concatenate([self.output_embeddings, self.new_output_embeddings], dim=0) + logits = F.linear(hidden_states[:, slice_indices, :], lm_head) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.vocab_size) + + 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 + 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=logits.device) + if outputs.patch_logits is not None: + selected_patches = torch.argmax(outputs.patch_logits, -1) + bs, seq, n_patches = outputs.patch_logits.shape + batch_idx = torch.arange(outputs.patch_logits.shape[0], device=self.device) + seq_ix = torch.arange(outputs.patch_logits.shape[1], device=self.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=logits.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=logits.device) + + logits = torch.concatenate([ + logits, argmax_patch_logits, subpatch_logits, location_logits], -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) \ No newline at end of file