Update modeling_gemma3_omni.py

modeling_gemma3_omni.py

@@ -1,20 +1,16 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-# ────────────────────────────────────────────────────────────────────────────────
-# 0. Monkey‑patch Gemma3TextScaledWordEmbedding.forward  ‑‑>  clone() to break view relation
-# ────────────────────────────────────────────────────────────────────────────────
 import torch
 from transformers.models.gemma3.modeling_gemma3 import Gemma3TextScaledWordEmbedding as _OrigEmb
 
+
 def _patched_forward(self: _OrigEmb, input_ids: torch.LongTensor):
     return super(_OrigEmb, self).forward(input_ids).clone()
 
+
 _OrigEmb.forward = _patched_forward
 
-# ────────────────────────────────────────────────────────────────────────────────
-# 1. Standard imports (the rest of your original file starts here)             │
-# ────────────────────────────────────────────────────────────────────────────────
 from typing import List, Optional, Tuple, Union, Callable
 
 from transformers import (
@@ -83,20 +79,24 @@ class Gemma3AudioProjectorConfig(PretrainedConfig):
 
 from torch import nn
 
+
 class LayerWiseWeightedSum(nn.Module):
     def __init__(self, num_layers: int, learnable: bool = True):
         super().__init__()
         self.num_layers = num_layers
         if learnable:
-            self.scalar = nn.Parameter(torch.zeros(num_layers))
+            self.scalar_weights = nn.Parameter(torch.zeros(num_layers))
         else:
-            self.register_buffer("scalar", torch.zeros(num_layers))
+            self.register_buffer("scalar_weights", torch.zeros(num_layers))
 
     def forward(self, hidden_states: list[torch.Tensor]) -> torch.Tensor:
-        assert len(hidden_states) == self.num_layers
-        norm_w = torch.softmax(self.scalar, dim=0).view(-1, 1, 1, 1)
-        stacked = torch.stack(hidden_states, dim=0)
-        return (norm_w * stacked).sum(dim=0)
+        if len(hidden_states) != self.num_layers:
+            raise ValueError(f"Expected {self.num_layers} hidden states, but got {len(hidden_states)}")
+
+        norm_weights = torch.softmax(self.scalar_weights, dim=0).view(-1, 1, 1, 1)
+        stacked_states = torch.stack(hidden_states, dim=0)
+        weighted_sum = (norm_weights * stacked_states).sum(dim=0)
+        return weighted_sum
 
 
 class Gemma3AudioProjector(PreTrainedModel):
@@ -143,19 +143,15 @@ class Gemma3AudioProjector(PreTrainedModel):
         self.layer_weighter = LayerWiseWeightedSum(
             num_layers=encoder_config["num_blocks"]
         )
+        self.norm = Gemma3RMSNorm(encoder_config['attention_dim'], eps=1e-6)
         self.proj = nn.Linear(encoder_config['attention_dim'], config.hidden_size, bias=False)
 
-    def forward(self, mel: torch.Tensor, mel_mask: torch.Tensor):
-        mel = mel.squeeze(1)  # (B, T, 80)
-        mel_mask = mel_mask.squeeze(1)  # (B, L)
-
-        if mel_mask.size(1) != mel.size(1):
-            mel_mask = mel_mask[..., : mel.size(1)]
-
+    def forward(self, mel: torch.Tensor, mel_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         _, out_mask, hidden_list = self.encoder(mel, mel_mask)
-        hidden_sum = self.layer_weighter(hidden_list)
-        hidden = self.proj(hidden_list[-1])
-        return hidden, out_mask
+        features = self.layer_weighter(hidden_list)
+        normalized_features = self.norm(features)
+        projected_features = self.proj(normalized_features)
+        return projected_features, out_mask
 
 
 class Gemma3VisionProjector(nn.Module):
@@ -188,6 +184,7 @@ def token_type_ids_mask_function(token_type_ids: Optional[torch.Tensor]) -> Opti
 
     def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
         return token_type_ids[batch_idx, kv_idx] != 0
+
     return inner_mask
 
 
@@ -199,7 +196,8 @@ class Gemma3OmniModel(Gemma3PreTrainedModel):
         self.vision_tower = AutoModel.from_config(config=config.vision_config)
         self.multi_modal_projector = Gemma3VisionProjector(config)
         self.audio_projector = Gemma3AudioProjector(
-            Gemma3AudioProjectorConfig(hidden_size=config.text_config.hidden_size)
+            Gemma3AudioProjectorConfig(hidden_size=config.text_config.hidden_size, n_mels=config.audio_config.n_mels,
+                                       num_hidden_layers=config.audio_config.num_hidden_layers)
         )
         self.vocab_size = config.text_config.vocab_size
 
@@ -235,7 +233,6 @@ class Gemma3OmniModel(Gemma3PreTrainedModel):
             **lm_kwargs,
     ) -> Union[Tuple, Gemma3ModelOutputWithPast]:
         if (input_ids is None) ^ (inputs_embeds is not None):
-            print("input_ids:", input_ids, "inputs_embeds:", inputs_embeds)
             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
@@ -261,15 +258,14 @@ class Gemma3OmniModel(Gemma3PreTrainedModel):
             )
 
         if pixel_values is not None and past_key_values is None:
-            image_features = self.get_image_features(pixel_values)
+            vision_outputs = self.vision_tower(pixel_values, output_hidden_states=True)
+            image_features = self.multi_modal_projector(vision_outputs.hidden_states[-1])
 
             if input_ids is None:
-                special_image_mask = inputs_embeds == self.get_input_embeddings()(
-                    torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
-                )
-            else:
-                special_image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1)
-                special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+                raise ValueError("`input_ids` are required when `pixel_values` are provided.")
+
+            special_image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1)
+            special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
 
             if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel():
                 image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0]
@@ -286,23 +282,14 @@ class Gemma3OmniModel(Gemma3PreTrainedModel):
                 input_audio_embeds, audio_attention_mask
             )
             if input_ids is None:
-                special_audio_mask = (
-                        inputs_embeds
-                        == self.get_input_embeddings()(
-                    torch.tensor(
-                        self.config.audio_token_index,
-                        dtype=torch.long,
-                        device=inputs_embeds.device,
-                    )
-                )
-                )
-            else:
-                special_audio_mask = (
-                        input_ids == self.config.audio_token_index
-                ).unsqueeze(-1)
-                special_audio_mask = special_audio_mask.expand_as(inputs_embeds).to(
-                    inputs_embeds.device
-                )
+                raise ValueError("`input_ids` are required when `input_audio_embeds` are provided.")
+
+            special_audio_mask = (
+                    input_ids == self.config.audio_token_index
+            ).unsqueeze(-1)
+            special_audio_mask = special_audio_mask.expand_as(inputs_embeds).to(
+                inputs_embeds.device
+            )
             if (
                     not is_torchdynamo_compiling()
                     and inputs_embeds[special_audio_mask].numel() != audio_features.numel()
@@ -314,9 +301,9 @@ class Gemma3OmniModel(Gemma3PreTrainedModel):
                     f"({audio_features.shape[0] * audio_features.shape[1]})."
                 )
             audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
-            inputs_embeds = inputs_embeds.masked_scatter(special_audio_mask, audio_features)
+            inputs_embeds = inputs_embeds.masked_scatter(special_audio_mask, audio_features.flatten(0, 1))
 
-        if not isinstance(causal_mask_mapping := attention_mask, dict):
+        if not isinstance(attention_mask, dict):
             mask_kwargs = {
                 "config": self.config.get_text_config(),
                 "input_embeds": inputs_embeds,
@@ -329,13 +316,13 @@ class Gemma3OmniModel(Gemma3PreTrainedModel):
                     token_type_ids.to(cache_position.device)
                 )
 
-            causal_mask_mapping = {
+            attention_mask = {
                 "full_attention": create_causal_mask(**mask_kwargs),
                 "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
             }
 
         outputs = self.language_model(
-            attention_mask=causal_mask_mapping,
+            attention_mask=attention_mask,
             position_ids=position_ids,
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
@@ -347,12 +334,16 @@ class Gemma3OmniModel(Gemma3PreTrainedModel):
             **lm_kwargs,
         )
 
+        image_hidden_states = None
+        if 'vision_outputs' in locals():
+            image_hidden_states = vision_outputs.hidden_states[-1]
+
         return Gemma3ModelOutputWithPast(
             last_hidden_state=outputs.last_hidden_state,
             past_key_values=outputs.past_key_values if use_cache else None,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-            image_hidden_states=image_features if pixel_values is not None else None,
+            image_hidden_states=image_hidden_states,
         )
 
 
@@ -476,13 +467,15 @@ class Gemma3OmniForConditionalGeneration(Gemma3PreTrainedModel, GenerationMixin)
             output = (logits,) + outputs[1:]
             return (loss,) + output if loss is not None else output
 
+        image_hidden_states = outputs.image_hidden_states if return_dict else outputs[4]
+
         return Gemma3CausalLMOutputWithPast(
             loss=loss,
             logits=logits,
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-            image_hidden_states=outputs.image_hidden_states,
+            image_hidden_states=image_hidden_states,
         )
 
 
@@ -492,4 +485,4 @@ __all__ = [
     "Gemma3VisionProjector",
     "Gemma3OmniModel",
     "Gemma3OmniForConditionalGeneration",
-]
+]