Update modeling_moondream3.py

modeling_moondream3.py

@@ -15,14 +15,10 @@
 
 from typing import Callable, Optional, Tuple, Union
 
-import numpy as np
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from PIL import Image
-
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.masking_utils import create_causal_mask
@@ -37,33 +33,35 @@ 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
+from .configuration_moondream3 import (
+    Moondream3Config,
+    Moondream3TextConfig,
+    Moondream3VisionConfig,
+    Moondream3RegionConfig,
+)
 
 logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "Moondream3Config"
 
-import torch
-
-DEBUG=True
 
 def apply_rotary_pos_emb(
-    q: torch.Tensor,  # [B, H, L, D]
-    k: torch.Tensor,  # [B, H, L, D]
-    cos: torch.Tensor,  # [B, L, rot_dim]
-    sin: torch.Tensor,  # [B, L, rot_dim]
+    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)
     """
@@ -77,13 +75,15 @@ def apply_rotary_pos_emb(
         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_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
 
@@ -112,15 +112,17 @@ class Moondream3RotaryEmbedding(nn.Module):
         """
         Computes the inverse frequencies according to the original RoPE implementation
         """
-        base = config.rope_parameters["rope_theta"]  # Should be 1500000.0 to match original
-        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        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
 
-        # Compute the inverse frequencies - matches your original formula
         inv_freq = 1.0 / (
-            base ** (torch.arange(0, dim, 2, dtype=torch.float64)[: (dim // 2)] / dim)
+            base ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim)
         )
         if device is not None:
             inv_freq = inv_freq.to(device=device)
@@ -129,38 +131,53 @@ class Moondream3RotaryEmbedding(nn.Module):
     @torch.no_grad()
     @dynamic_rope_update
     def forward(self, x, position_ids):
-        # inv_freq shape: [dim//2]
-        # position_ids shape: [batch_size, seq_len]
-        
-        inv_freq_expanded = self.inv_freq[None, :, None].to(torch.float64).expand(position_ids.shape[0], -1, 1).to(x.device)
-        position_ids_expanded = position_ids[:, None, :].to(torch.float64)
-
-        freqs = (inv_freq_expanded.to(torch.float64) @ position_ids_expanded.to(torch.float64)).transpose(1, 2)
-        cfreqs = torch.exp(1j * freqs).unsqueeze(1).expand(-1, self.config.num_attention_heads, -1, -1)
+        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):
+    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)
+        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
 
-        # Initialize parameters based on config type
         if isinstance(config, Moondream3TextConfig):
             self.is_causal = True
-        elif isinstance(config, Moondream3VisionConfig):  # Moondream3VisionConfig
+        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
 
@@ -170,15 +187,29 @@ class Moondream3Attention(nn.Module):
                 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)
-        
-        # Tau parameters for token-level attention (from original Moondream) - only for text model
+        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
+            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))
@@ -196,65 +227,65 @@ class Moondream3Attention(nn.Module):
         **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
         input_shape = hidden_states.shape[:-1]
-        if isinstance(self.config, Moondream3TextConfig) and DEBUG:
-            torch.save(hidden_states, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_input_states")
+
         bsz, q_len, _ = hidden_states.size()
 
-        # Get qkv combined for tau (before splitting)
         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)  # (bsz, n_heads, seq_len)
-            tok_v = torch.tanh(self.tau_wv(tok_feat)).permute(0, 2, 1)  # (bsz, n_heads, seq_len)
+            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)  # (n_heads, seq_len)
-            tau_q = (tok_q + tau_pos).unsqueeze(-1)  # (bsz, n_heads, seq_len, 1)
-            tau_v = (tok_v + tau_pos).unsqueeze(-1)  # (bsz, n_heads, seq_len, 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 isinstance(self.config, Moondream3TextConfig) and DEBUG:
-            torch.save(value_states, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_pre_tau_value")
+            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, :, :]
+                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
 
-        if isinstance(self.config, Moondream3TextConfig) and DEBUG:
-            torch.save(value_states, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_post_tau_value")
-            torch.save(key_states, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_pre_rope_key")
-
         cos, sin = None, None
         if position_embeddings is not None:
             cos, sin = position_embeddings
-            if isinstance(self.config, Moondream3TextConfig) and DEBUG:
-                torch.save(cos, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_cos")
-                torch.save(sin, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_sin")
-            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
-        if isinstance(self.config, Moondream3TextConfig) and DEBUG:
-            torch.save(key_states, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_post_rope_key")
-        query_states, key_states = query_states.to(value_states.dtype), key_states.to(value_states.dtype)
+            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)
-
-        if isinstance(self.config, Moondream3TextConfig) and DEBUG:
-            torch.save(key_states, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_post_cache_key")
-            torch.save(attention_mask, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_attn_mask")
+            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()
@@ -272,13 +303,19 @@ class Moondream3Attention(nn.Module):
         attn_output = attn_output.reshape(*input_shape, -1).contiguous()
         attn_output = self.o_proj(attn_output)
 
-        if isinstance(self.config, Moondream3TextConfig) and DEBUG:
-            torch.save(attn_output, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_attn_out")
-
         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):
+    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
@@ -289,15 +326,15 @@ class Moondream3MLP(nn.Module):
         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.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:
-            # separate up and gate causes precision issues
-            combined_weight = torch.cat([self.up_proj.weight, self.gate_proj.weight], dim=0)
-            h_full = F.linear(x, combined_weight)
-            h, g = h_full.chunk(2, dim=-1)
+            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))
@@ -305,7 +342,7 @@ class Moondream3MLP(nn.Module):
 
 
 class Moondream3SparseMoeBlock(nn.Module):
-    def __init__(self, config: Moondream3TextConfig, layer_idx = None):
+    def __init__(self, config: Moondream3TextConfig, layer_idx=None):
         super().__init__()
         self.layer_idx = layer_idx
         self.hidden_size = config.hidden_size
@@ -314,18 +351,35 @@ class Moondream3SparseMoeBlock(nn.Module):
         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, hidden_act="gelu", gated=True, bias=False) for _ in range(self.num_experts)])
+        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]:
+    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, 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
+            (batch_size * sequence_length, hidden_dim),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
         )
 
         for expert_idx in range(self.num_experts):
@@ -336,12 +390,16 @@ class Moondream3SparseMoeBlock(nn.Module):
                 continue
 
             current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
-            # torch.save(current_state, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_e{expert_idx}")
-            current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
-            # torch.save(current_hidden_states, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_e{expert_idx}")
-            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+            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)
+        final_hidden_states = final_hidden_states.reshape(
+            batch_size, sequence_length, hidden_dim
+        )
         return final_hidden_states, router_logits
 
 
@@ -358,7 +416,10 @@ class Moondream3DecoderLayer(nn.Module):
         if self.is_moe_layer:
             self.mlp = Moondream3SparseMoeBlock(config, layer_idx=layer_idx)
         else:
-            self.mlp = Moondream3MLP(self.hidden_size, self.intermediate_size)
+            self.mlp = Moondream3MLP(
+                self.hidden_size,
+                self.intermediate_size,
+            )
 
     def forward(
         self,
@@ -373,16 +434,10 @@ class Moondream3DecoderLayer(nn.Module):
         position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
         **kwargs,
     ) -> Tuple:
-        residual = hidden_states
+        hidden_states_ln = self.input_layernorm(hidden_states)
 
-        # Apply layer norm like original
-        l_in = self.input_layernorm(hidden_states)
-        if DEBUG:
-            torch.save(l_in, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_ln_out")
-
-        # Attention
         hidden_states_attn, self_attn_weights = self.self_attn(
-            hidden_states=l_in,
+            hidden_states=hidden_states_ln,
             attention_mask=attention_mask,
             position_ids=position_ids,
             past_key_values=past_key_values,
@@ -390,20 +445,19 @@ class Moondream3DecoderLayer(nn.Module):
             use_cache=use_cache,
             cache_position=cache_position,
             position_embeddings=position_embeddings,
-            **kwargs
+            **kwargs,
         )
 
-        # MLP
         if self.is_moe_layer:
-            hidden_states_mlp, router_logits = self.mlp(l_in, cache_position=cache_position)
+            hidden_states_mlp, router_logits = self.mlp(
+                hidden_states_ln, cache_position=cache_position
+            )
         else:
-            hidden_states_mlp = self.mlp(l_in)
+            hidden_states_mlp = self.mlp(hidden_states_ln)
             router_logits = None
-        if DEBUG:
-            torch.save(hidden_states_mlp, f"dbg/hf_l{self.layer_idx}_c{cache_position[-1].item()}_mlp_out")
 
         # Add both attention and MLP to residual like original
-        hidden_states = residual + hidden_states_attn + hidden_states_mlp
+        hidden_states = hidden_states + hidden_states_attn + hidden_states_mlp
 
         outputs = (hidden_states,)
 
@@ -427,14 +481,15 @@ class Moondream3PreTrainedModel(PreTrainedModel):
     _supports_cache_class = True
 
     def _init_weights(self, module):
-        # Use text_config initializer_range if available, otherwise use default
-        if hasattr(self.config, 'text_config') and hasattr(self.config.text_config, 'initializer_range'):
+        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'):
+        elif hasattr(self.config, "initializer_range"):
             std = self.config.initializer_range
         else:
-            std = 0.02  # Default initialization range
-            
+            std = 0.02
+
         if isinstance(module, nn.Linear):
             module.weight.data.normal_(mean=0.0, std=std)
             if module.bias is not None:
@@ -453,15 +508,19 @@ class Moondream3TextModel(Moondream3PreTrainedModel):
         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.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)]
+            [
+                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(
@@ -478,19 +537,31 @@ class Moondream3TextModel(Moondream3PreTrainedModel):
         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_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_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
+            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
+        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")
+            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)
@@ -509,9 +580,13 @@ class Moondream3TextModel(Moondream3PreTrainedModel):
             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
+            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
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
             )
 
         if position_ids is None:
@@ -538,7 +613,7 @@ class Moondream3TextModel(Moondream3PreTrainedModel):
                     output_router_logits,
                     use_cache,
                     cache_position,
-                    position_embeddings
+                    position_embeddings,
                 )
             else:
                 layer_outputs = decoder_layer(
@@ -550,7 +625,7 @@ class Moondream3TextModel(Moondream3PreTrainedModel):
                     output_router_logits=output_router_logits,
                     use_cache=use_cache,
                     cache_position=cache_position,
-                    position_embeddings=position_embeddings
+                    position_embeddings=position_embeddings,
                 )
 
             hidden_states = layer_outputs[0]
@@ -571,7 +646,13 @@ class Moondream3TextModel(Moondream3PreTrainedModel):
         if not return_dict:
             return tuple(
                 v
-                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
+                for v in [
+                    hidden_states,
+                    next_cache,
+                    all_hidden_states,
+                    all_self_attns,
+                    all_router_logits,
+                ]
                 if v is not None
             )
 
@@ -594,8 +675,14 @@ class Moondream3VisionPatchEmbeddings(nn.Module):
         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))
+        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
@@ -610,17 +697,23 @@ class Moondream3VisionPatchEmbeddings(nn.Module):
         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.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)
+        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
@@ -635,6 +728,7 @@ class Moondream3VisionEncoderLayer(nn.Module):
 
         return hidden_states
 
+
 class Moondream3VisionModel(Moondream3PreTrainedModel):
     config_class = Moondream3VisionConfig
     main_input_name = "pixel_values"
@@ -649,9 +743,18 @@ class Moondream3VisionModel(Moondream3PreTrainedModel):
         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.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.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()
 
@@ -688,7 +791,6 @@ class Moondream3VisionModel(Moondream3PreTrainedModel):
         crop_height, crop_width = crops[0].shape[:2]
         margin_pixels = overlap_margin * patch_size
 
-        # Calculate output size (only adding margins once)
         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(
@@ -701,21 +803,16 @@ class Moondream3VisionModel(Moondream3PreTrainedModel):
             tile_y = i // tiling_w
             tile_x = i % tiling_w
 
-            # For each tile, determine which part to keep
-            # Keep left margin only for first column
             x_start = 0 if tile_x == 0 else margin_pixels
-            # Keep right margin only for last column
             x_end = crop_width if tile_x == tiling_w - 1 else crop_width - margin_pixels
-            # Keep top margin only for first row
             y_start = 0 if tile_y == 0 else margin_pixels
-            # Keep bottom margin only for last row
-            y_end = crop_height if tile_y == tiling_h - 1 else crop_height - margin_pixels
+            y_end = (
+                crop_height if tile_y == tiling_h - 1 else crop_height - margin_pixels
+            )
 
-            # Calculate where this piece belongs in the output
             out_x = tile_x * (crop_width - 2 * margin_pixels)
             out_y = tile_y * (crop_height - 2 * margin_pixels)
 
-            # Place the piece
             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]
@@ -725,16 +822,24 @@ class Moondream3VisionModel(Moondream3PreTrainedModel):
     def forward(
         self,
         pixel_values: torch.FloatTensor,
-        tiling: Tuple[int,int],
+        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_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
+            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
         )
-        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
@@ -749,17 +854,19 @@ class Moondream3VisionModel(Moondream3PreTrainedModel):
                 all_hidden_states += (hidden_states,)
 
             if self.gradient_checkpointing and self.training:
-                layer_outputs = self._gradient_checkpointing_func(encoder_layer.__call__, hidden_states)
+                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)
-        # B, _, _
 
-        # back out into batch_size, num_crops
-        hidden_states = hidden_states.view(batch_size, num_crops, *hidden_states.shape[1:])
+        hidden_states = hidden_states.view(
+            batch_size, num_crops, *hidden_states.shape[1:]
+        )
         outputs = []
         for b in range(batch_size):
             hs = hidden_states[b]
@@ -782,9 +889,12 @@ class Moondream3VisionModel(Moondream3PreTrainedModel):
 
             reconstructed = reconstructed.permute(2, 0, 1)
             reconstructed = F.adaptive_avg_pool2d(
-                reconstructed, output_size=(self.num_hidden_layers, self.num_hidden_layers)
+                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
             )
-            reconstructed = reconstructed.permute(1, 2, 0).view(729, self.hidden_size)
             final_features = torch.cat([global_features, reconstructed], dim=-1)
             outputs.append(final_features)
         output = torch.stack(outputs, 0)
@@ -795,7 +905,11 @@ class Moondream3VisionModel(Moondream3PreTrainedModel):
             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 tuple(
+                v
+                for v in [hidden_states, all_hidden_states, all_attentions]
+                if v is not None
+            )
 
         return BaseModelOutputWithPast(
             last_hidden_state=hidden_states,
@@ -803,12 +917,13 @@ class Moondream3VisionModel(Moondream3PreTrainedModel):
             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)
@@ -826,6 +941,7 @@ class Moondream3RegionEncoder(nn.Module):
         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__()
@@ -836,7 +952,8 @@ class Moondream3RegionDecoder(nn.Module):
         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)
+        return self.size_decoder(hidden_state).view(hidden_state.shape[0], 2, -1)
+
 
 class Moondream3Model(Moondream3PreTrainedModel):
     def __init__(self, config: Moondream3Config):
@@ -844,7 +961,7 @@ class Moondream3Model(Moondream3PreTrainedModel):
         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()
@@ -865,7 +982,7 @@ class Moondream3Model(Moondream3PreTrainedModel):
         self,
         input_ids: torch.LongTensor = None,
         pixel_values: torch.FloatTensor = None,
-        tiling: Tuple[int,int] = None,
+        tiling: Tuple[int, int] = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
         past_key_values: Optional[Cache] = None,
@@ -878,11 +995,19 @@ class Moondream3Model(Moondream3PreTrainedModel):
         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_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
+            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
         )
-        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.")
@@ -890,8 +1015,9 @@ class Moondream3Model(Moondream3PreTrainedModel):
         if not ((pixel_values is not None) ^ (tiling is None)):
             raise ValueError("You must specify both pixel_values and tiling")
 
-        # Case A: inputs_embeds provided -> assume it already contains BOS+image+text in correct order.
-        if inputs_embeds is not None and (pixel_values is not None or tiling is not None):
+        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)."
@@ -904,7 +1030,9 @@ class Moondream3Model(Moondream3PreTrainedModel):
             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
+            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
             )
@@ -912,20 +1040,24 @@ class Moondream3Model(Moondream3PreTrainedModel):
         if position_ids is None:
             position_ids = cache_position.unsqueeze(0)
 
-        def image_mask_function(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int):
-            # set all up to `self.config.vision_config.prefix_len` to true
-            return kv_idx <= q_idx
-
         if pixel_values is not None:
-            # Vision embeds
-            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"]  # [B,P,D]
+            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), self.config.text_config.bos_token_id, dtype=input_ids.dtype, device=input_ids.device)
+                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)
+            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,
@@ -949,11 +1081,21 @@ class Moondream3Model(Moondream3PreTrainedModel):
         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),
+            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,
-            and_mask_function=image_mask_function
         )
 
         outputs = self.text_model(
@@ -970,12 +1112,16 @@ class Moondream3Model(Moondream3PreTrainedModel):
         )
 
         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 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,
@@ -984,9 +1130,10 @@ class Moondream3Model(Moondream3PreTrainedModel):
             attentions=getattr(outputs, "attentions", None),
         )
 
+
 @dataclass
 class Moondream3GenerateOutput(GenerateDecoderOnlyOutput):
-    objects: Optional[list[dict[str,float]]] = None
+    objects: Optional[list[dict[str, float]]] = None
 
 
 class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMixin):
@@ -997,7 +1144,9 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
         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.lm_head = nn.Linear(
+            config.text_config.hidden_size, config.text_config.vocab_size, bias=True
+        )
         self.post_init()
 
     def get_input_embeddings(self):
@@ -1023,8 +1172,10 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
         generation_config,
         **kwargs,
     ):
-        generation_config = super()._prepare_generated_length(generation_config, **kwargs)
-        generation_config.max_length += 730
+        generation_config = super()._prepare_generated_length(
+            generation_config, **kwargs
+        )
+        generation_config.max_length += self.config.vision_config.prefix_len
         return generation_config
 
     def forward(
@@ -1048,7 +1199,7 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
         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
-        # Get hidden states from the base model (it already builds the multimodal prefix)
+        
         model_outputs = self.model(
             input_ids=input_ids,
             pixel_values=pixel_values,
@@ -1065,9 +1216,7 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
             cache_position=cache_position,
             logits_to_keep=logits_to_keep,
         )
-        hidden_states = model_outputs.last_hidden_state  # [B, T, D]
-
-        # Compute logits; only keep the tail if requested
+        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):
@@ -1076,21 +1225,24 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
             hs = hidden_states
 
         hs = self.model.text_model.norm(hs)
-        logits = self.lm_head(hs)  # [B, T', V]
+        logits = self.lm_head(hs)
 
         pred = torch.argmax(logits, dim=-1)
 
-        pos_ids = position_ids[:,-1:] + 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
         )
-        while torch.any(pred == 5):
-            batch_mask = (pred[:, -1] == 5)
+        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)
+            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,
@@ -1107,16 +1259,18 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
                 cache_position=cache_pos,
                 logits_to_keep=logits_to_keep,
             )
-            hidden_states = model_outputs.last_hidden_state  # [B, T, D]
+            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)
+            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[0,1] == 7235:
+            if input_ids.shape[-1] > 1 and input_ids[0, 1] == 7235:
                 model_outputs = self.model(
                     input_ids=None,
                     pixel_values=None,
@@ -1133,18 +1287,18 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
                     cache_position=cache_pos,
                     logits_to_keep=logits_to_keep,
                 )
-                hidden_states = model_outputs.last_hidden_state  # [B, T, D]
+                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_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)
+                        torch.stack([w, h], dim=-1).to(size_logits.dtype)
                     )
                 ).unsqueeze(1)
                 x_center = x_center.squeeze(1)
@@ -1155,7 +1309,7 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
                     x_center + w / 2,
                     y_center + h / 2,
                 ]
-                bbox = torch.stack(bbox, dim=1)
+                bbox = torch.stack(bbox, dim=1)  # shape (B, 4)
                 bbox = bbox * (batch_mask).unsqueeze(1)
                 pos_ids += 1
                 cache_pos = cache_pos + 1
@@ -1183,26 +1337,38 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
             )
             pos_ids += 1
             cache_pos = cache_pos + 1
-            hidden_states = model_outputs.last_hidden_state  # [B, T, D]
+            hidden_states = model_outputs.last_hidden_state
 
             indices = torch.tensor(
-                [self.config.text_config.coord_token_id, self.config.text_config.eos_token_id],
+                [
+                    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 = (
+                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_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:
-            # Shift if your training uses standard LM convention; here we assume labels aligned with hs
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.vocab_size)
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.vocab_size
+            )
 
         return CausalLMOutputWithPast(
             loss=loss,
@@ -1214,25 +1380,20 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
 
     def generate(self, **kwargs) -> Union[Moondream3GenerateOutput, torch.LongTensor]:
         outputs = super().generate(**kwargs)
-        if len(self.objects) > 0:
+        if len(self.objects if self.objects is not None else []) > 0:
             if isinstance(outputs, torch.Tensor):
                 outputs = self.objects
-                self.objects = []
+                self.objects = None
             else:
-                outputs = Moondream3GenerateOutput(
-                    **outputs,
-                    objects=self.objects
-                )
-                self.objects = []
+                outputs = Moondream3GenerateOutput(**outputs, objects=self.objects)
+                self.objects = None
         return outputs
 
-    def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        **model_kwargs
-    ):
+    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"]
+        model_inputs["position_ids"] += (
+            model_inputs["cache_position"].unsqueeze(0) - model_inputs["position_ids"]
+        )
         return model_inputs
 
     def _update_model_kwargs_for_generation(
@@ -1252,13 +1413,15 @@ class Moondream3ForConditionalGeneration(Moondream3PreTrainedModel, GenerationMi
         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),
+                tuple(
+                    past_state.index_select(0, beam_idx.to(past_state.device))
+                    for past_state in layer_past
+                ),
             )
         return reordered_past