# coding=utf-8 # Copyright 2025 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import Callable, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache from transformers.masking_utils import create_causal_mask from dataclasses import dataclass from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, ) from transformers.processing_utils import Unpack from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from transformers.generation import GenerationMixin from transformers.generation.utils import GenerateDecoderOnlyOutput from transformers.utils import logging, TransformersKwargs from .configuration_moondream3 import ( Moondream3Config, Moondream3TextConfig, Moondream3VisionConfig, Moondream3RegionConfig, ) logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "Moondream3Config" def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, rot_dim: int = 32, ): """ Apply rotary position embeddings to query and key tensors. Args: q: Query tensor [batch, num_heads, seq_len, head_dim] k: Key tensor [batch, num_heads, seq_len, head_dim] cos: Cosine frequencies [batch, seq_len, rot_dim] sin: Sine frequencies [batch, seq_len, rot_dim] rot_dim: Number of dimensions to apply rotation to (default: 32) Returns: Tuple of (rotated_q, rotated_k) """ def apply_rope(x): dtype = x.dtype x = x.to(torch.float64) x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:] d_q = x_rot.shape[-1] // 2 xq_r, xq_i = x_rot[..., :d_q], x_rot[..., d_q:] xq_out_r = xq_r * cos - xq_i * sin xq_out_i = xq_r * sin + xq_i * cos xq_out = torch.stack((xq_out_r, xq_out_i), dim=-1).flatten(-2) return torch.cat([xq_out, x_pass], dim=-1) return apply_rope(q), apply_rope(k) class Moondream3RotaryEmbedding(nn.Module): inv_freq: torch.Tensor def __init__(self, config: Moondream3Config, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_type = self.config.rope_parameters["rope_type"] rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": 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 = inv_freq @staticmethod def compute_default_rope_parameters( config: Optional[Moondream3Config] = None, device: Optional["torch.device"] = None, seq_len: Optional[int] = None, ) -> tuple["torch.Tensor", float]: """ Computes the inverse frequencies according to the original RoPE implementation """ base = config.rope_parameters["rope_theta"] dim = ( getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads ) dim //= 2 attention_factor = 1.0 inv_freq = 1.0 / ( base ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim) ) if device is not None: inv_freq = inv_freq.to(device=device) return inv_freq, attention_factor @torch.no_grad() @dynamic_rope_update def forward(self, x, position_ids): inv_freq_expanded = ( self.inv_freq[None, :, None] .to(torch.float32) .expand(position_ids.shape[0], -1, 1) .to(x.device) ) position_ids_expanded = position_ids[:, None, :].to(torch.float32) freqs = ( inv_freq_expanded.to(torch.float32) @ position_ids_expanded.to(torch.float32) ).transpose(1, 2) cfreqs = ( torch.exp(1j * freqs) .unsqueeze(1) .expand(-1, self.config.num_attention_heads, -1, -1) ) return cfreqs.real, cfreqs.imag class Moondream3Attention(nn.Module): def __init__( self, config: Moondream3TextConfig | Moondream3VisionConfig, layer_idx: Optional[int] = None, use_tau: bool = True, ): super().__init__() self.config = config self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) self.num_key_value_heads = getattr( config, "num_key_value_heads", self.num_heads ) attention_bias = config.attention_bias self.attention_dropout = config.attention_dropout if isinstance(config, Moondream3TextConfig): self.is_causal = True elif isinstance(config, Moondream3VisionConfig): self.is_causal = False else: raise TypeError(f"Unsupported config type: {type(config)}") self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.use_tau = use_tau if (self.head_dim * self.num_heads) != self.hidden_size: raise ValueError( f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" f" and `num_heads`: {self.num_heads})." ) self.q_proj = nn.Linear( self.hidden_size, self.num_heads * self.head_dim, bias=attention_bias ) self.k_proj = nn.Linear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=attention_bias, ) self.v_proj = nn.Linear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=attention_bias, ) self.o_proj = nn.Linear( self.num_heads * self.head_dim, self.hidden_size, bias=attention_bias ) if self.use_tau: # In original, tau weights are (n_heads, qkv_dim) where qkv_dim is the combined QKV dimension qkv_dim = ( self.num_heads * self.head_dim + 2 * self.num_key_value_heads * self.head_dim ) self.tau_wq = nn.Linear(qkv_dim, self.num_heads, bias=False) self.tau_wv = nn.Linear(qkv_dim, self.num_heads, bias=False) self.tau_alpha = nn.Parameter(torch.empty(self.num_heads)) def forward( self, hidden_states: torch.Tensor, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: input_shape = hidden_states.shape[:-1] bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) if self.use_tau: qkv_out = torch.cat([query_states, key_states, value_states], dim=-1) tok_feat = F.gelu(qkv_out) tok_q = torch.tanh(self.tau_wq(tok_feat)).permute(0, 2, 1) tok_v = torch.tanh(self.tau_wv(tok_feat)).permute(0, 2, 1) pos = position_ids.to(tok_q.dtype) + 1 alpha = self.tau_alpha.to(tok_q.dtype) tau_pos = 1 + ( torch.sigmoid(alpha[None, :, None] * pos[:, None, :].log()) - 0.5 ) tau_q = (tok_q + tau_pos).unsqueeze(-1) tau_v = (tok_v + tau_pos).unsqueeze(-1) query_states = query_states.view( bsz, q_len, self.num_heads, self.head_dim ).transpose(1, 2) key_states = key_states.view( bsz, q_len, self.num_key_value_heads, self.head_dim ).transpose(1, 2) value_states = value_states.view( bsz, q_len, self.num_key_value_heads, self.head_dim ).transpose(1, 2) if self.use_tau: query_states = query_states * tau_q if self.num_key_value_groups > 1: tau_v_repeated = tau_v.repeat(1, self.num_key_value_groups, 1, 1)[ :, : self.num_key_value_heads, :, : ] else: tau_v_repeated = tau_v value_states = value_states * tau_v_repeated cos, sin = None, None if position_embeddings is not None: cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb( query_states, key_states, cos, sin ) query_states, key_states = ( query_states.to(value_states.dtype), key_states.to(value_states.dtype), ) if past_key_values is not None: 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 ) query_states = query_states.contiguous() key_states = key_states.contiguous() value_states = value_states.contiguous() attn_output, attn_weights = ALL_ATTENTION_FUNCTIONS["sdpa"]( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class Moondream3MLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str = "gelu_pytorch_tanh", out_size: int | None = None, gated: bool = False, bias: bool = True, ): super().__init__() self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.out_size = self.hidden_size if out_size is None else out_size self.hidden_act = hidden_act self.gated = gated self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) self.down_proj = nn.Linear(self.intermediate_size, self.out_size, bias=bias) self.gate_proj = None if self.gated: self.gate_proj = nn.Linear( self.hidden_size, self.intermediate_size, bias=bias ) self.act_fn = ACT2FN[self.hidden_act] def forward(self, x) -> torch.Tensor: if self.gated: h = self.up_proj(x) g = self.gate_proj(x) x = self.act_fn(h) * (g + 1) else: x = self.act_fn(self.up_proj(x)) return self.down_proj(x) class Moondream3SparseMoeBlock(nn.Module): def __init__(self, config: Moondream3TextConfig, layer_idx=None): super().__init__() self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.moe_intermediate_size = config.moe_intermediate_size self.num_experts = config.num_experts self.top_k = config.num_experts_per_tok self.gate = nn.Linear(self.hidden_size, self.num_experts, bias=True) self.experts = nn.ModuleList( [ Moondream3MLP( hidden_size=self.hidden_size, intermediate_size=self.moe_intermediate_size, gated=True, bias=False, hidden_act="gelu" ) for _ in range(self.num_experts) ] ) def forward( self, hidden_states: torch.Tensor, cache_position=None ) -> Tuple[torch.Tensor, torch.Tensor]: batch_size, sequence_length, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) router_logits: torch.Tensor = self.gate(hidden_states) routing_weights, selected_experts = torch.topk( router_logits, self.top_k, dim=-1 ) routing_weights = F.softmax(routing_weights, dim=-1, dtype=torch.float32) routing_weights = routing_weights.to(hidden_states.dtype) final_hidden_states = torch.zeros( (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device, ) for expert_idx in range(self.num_experts): expert_layer = self.experts[expert_idx] top_x, idx = (selected_experts == expert_idx).nonzero(as_tuple=True) if top_x.shape[0] == 0: continue current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) current_hidden_states = ( expert_layer(current_state) * routing_weights[top_x, idx, None] ) final_hidden_states.index_add_( 0, top_x, current_hidden_states.to(hidden_states.dtype) ) final_hidden_states = final_hidden_states.reshape( batch_size, sequence_length, hidden_dim ) return final_hidden_states, router_logits class Moondream3DecoderLayer(nn.Module): def __init__(self, config: Moondream3TextConfig, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.self_attn = Moondream3Attention(config, layer_idx, use_tau=True) self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps) self.is_moe_layer = layer_idx >= config.moe_start_layer if self.is_moe_layer: self.mlp = Moondream3SparseMoeBlock(config, layer_idx=layer_idx) else: self.mlp = Moondream3MLP( self.hidden_size, self.intermediate_size, ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, output_attentions: bool = False, output_router_logits: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, **kwargs, ) -> Tuple: hidden_states_ln = self.input_layernorm(hidden_states) hidden_states_attn, self_attn_weights = self.self_attn( hidden_states=hidden_states_ln, 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, position_embeddings=position_embeddings, **kwargs, ) if self.is_moe_layer: hidden_states_mlp, router_logits = self.mlp( hidden_states_ln, cache_position=cache_position ) else: hidden_states_mlp = self.mlp(hidden_states_ln) router_logits = None # Add both attention and MLP to residual like original hidden_states = hidden_states + hidden_states_attn + hidden_states_mlp outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) if output_router_logits: outputs += (router_logits,) return outputs class Moondream3PreTrainedModel(PreTrainedModel): config_class = Moondream3Config base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["Moondream3DecoderLayer", "Moondream3SparseMoeBlock"] _skip_keys_device_placement = "past_key_values" _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True def _init_weights(self, module): if hasattr(self.config, "text_config") and hasattr( self.config.text_config, "initializer_range" ): std = self.config.text_config.initializer_range elif hasattr(self.config, "initializer_range"): std = self.config.initializer_range else: std = 0.02 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, 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_() class Moondream3TextModel(Moondream3PreTrainedModel): config_class = Moondream3TextConfig def __init__(self, config: Moondream3TextConfig): super().__init__(config) self.padding_idx = config.pad_token_id if hasattr(config, "pad_token_id") else 0 self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) self.layers = nn.ModuleList( [ Moondream3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers) ] ) self.norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Moondream3RotaryEmbedding(config=config) self.gradient_checkpointing = False self.post_init() 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_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_router_logits = ( output_router_logits if output_router_logits is not None else self.config.output_router_logits ) 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 return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" ) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) hidden_states = inputs_embeds batch_size = hidden_states.shape[0] if self.gradient_checkpointing and self.training: if use_cache: logger.warning( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." ) use_cache = False if use_cache and past_key_values is None: past_key_values = DynamicCache() 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) position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None all_router_logits = () if output_router_logits else None for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, hidden_states, attention_mask, position_ids, past_key_values, output_attentions, output_router_logits, use_cache, cache_position, position_embeddings, ) else: layer_outputs = decoder_layer( hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, output_attentions=output_attentions, output_router_logits=output_router_logits, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, ) hidden_states = layer_outputs[0] if output_attentions: all_self_attns += (layer_outputs[1],) if output_router_logits and layer_outputs[-1] is not None: all_router_logits += (layer_outputs[-1],) if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = None if use_cache: next_cache = past_key_values if not return_dict: return tuple( v for v in [ hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits, ] if v is not None ) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, ) class Moondream3VisionPatchEmbeddings(nn.Module): def __init__(self, config: Moondream3VisionConfig): super().__init__() self.patch_size = config.patch_size self.num_channels = config.in_channels self.hidden_size = config.hidden_size self.crop_size = config.crop_size self.patch_size = config.patch_size self.grid_size = self.crop_size // self.patch_size self.num_patches = self.grid_size * self.grid_size self.projection = nn.Linear( self.patch_size * self.patch_size * self.num_channels, self.hidden_size, bias=True, ) self.position_embeddings = nn.Parameter( torch.zeros(1, self.num_patches, config.hidden_size) ) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: B, C, H, W = pixel_values.shape P1 = P2 = self.patch_size x = pixel_values.reshape(B, C, H // P1, P1, W // P2, P2) x = x.permute(0, 2, 4, 1, 3, 5) x = x.reshape(B, (H // P1) * (W // P2), C * P1 * P2) x = self.projection(x) return x + self.position_embeddings class Moondream3VisionEncoderLayer(nn.Module): def __init__(self, config: Moondream3VisionConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.layer_idx = layer_idx self.self_attn = Moondream3Attention( config, layer_idx=self.layer_idx, use_tau=False ) self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=1e-5) self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=1e-5) self.mlp = Moondream3MLP( hidden_size=self.hidden_size, intermediate_size=self.intermediate_size, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states, _ = self.self_attn(hidden_states=hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class Moondream3VisionModel(Moondream3PreTrainedModel): config_class = Moondream3VisionConfig main_input_name = "pixel_values" _no_split_modules = ["Moondream3VisionEncoderLayer"] def __init__(self, config: Moondream3VisionConfig): super().__init__(config) self.config = config self.hidden_size = self.config.hidden_size self.num_hidden_layers = self.config.num_hidden_layers self.proj_inner_dim = self.config.proj_inner_dim self.proj_out_dim = self.config.proj_out_dim self.embeddings = Moondream3VisionPatchEmbeddings(config) self.layers = nn.ModuleList( [ Moondream3VisionEncoderLayer(config, layer_idx) for layer_idx in range(self.num_hidden_layers) ] ) self.post_layernorm = nn.LayerNorm(self.hidden_size, eps=1e-5) self.vision_projection = Moondream3MLP( hidden_size=self.hidden_size * 2, intermediate_size=self.proj_inner_dim, out_size=self.proj_out_dim, ) self.gradient_checkpointing = False self.post_init() def _reconstruct_from_crops( self, crops: torch.Tensor, tiling: tuple[int, int], overlap_margin: int = 4, patch_size: int = 14, ) -> torch.Tensor: """ Reconstruct the original image from overlapping crops into a single seamless image. Takes a list of overlapping image crops along with their positional metadata and reconstructs them into a single coherent image by carefully stitching together non-overlapping regions. Handles both numpy arrays and PyTorch tensors. Args: crops: List of image crops as numpy arrays or PyTorch tensors with shape (H,W,C) tiling: Tuple of (height,width) indicating crop grid layout patch_size: Size in pixels of each patch, default 14 overlap_margin: Number of overlapping patches on each edge, default 4 Returns: Reconstructed image as numpy array or PyTorch tensor matching input type, with shape (H,W,C) where H,W are the original image dimensions """ if isinstance(tiling, torch.Tensor): tiling_h, tiling_w = tiling[0].item(), tiling[1].item() else: tiling_h, tiling_w = tiling tiling_h, tiling_w = int(tiling_h), int(tiling_w) crop_height, crop_width = crops[0].shape[:2] margin_pixels = overlap_margin * patch_size output_h = (crop_height - 2 * margin_pixels) * tiling_h + 2 * margin_pixels output_w = (crop_width - 2 * margin_pixels) * tiling_w + 2 * margin_pixels reconstructed = torch.zeros( (output_h, output_w, crops[0].shape[2]), device=crops[0].device, dtype=crops[0].dtype, ) for i, crop in enumerate(crops): tile_y = i // tiling_w tile_x = i % tiling_w x_start = 0 if tile_x == 0 else margin_pixels x_end = crop_width if tile_x == tiling_w - 1 else crop_width - margin_pixels y_start = 0 if tile_y == 0 else margin_pixels y_end = ( crop_height if tile_y == tiling_h - 1 else crop_height - margin_pixels ) out_x = tile_x * (crop_width - 2 * margin_pixels) out_y = tile_y * (crop_height - 2 * margin_pixels) reconstructed[ out_y + y_start : out_y + y_end, out_x + x_start : out_x + x_end ] = crop[y_start:y_end, x_start:x_end] return reconstructed def forward( self, pixel_values: torch.FloatTensor, tiling: Tuple[int, int], output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) batch_size, num_crops = pixel_values.shape[:2] # flatten batch_size and num_crops into same dim pixel_values = pixel_values.view(-1, *pixel_values.shape[2:]) hidden_states: torch.Tensor = self.embeddings(pixel_values) all_hidden_states = () if output_hidden_states else None all_attentions = () if output_attentions else None for encoder_layer in self.layers: if output_hidden_states and all_hidden_states is not None: all_hidden_states += (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( encoder_layer.__call__, hidden_states ) else: layer_outputs = encoder_layer(hidden_states) hidden_states = layer_outputs hidden_states = self.post_layernorm(hidden_states) hidden_states = hidden_states.view( batch_size, num_crops, *hidden_states.shape[1:] ) outputs = [] for b in range(batch_size): hs = hidden_states[b] t = tiling[b] global_features = hs[0] local_features = hs[1:].view( -1, self.num_hidden_layers, self.num_hidden_layers, self.hidden_size, ) reconstructed = self._reconstruct_from_crops( local_features, t, patch_size=1, overlap_margin=self.config.overlap_margin, ) reconstructed = reconstructed.permute(2, 0, 1) reconstructed = F.adaptive_avg_pool2d( reconstructed, output_size=(self.num_hidden_layers, self.num_hidden_layers), ) reconstructed = reconstructed.permute(1, 2, 0).view( self.num_hidden_layers * self.num_hidden_layers, self.hidden_size ) final_features = torch.cat([global_features, reconstructed], dim=-1) outputs.append(final_features) output = torch.stack(outputs, 0) hidden_states = self.vision_projection(output) if output_hidden_states and all_hidden_states is not None: all_hidden_states += (hidden_states,) if not return_dict: return tuple( v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None ) return BaseModelOutputWithPast( last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions, ) class Moondream3RegionEncoder(nn.Module): def __init__(self, config: Moondream3RegionConfig): super().__init__() self.coord_encoder = nn.Linear(config.coord_feat_dim, config.hidden_size) self.size_encoder = nn.Linear(config.size_feat_dim, config.hidden_size) coord_freq = torch.randn(config.coord_feat_dim // 2, 1) * 10.0 size_freq = torch.randn(config.size_feat_dim // 2, 2) * 10.0 self.register_buffer("coord_freq", coord_freq.T) self.register_buffer("size_freq", size_freq.T) def fourier_features(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: x_proj = 2 * torch.pi * x @ w return torch.cat([x_proj.cos(), x_proj.sin()], dim=-1) def encode_coordinate(self, coord: torch.Tensor) -> torch.Tensor: fourier_features = self.fourier_features(coord, self.coord_freq) return self.coord_encoder(fourier_features) def encode_size(self, size: torch.Tensor) -> torch.Tensor: fourier_features = self.fourier_features(size, self.size_freq) return self.size_encoder(fourier_features) class Moondream3RegionDecoder(nn.Module): def __init__(self, config: Moondream3RegionConfig): super().__init__() self.coord_decoder = nn.Linear(config.hidden_size, config.coord_out_dim) self.size_decoder = nn.Linear(config.hidden_size, config.size_out_dim) def decode_coordinate(self, hidden_state: torch.Tensor) -> torch.Tensor: return self.coord_decoder(hidden_state) def decode_size(self, hidden_state: torch.Tensor) -> torch.Tensor: return self.size_decoder(hidden_state).view(hidden_state.shape[0], 2, -1) class Moondream3Model(Moondream3PreTrainedModel): def __init__(self, config: Moondream3Config): super().__init__(config) self.text_model = Moondream3TextModel(config.text_config) self.vision_model = Moondream3VisionModel(config.vision_config) self.vocab_size = config.text_config.vocab_size self.region_encoder = Moondream3RegionEncoder(config.region_config) self.region_decoder = Moondream3RegionDecoder(config.region_config) self.post_init() def get_input_embeddings(self): return self.text_model.embed_tokens def set_input_embeddings(self, value): self.text_model.embed_tokens = value def set_decoder(self, decoder): self.text_model = decoder def get_decoder(self): return self.text_model def forward( self, input_ids: torch.LongTensor = None, pixel_values: torch.FloatTensor = None, tiling: Tuple[int, int] = 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, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: int = 0, ) -> Union[Tuple, BaseModelOutputWithPast]: output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) if (input_ids is not None) == (inputs_embeds is not None): raise ValueError("Provide exactly one of input_ids or inputs_embeds.") if not ((pixel_values is not None) ^ (tiling is None)): raise ValueError("You must specify both pixel_values and tiling") if inputs_embeds is not None and ( pixel_values is not None or tiling is not None ): raise ValueError( "When inputs_embeds is provided, do not pass pixel_values/tiling; " "inputs_embeds must already include BOS+image(+text)." ) if inputs_embeds is None: inputs_embeds: torch.Tensor = self.text_model.embed_tokens(input_ids) if use_cache and past_key_values is None: 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.Tensor = torch.arange( past_seen_tokens, past_seen_tokens, device=inputs_embeds.device ) if position_ids is None: position_ids = cache_position.unsqueeze(0) if pixel_values is not None: pixel_values = pixel_values.to( dtype=self.vision_model.embeddings.projection.weight.dtype ) image_embeds = self.vision_model(pixel_values, tiling=tiling)[ "last_hidden_state" ] prefix = self.text_model.embed_tokens( torch.full( (input_ids.shape[0], 1), 0, dtype=input_ids.dtype, device=input_ids.device, ) ) embeds = torch.cat([prefix, image_embeds], dim=1) cache_pos = torch.arange(embeds.shape[-2], device=embeds.device) pos = cache_pos.unsqueeze(0).expand(embeds.shape[0], -1) attn_mask = torch.full( (embeds.shape[0], 1, embeds.shape[-2], pos.shape[-1]), True, dtype=torch.bool, device=embeds.device, ) outputs = self.text_model( input_ids=None, attention_mask=attn_mask, position_ids=pos, past_key_values=past_key_values, inputs_embeds=embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_pos, ) attn_mask = create_causal_mask( config=self.config, input_embeds=inputs_embeds, attention_mask=torch.cat( [ torch.ones( attention_mask.shape[0], cache_position[-1] + 1 - attention_mask.shape[-1], device=attention_mask.device, dtype=attention_mask.dtype, ), attention_mask, ], dim=-1, ), cache_position=cache_position, past_key_values=past_key_values, position_ids=position_ids, ) outputs = self.text_model( input_ids=None, attention_mask=attn_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, ) if not return_dict: return tuple( v for v in [ outputs.last_hidden_state, getattr(outputs, "past_key_values", None), getattr(outputs, "hidden_states", None), getattr(outputs, "attentions", None), ] if v is not None ) return BaseModelOutputWithPast( last_hidden_state=outputs.last_hidden_state, past_key_values=getattr(outputs, "past_key_values", None), hidden_states=getattr(outputs, "hidden_states", None), attentions=getattr(outputs, "attentions", None), ) @dataclass class Moondream3GenerateOutput(GenerateDecoderOnlyOutput): objects: Optional[list[dict[str, float]]] = None class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] def __init__(self, config: Moondream3Config): super().__init__(config) self.objects = None self.model = Moondream3Model(config) self.vocab_size = config.text_config.vocab_size self.lm_head = nn.Linear( config.text_config.hidden_size, config.text_config.vocab_size, bias=True ) self.post_init() def get_input_embeddings(self): return self.model.text_model.embed_tokens def set_input_embeddings(self, value): self.model.text_model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def set_decoder(self, decoder): self.model.text_model = decoder def get_decoder(self): return self.model.text_model def _prepare_generated_length( self, generation_config, **kwargs, ): generation_config = super()._prepare_generated_length( generation_config, **kwargs ) generation_config.max_length += self.config.vision_config.prefix_len return generation_config def forward( self, input_ids: torch.LongTensor = None, pixel_values: torch.FloatTensor = None, tiling: 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, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: int = 0, **kwargs: Unpack[TransformersKwargs], ) -> Union[Tuple, CausalLMOutputWithPast]: if pixel_values is not None and inputs_embeds is None: position_ids += self.config.vision_config.prefix_len cache_position += self.config.vision_config.prefix_len model_outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, tiling=tiling, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=None, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, logits_to_keep=logits_to_keep, ) hidden_states = model_outputs.last_hidden_state if isinstance(logits_to_keep, int) and logits_to_keep > 0: hs = hidden_states[:, -logits_to_keep:, :] elif isinstance(logits_to_keep, slice): hs = hidden_states[:, logits_to_keep, :] else: hs = hidden_states hs = self.model.text_model.norm(hs) logits = self.lm_head(hs) pred = torch.argmax(logits, dim=-1) pos_ids = position_ids[:, -1:] + 1 cache_pos = cache_position[-1:] + 1 mask = torch.ones( hidden_states.shape[0], 1, device=self.device, dtype=torch.long ) is_processing_point = torch.any(pred == 5) while is_processing_point: batch_mask = pred[:, -1] == 5 hidden_states = hidden_states[:, -1:, :] x_logits = self.model.region_decoder.decode_coordinate(hidden_states) x_center = torch.argmax(x_logits, dim=-1) / x_logits.size(-1) next_embeds = self.model.region_encoder.encode_coordinate( x_center.to(x_logits.dtype) ).unsqueeze(1) model_outputs = self.model( input_ids=None, pixel_values=None, tiling=None, attention_mask=mask, position_ids=pos_ids, past_key_values=past_key_values, inputs_embeds=next_embeds, labels=None, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_pos, logits_to_keep=logits_to_keep, ) hidden_states = model_outputs.last_hidden_state y_logits = self.model.region_decoder.decode_coordinate(hidden_states) y_center = torch.argmax(y_logits, dim=-1) / y_logits.size(-1) next_embeds = self.model.region_encoder.encode_coordinate( y_center.to(y_logits.dtype) ).unsqueeze(1) coords = torch.cat([x_center, y_center], dim=1) coords = coords * (batch_mask).unsqueeze(1) pos_ids += 1 cache_pos = cache_pos + 1 bbox = None if input_ids.shape[-1] > 1 and input_ids[0, 1] == 7235: model_outputs = self.model( input_ids=None, pixel_values=None, tiling=None, attention_mask=mask, position_ids=pos_ids, past_key_values=past_key_values, inputs_embeds=next_embeds, labels=None, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_pos, logits_to_keep=logits_to_keep, ) hidden_states = model_outputs.last_hidden_state size_logits = self.model.region_decoder.decode_size(hidden_states) bins = torch.argmax(size_logits, dim=-1) w_bin = bins[:, 0] h_bin = bins[:, 1] w = torch.pow(2.0, (w_bin.float() / 1023.0) * 10.0 - 10.0) h = torch.pow(2.0, (h_bin.float() / 1023.0) * 10.0 - 10.0) next_embeds = ( self.model.region_encoder.encode_size( torch.stack([w, h], dim=-1).to(size_logits.dtype) ) ).unsqueeze(1) x_center = x_center.squeeze(1) y_center = y_center.squeeze(1) bbox = [ x_center - w / 2, y_center - h / 2, x_center + w / 2, y_center + h / 2, ] bbox = torch.stack(bbox, dim=1) # shape (B, 4) bbox = bbox * (batch_mask).unsqueeze(1) pos_ids += 1 cache_pos = cache_pos + 1 new = coords.unsqueeze(1) if bbox is None else bbox.unsqueeze(1) if self.objects is None: self.objects = new else: self.objects = torch.cat([self.objects, new], dim=1) model_outputs = self.model( input_ids=None, pixel_values=None, tiling=None, attention_mask=mask, position_ids=pos_ids, past_key_values=past_key_values, inputs_embeds=next_embeds, labels=None, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_pos, logits_to_keep=logits_to_keep, ) pos_ids += 1 cache_pos = cache_pos + 1 hidden_states = model_outputs.last_hidden_state indices = torch.tensor( [ self.config.text_config.coord_token_id, 0, ], device=self.device, ) hidden_states = self.model.text_model.norm(hidden_states) logits = ( hidden_states @ self.lm_head.weight[indices].T + self.lm_head.bias[indices] ) logits_full = torch.full( (logits.shape[0], logits.shape[1], self.config.text_config.vocab_size), float("-inf"), device=logits.device, dtype=logits.dtype, ) logits_full[:, :, torch.tensor([5, 0])] = logits logits = logits_full pred[batch_mask] = torch.argmax(logits, dim=-1)[batch_mask] is_processing_point = torch.any(pred == 5) loss = None if labels is not None: loss = self.loss_function( logits=logits, labels=labels, vocab_size=self.vocab_size ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=getattr(model_outputs, "past_key_values", None), hidden_states=getattr(model_outputs, "hidden_states", None), attentions=getattr(model_outputs, "attentions", None), ) def generate(self, **kwargs) -> Union[Moondream3GenerateOutput, torch.LongTensor]: outputs = super().generate(**kwargs) if len(self.objects if self.objects is not None else []) > 0: if isinstance(outputs, torch.Tensor): outputs = self.objects self.objects = None else: outputs = Moondream3GenerateOutput(**outputs, objects=self.objects) self.objects = None return outputs def prepare_inputs_for_generation(self, input_ids, **model_kwargs): model_inputs = super().prepare_inputs_for_generation(input_ids, **model_kwargs) model_inputs["position_ids"] += ( model_inputs["cache_position"].unsqueeze(0) - model_inputs["position_ids"] ) return model_inputs def _update_model_kwargs_for_generation( self, outputs, model_kwargs, is_encoder_decoder, num_new_tokens: int = 1, ): model_kwargs = super()._update_model_kwargs_for_generation( outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder, num_new_tokens=num_new_tokens, ) model_kwargs["pixel_values"] = None model_kwargs["tiling"] = None return model_kwargs @staticmethod def _reorder_cache(past_key_values, beam_idx): reordered_past = () for layer_past in past_key_values: reordered_past += ( tuple( past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past ), ) return reordered_past __all__ = [ "Moondream3Config", "Moondream3TextConfig", "Moondream3VisionConfig", "Moondream3RegionConfig", "Moondream3PreTrainedModel", "Moondream3Model", "Moondream3TextModel", "Moondream3VisionModel", "Moondream3ForConditionalGeneration", ]