Fix broken PaliGemma2

- PaliGemma2 doesn't use diagonal mask any more
- Set image embedding scaling factor correctly

PiperOrigin-RevId: 723233385

Author: ai-edge-bot

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -13,11 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Example of converting a PaliGemma model to multi-signature tflite model.
-
-DISCLAIMER: It works only with ODML Torch conversion backend. Refer to
-https://github.com/google-ai-edge/ai-edge-torch/blob/main/docs/pytorch_converter/README.md#use-odml-torch-conversion-backend-experimental.
-"""
+"""Example of converting a PaliGemma model to multi-signature tflite model."""
 
 import os
 import pathlib

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -55,7 +55,6 @@ def forward(
       input_embeds: torch.Tensor = None,
       mask: Optional[torch.Tensor] = None,
       export_config: Optional[model_builder.ExportConfig] = None,
-      called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       return super().forward(
@@ -64,11 +63,11 @@ def forward(
 
     assert input_embeds is not None
 
-    repo_pos = input_pos + 1  # PaliGemma position is 1-based.
+    rope_pos = input_pos + 1  # PaliGemma position is 1-based.
     # ROPE parameters for all attn_configs are the same. Take the first one.
     attn_config = self.config.block_config(0).attn_config
     n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
-    rope = self.config.build_rope(repo_pos, n_elem, attn_config.rotary_base)
+    rope = self.config.build_rope(rope_pos, n_elem, attn_config.rotary_base)
 
     # The first part of input_embeds are image embeddings. Diagonal causal mask
     # doesn't work here.

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -58,34 +58,23 @@ def forward(
       input_embeds: torch.Tensor = None,
       mask: Optional[torch.Tensor] = None,
       export_config: Optional[model_builder.ExportConfig] = None,
-      called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       return super().forward(tokens, input_pos, kv_cache, mask, export_config)
 
     assert input_embeds is not None
 
-    repo_pos = input_pos + 1  # PaliGemma2 position is 1-based.
+    rope_pos = input_pos + 1  # PaliGemma2 position is 1-based.
     # ROPE parameters for all attn_configs are the same. Take the first one.
     attn_config = self.config.block_config(0).attn_config
     n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
-    rope = self.config.build_rope(repo_pos, n_elem, attn_config.rotary_base)
+    rope = self.config.build_rope(rope_pos, n_elem, attn_config.rotary_base)
 
     if mask is None:
-      if called_by_generate:
-        # PaliGemma2 generate() uses a diagonal causal mask even with image
-        # embeds.
-        mask = [
-            self.get_attention_mask(
-                self.config.block_config(i).attn_config.attn_type, input_pos
-            )
-            for i in range(self.config.num_layers)
-        ]
-      else:
-        # By default, don't mask image embeds with a diagonal causal mask.
-        embeds_len = input_embeds.shape[1]
-        mask = torch.zeros(embeds_len, self.config.kv_cache_max)
-        mask[:, embeds_len:] = float("-inf")
+      # By default, don't mask image embeds with a diagonal causal mask.
+      embeds_len = input_embeds.shape[1]
+      mask = torch.zeros(embeds_len, self.config.kv_cache_max)
+      mask[:, embeds_len:] = float("-inf")
 
     return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -15,7 +15,7 @@
 
 """Example of building a full-stack of PaliGemma model."""
 
-from dataclasses import dataclass
+import dataclasses
 from typing import Optional
 
 from ai_edge_torch.generative.examples.paligemma import decoder
@@ -31,15 +31,14 @@
 PROJECTION_TENSOR_NAME = "multi_modal_projector.linear"
 
 
-@dataclass
+@dataclasses.dataclass
 class PaliGemmaConfig:
   """PaliGemma model configurations."""
 
   image_encoder_config: cfg.ModelConfig
   decoder_config: cfg.ModelConfig
 
   image_token_id: int
-  image_projection_scale: float
   image_projection_use_bias: bool = False
 
 
@@ -73,7 +72,6 @@ def forward(
       mask: Optional[torch.Tensor] = None,
       pixel_values: torch.Tensor = None,
       export_config: Optional[model_builder.ExportConfig] = None,
-      called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if pixel_values is None:
       return self.decoder(
@@ -83,14 +81,13 @@ def forward(
           mask=mask,
           input_embeds=None,
           export_config=export_config,
-          called_by_generate=called_by_generate,
       )
 
     input_embeds = self.decoder.tok_embedding(tokens)
 
     image_encoded = self.image_encoder(pixel_values=pixel_values)
     image_embeds = self.image_projection(image_encoded)
-    image_embeds = image_embeds / self.config.image_projection_scale
+    image_embeds = image_embeds / self.config.decoder_config.embedding_scale
 
     # Merging image_embeds into text_embeds as PaliGemmaForConditionalGeneration
     # can be done like:
@@ -116,7 +113,6 @@ def forward(
         mask=mask,
         input_embeds=input_embeds,
         export_config=export_config,
-        called_by_generate=called_by_generate,
     )
 
 
@@ -130,7 +126,6 @@ def get_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
       image_encoder_config=image_encoder.get_image_encoder_config(),
       decoder_config=get_decoder_config(**kwargs),
       image_token_id=257152,
-      image_projection_scale=2048**0.5,
       image_projection_use_bias=True,
   )
 
@@ -140,7 +135,6 @@ def get_fake_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
       image_encoder_config=image_encoder.get_fake_image_encoder_config(),
       decoder_config=get_decoder_config(**kwargs),
       image_token_id=127,
-      image_projection_scale=128**0.5,
       image_projection_use_bias=True,
   )
 

ai_edge_torch/generative/examples/paligemma/verify.py

@@ -41,7 +41,7 @@
 )
 _PROMPTS = flags.DEFINE_string(
     "prompts",
-    "describe en",
+    "<image><bos>describe en",
     "The input prompts to generate answers.",
 )
 _MAX_NEW_TOKENS = flags.DEFINE_integer(
@@ -59,16 +59,9 @@
 class ReauthoredPaliGemmaWrapper(verifier.ReauthoredModelWrapper):
   """Reauthored PaliGemma model wrapper."""
 
-  def __init__(self, model: torch.nn.Module):
-    super().__init__(model)
-    self.forward_called_by_generate = False
-
   def _init_kv_cache(self):
     return kv_cache.KVCache.from_model_config(self.model.config.decoder_config)
 
-  def _get_extra_args_for_forward(self):
-    return {"called_by_generate": self.forward_called_by_generate}
-
 
 def main(_):
   if _VERSION.value == "1":
@@ -137,7 +130,6 @@ def main(_):
   logging.info("outputs_from_original_model: [[%s]]", response_original)
 
   logging.info("Generating answer with the reauthored model...")
-  wrapped_reauthored_model.forward_called_by_generate = True
   outputs_reauthored = wrapped_reauthored_model.generate(
       prompts=inputs["input_ids"],
       pixel_values=inputs["pixel_values"],