Update modeling_gemma3_tiled.py

modeling_gemma3_tiled.py

@@ -6,12 +6,9 @@ are tiled into grids, with spatial rearrangement of embeddings and
 linebreak tokens between rows.
 """
 
-from typing import Optional, Union
-
 import torch
-import torch.nn as nn
-
-from transformers import Gemma3ForConditionalGeneration, Gemma3Model, AutoTokenizer
+from torch import nn
+from transformers import AutoTokenizer, Gemma3ForConditionalGeneration, Gemma3Model
 from transformers.cache_utils import Cache
 
 from .configuration_gemma3_tiled import Gemma3TiledConfig
@@ -20,171 +17,167 @@ from .configuration_gemma3_tiled import Gemma3TiledConfig
 class Gemma3TiledModel(Gemma3Model):
     """
     Gemma3 model with tiled image support.
-    
+
     Key differences from Gemma3Model:
     - get_image_features() handles tile grids and spatial rearrangement
     - get_placeholder_mask() validates tiled structure
     - Inserts linebreak embeddings (from "\n" token) between rows
     """
-    
+
     config_class = Gemma3TiledConfig
-    
+
     def __init__(self, config: Gemma3TiledConfig):
         super().__init__(config)
         self.tokens_per_tile = config.mm_tokens_per_image  # 256
-        self.tokens_per_tile_side = int(self.tokens_per_tile ** 0.5)  # 16
+        self.tokens_per_tile_side = int(self.tokens_per_tile**0.5)  # 16
 
         # Look up newline token ID from tokenizer vocab
         tokenizer = AutoTokenizer.from_pretrained(config._name_or_path)
         vocab = tokenizer.get_vocab()
         if "\n" not in vocab:
-            raise ValueError(f"Tokenizer vocab does not contain '\\n' token")
+            raise ValueError("Tokenizer vocab does not contain '\\n' token")
         self._linebreak_token_id = vocab["\n"]
 
     def get_linebreak_embedding(self) -> torch.Tensor:
         """Get the embedding for the linebreak token."""
         embedding_layer = self.get_input_embeddings()
         return embedding_layer.weight[self._linebreak_token_id]
-    
-    def get_image_features_tiled(
+
+    def _process_tiled_image(
         self,
         pixel_values: torch.Tensor,
-        tile_grid_shape: tuple[int, int],
+        grid_h: int,
+        grid_w: int,
     ) -> torch.Tensor:
         """
-        Process tiled image and return spatially arranged embeddings with linebreaks.
-        
+        Process a single tiled image and return spatially arranged embeddings with linebreaks.
+
         Args:
             pixel_values: Tensor of shape [num_tiles, 3, 896, 896]
-            tile_grid_shape: Tuple of (grid_h, grid_w)
-            
+            grid_h: Number of tile rows
+            grid_w: Number of tile columns
+
         Returns:
             Tensor of shape [total_tokens, hidden_size] where:
             total_tokens = (grid_h * 16) * (grid_w * 16) + (grid_h * 16 - 1)
         """
-        grid_h, grid_w = tile_grid_shape
         num_tiles = grid_h * grid_w
-        
+
         assert pixel_values.shape[0] == num_tiles, (
-            f"Expected {num_tiles} tiles for {grid_h}x{grid_w} grid, "
-            f"got {pixel_values.shape[0]}"
+            f"Expected {num_tiles} tiles for {grid_h}x{grid_w} grid, got {pixel_values.shape[0]}"
         )
-        
+
         # Process each tile through vision tower
         vision_outputs = self.vision_tower(pixel_values=pixel_values).last_hidden_state
-        
+
         # Project through multimodal projector
         # Output shape: [num_tiles, 256, hidden_size]
         tile_embeds = self.multi_modal_projector(vision_outputs)
-        
+
         # Reshape to spatial grid
         # [num_tiles, 256, hidden] -> [grid_h, grid_w, 16, 16, hidden]
         hidden_size = tile_embeds.shape[-1]
         tile_embeds = tile_embeds.view(
-            grid_h, grid_w, 
-            self.tokens_per_tile_side, self.tokens_per_tile_side, 
-            hidden_size
+            grid_h, grid_w, self.tokens_per_tile_side, self.tokens_per_tile_side, hidden_size
         )
-        
+
         # Rearrange to merge tiles spatially
         # We want: for each row of tiles, merge their columns
         # [grid_h, grid_w, 16, 16, hidden] -> [grid_h, 16, grid_w, 16, hidden]
         tile_embeds = tile_embeds.permute(0, 2, 1, 3, 4)
-        
+
         # Merge into full spatial grid
         # [grid_h, 16, grid_w, 16, hidden] -> [grid_h * 16, grid_w * 16, hidden]
         total_rows = grid_h * self.tokens_per_tile_side
         total_cols = grid_w * self.tokens_per_tile_side
         tile_embeds = tile_embeds.reshape(total_rows, total_cols, hidden_size)
-        
+
         # Now insert linebreak embeddings between rows
         linebreak_emb = self.get_linebreak_embedding()  # [hidden_size]
-        
+
         # Build output by interleaving rows with linebreaks
         output_parts = []
         for row_idx in range(total_rows):
             # Add the row (all columns)
             row = tile_embeds[row_idx]  # [total_cols, hidden_size]
             output_parts.append(row)
-            
+
             # Add linebreak after each row except the last
             if row_idx < total_rows - 1:
                 output_parts.append(linebreak_emb.unsqueeze(0))  # [1, hidden_size]
-        
+
         # Concatenate all parts
         output = torch.cat(output_parts, dim=0)  # [total_tokens, hidden_size]
-        
+
         return output
-    
+
     def get_image_features(
         self,
-        pixel_values,
-        tile_grid_shape=None,
+        pixel_values: torch.Tensor,
+        tile_grid_shape: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """
-        Get image features, handling both single images and tiled images.
-        Supports batched inputs where pixel_values is a list and tile_grid_shape is a list of tuples.
-        
+        Get image features for tiled images.
+
         Args:
-            pixel_values: Image tensor(s) - can be a single tensor or list of tensors
-            tile_grid_shape: If provided, treats input as tiled. Can be a single tuple or list of tuples.
-            
+            pixel_values: Concatenated tiles tensor of shape [total_tiles, 3, H, W]
+            tile_grid_shape: Tensor of shape [num_images, 2] where each row is (grid_h, grid_w).
+                If None, falls back to parent's non-tiled processing.
+
         Returns:
-            Image features tensor
+            Image features tensor of shape [total_tokens, hidden_size]
         """
         if tile_grid_shape is None:
-            # Standard single-image processing
+            # Standard single-image processing (non-tiled)
             return super().get_image_features(pixel_values)
-        
+
         # Get device and dtype from vision tower weights
         vision_weight = self.vision_tower.vision_model.embeddings.patch_embedding.weight
         target_device = vision_weight.device
         target_dtype = vision_weight.dtype
-        
-        # Handle batched inputs: pixel_values is list of tensors, tile_grid_shape is list of tuples
+
+        # Normalize tile_grid_shape: list -> tensor
         if isinstance(tile_grid_shape, list):
-            all_features = []
-            # pixel_values can be list of numpy arrays or tensors
-            if isinstance(pixel_values, (list, tuple)):
-                for pv, grid_shape in zip(pixel_values, tile_grid_shape):
-                    # Convert to tensor if needed and move to correct device/dtype
-                    if not isinstance(pv, torch.Tensor):
-                        pv = torch.tensor(pv, dtype=target_dtype, device=target_device)
-                    else:
-                        pv = pv.to(device=target_device, dtype=target_dtype)
-                    features = self.get_image_features_tiled(pv, grid_shape)
-                    all_features.append(features)
-                # Concatenate all image features
-                return torch.cat(all_features, dim=0)
-            else:
-                # pixel_values is already concatenated, but we have multiple grid shapes
-                # This shouldn't happen with proper preprocessing, fall back to first grid shape
-                return self.get_image_features_tiled(pixel_values, tile_grid_shape[0])
+            tile_grid_shape = torch.tensor(tile_grid_shape, device=target_device)
+
+        # Ensure pixel_values is tensor on correct device/dtype
+        if not isinstance(pixel_values, torch.Tensor):
+            pixel_values = torch.tensor(pixel_values, dtype=target_dtype, device=target_device)
         else:
-            # Single image case - ensure correct device/dtype
-            if not isinstance(pixel_values, torch.Tensor):
-                pixel_values = torch.tensor(pixel_values, dtype=target_dtype, device=target_device)
-            else:
-                pixel_values = pixel_values.to(device=target_device, dtype=target_dtype)
-            return self.get_image_features_tiled(pixel_values, tile_grid_shape)
-    
+            pixel_values = pixel_values.to(device=target_device, dtype=target_dtype)
+
+        # Calculate tile counts per image for splitting concatenated pixel_values
+        tile_counts = (tile_grid_shape[:, 0] * tile_grid_shape[:, 1]).tolist()
+
+        # Split concatenated pixel_values by image
+        pixel_splits = torch.split(pixel_values, tile_counts, dim=0)
+
+        # Process each image
+        all_features = []
+        for pv, grid_shape in zip(pixel_splits, tile_grid_shape.tolist()):
+            grid_h, grid_w = int(grid_shape[0]), int(grid_shape[1])
+            features = self._process_tiled_image(pv, grid_h, grid_w)
+            all_features.append(features)
+
+        return torch.cat(all_features, dim=0)
+
     def get_placeholder_mask(
         self,
         input_ids: torch.LongTensor,
         inputs_embeds: torch.FloatTensor,
         image_features: torch.FloatTensor,
-        tile_grid_shape=None,
+        tile_grid_shape: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """
         Get mask for placeholder tokens, with validation for tiled images.
-        
+
         Args:
             input_ids: Input token IDs
             inputs_embeds: Input embeddings
             image_features: Image feature embeddings
-            tile_grid_shape: If provided, validates against expected tiled structure.
-                            Can be a single tuple or list of tuples.
-            
+            tile_grid_shape: Tensor of shape [num_images, 2] where each row is (grid_h, grid_w).
+                If provided, validates against expected tiled structure.
+
         Returns:
             Boolean mask tensor
         """
@@ -195,76 +188,73 @@ class Gemma3TiledModel(Gemma3Model):
             special_image_mask = special_image_mask.all(-1)
         else:
             special_image_mask = input_ids == self.config.image_token_id
-        
+
         n_image_tokens = special_image_mask.sum().item()
-        
+
         # Validate tiled structure if applicable
         if tile_grid_shape is not None:
-            tokens_per_tile_side = int(self.config.mm_tokens_per_image ** 0.5)
-            
-            # Handle list of grid shapes (multiple images)
+            tokens_per_tile_side = int(self.config.mm_tokens_per_image**0.5)
+
+            # Normalize to tensor if list
             if isinstance(tile_grid_shape, list):
-                expected_total = 0
-                for grid_h, grid_w in tile_grid_shape:
-                    total_rows = grid_h * tokens_per_tile_side
-                    total_cols = grid_w * tokens_per_tile_side
-                    expected_img_tokens = total_rows * total_cols
-                    expected_linebreaks = total_rows - 1
-                    expected_total += expected_img_tokens + expected_linebreaks
-            else:
-                grid_h, grid_w = tile_grid_shape
+                tile_grid_shape = torch.tensor(tile_grid_shape)
+
+            # Calculate expected tokens for all images
+            expected_total = 0
+            for grid_shape in tile_grid_shape.tolist():
+                grid_h, grid_w = int(grid_shape[0]), int(grid_shape[1])
                 total_rows = grid_h * tokens_per_tile_side
                 total_cols = grid_w * tokens_per_tile_side
                 expected_img_tokens = total_rows * total_cols
                 expected_linebreaks = total_rows - 1
-                expected_total = expected_img_tokens + expected_linebreaks
-            
+                expected_total += expected_img_tokens + expected_linebreaks
+
             if n_image_tokens != expected_total:
                 raise ValueError(
                     f"Tiled image validation failed: expected {expected_total} tokens "
-                    f"for tile grid(s) {tile_grid_shape}, but found {n_image_tokens} placeholder tokens"
+                    f"for tile grid(s) {tile_grid_shape.tolist()}, but found {n_image_tokens} placeholder tokens"
                 )
-        
+
         # Standard validation
         special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
-        
+
         if inputs_embeds[special_image_mask].numel() != image_features.numel():
             raise ValueError(
                 f"Image features and image tokens do not match: "
                 f"tokens: {n_image_tokens}, features: {image_features.numel() // image_features.shape[-1]}"
             )
-        
+
         return special_image_mask
-    
+
     def forward(
         self,
-        input_ids: Optional[torch.LongTensor] = None,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Cache] = None,
-        token_type_ids: Optional[torch.LongTensor] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        tile_grid_shape: Optional[tuple[int, int]] = None,  # NEW
+        input_ids: torch.LongTensor | None = None,
+        pixel_values: torch.FloatTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        token_type_ids: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        tile_grid_shape: torch.Tensor | None = None,
         **lm_kwargs,
     ):
         """Forward pass with support for tiled images."""
-        
+
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-        
+
         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
-        
+
         # Replace image id with PAD if the image token is OOV
         if input_ids is not None and self.config.image_token_id >= self.vocab_size:
             special_image_mask = input_ids == self.config.image_token_id
@@ -272,37 +262,38 @@ class Gemma3TiledModel(Gemma3Model):
             llm_input_ids[special_image_mask] = 0
         else:
             llm_input_ids = input_ids
-        
+
         if inputs_embeds is None:
             inputs_embeds = self.get_input_embeddings()(llm_input_ids)
-        
+
         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
             )
-        
+
         # Merge text and images
         image_features = None
-        if pixel_values is not None:
+        # Check for non-empty pixel_values (empty list would pass "is not None" check)
+        has_images = pixel_values is not None and (not isinstance(pixel_values, (list, tuple)) or len(pixel_values) > 0)
+        if has_images:
             # Get image features (handles tiled if tile_grid_shape provided)
             image_features = self.get_image_features(pixel_values, tile_grid_shape)
             image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
-            
+
             # Ensure correct shape for scatter
             if image_features.dim() == 2:
                 # [total_tokens, hidden] -> [1, total_tokens, hidden] for batch dim
                 image_features = image_features.unsqueeze(0)
-            
+
             special_image_mask = self.get_placeholder_mask(
-                input_ids, inputs_embeds=inputs_embeds, image_features=image_features,
-                tile_grid_shape=tile_grid_shape
+                input_ids, inputs_embeds=inputs_embeds, image_features=image_features, tile_grid_shape=tile_grid_shape
             )
             inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
-        
+
         # Rest is same as parent - create attention masks and run through LM
         # ... (inheriting the attention mask logic from parent)
-        
+
         return super().forward(
             input_ids=None,  # We've already embedded
             pixel_values=None,  # Already processed
@@ -324,44 +315,44 @@ class Gemma3TiledModel(Gemma3Model):
 class Gemma3TiledForConditionalGeneration(Gemma3ForConditionalGeneration):
     """
     Gemma3 model for conditional generation with tiled image support.
-    
+
     This is the main model class to use for both training and inference.
     """
-    
+
     config_class = Gemma3TiledConfig
-    
+
     def __init__(self, config: Gemma3TiledConfig):
         super().__init__(config)
         # Replace the model with our tiled version
         self.model = Gemma3TiledModel(config)
-    
+
     def forward(
         self,
-        input_ids: Optional[torch.LongTensor] = None,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Cache] = None,
-        token_type_ids: Optional[torch.LongTensor] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        logits_to_keep: Union[int, torch.Tensor] = 0,
-        tile_grid_shape: Optional[tuple[int, int]] = None,  # NEW
+        input_ids: torch.LongTensor | None = None,
+        pixel_values: torch.FloatTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        token_type_ids: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        tile_grid_shape: torch.Tensor | None = None,
         **lm_kwargs,
     ):
         """Forward pass with tiled image support."""
-        
+
         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
-        
+
         outputs = self.model(
             input_ids=input_ids,
             pixel_values=pixel_values,
@@ -379,13 +370,13 @@ class Gemma3TiledForConditionalGeneration(Gemma3ForConditionalGeneration):
             tile_grid_shape=tile_grid_shape,  # Pass through
             **lm_kwargs,
         )
-        
+
         hidden_states = outputs[0]
-        
+
         # Compute logits
         slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
         logits = self.lm_head(hidden_states[:, slice_indices, :])
-        
+
         loss = None
         if labels is not None:
             # Use parent's loss computation logic
@@ -393,35 +384,35 @@ class Gemma3TiledForConditionalGeneration(Gemma3ForConditionalGeneration):
             shift_logits = logits_float[..., :-1, :]
             shift_labels = labels[..., 1:]
             if attention_mask is not None:
-                shift_attention_mask = attention_mask[:, -shift_logits.shape[1]:].to(logits.device)
+                shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(logits.device)
                 shift_logits = shift_logits[shift_attention_mask != 0].contiguous()
                 shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous()
             else:
                 shift_logits = shift_logits.contiguous()
                 shift_labels = shift_labels.contiguous()
-            
+
             loss_fct = nn.CrossEntropyLoss()
             flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
             flat_labels = shift_labels.view(-1).to(shift_logits.device)
             loss = loss_fct(flat_logits, flat_labels)
-        
+
         if not return_dict:
             output = (logits,) + outputs[1:]
             return (loss,) + output if loss is not None else output
-        
+
         from transformers.models.gemma3.modeling_gemma3 import Gemma3CausalLMOutputWithPast
-        
+
         return Gemma3CausalLMOutputWithPast(
             loss=loss,
             logits=logits,
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-            image_hidden_states=getattr(outputs, 'image_hidden_states', None),
+            image_hidden_states=getattr(outputs, "image_hidden_states", None),
         )
 
 
 __all__ = [
-    "Gemma3TiledModel",
     "Gemma3TiledForConditionalGeneration",
+    "Gemma3TiledModel",
 ]