ai-edge-torch-nightly 0.3.0.dev20250204__py3-none-any.whl → 0.3.0.dev20250207__py3-none-any.whl

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 @@ class Decoder(model_builder.DecoderOnlyModel):
       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 @@ class Decoder(model_builder.DecoderOnlyModel):
 
     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 @@ class Decoder2(gemma2.Gemma2):
       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,7 +31,7 @@ from torch import nn
 PROJECTION_TENSOR_NAME = "multi_modal_projector.linear"
 
 
-@dataclass
+@dataclasses.dataclass
 class PaliGemmaConfig:
   """PaliGemma model configurations."""
 
@@ -39,7 +39,6 @@ class PaliGemmaConfig:
   decoder_config: cfg.ModelConfig
 
   image_token_id: int
-  image_projection_scale: float
   image_projection_use_bias: bool = False
 
 
@@ -73,7 +72,6 @@ class PaliGemma(nn.Module):
       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 @@ class PaliGemma(nn.Module):
           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 @@ class PaliGemma(nn.Module):
         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 @@ _IMAGE_URL = flags.DEFINE_string(
 )
 _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 @@ _CHECKPOINT = {
 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"],

ai_edge_torch/generative/examples/qwen_vl/decoder.py

@@ -15,16 +15,61 @@
 
 """Example of building decoder for Qwen 2.5 VL models."""
 
+from typing import Optional, Tuple
+
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
-from torch import nn
+import torch
 
 TENSOR_NAMES = model_builder.TENSOR_NAMES
 
 
 class Decoder(model_builder.DecoderOnlyModel):
-  """A decoder for Qwen-VL model built from the Edge Generative API layers."""
-  pass
+  """A decoder for Qwen-VL model built from the Edge Generative API layers.
+
+  Besides a tensor of text token IDs, forward() can also take a tensor of
+  embeddings which may include text or image or both.
+  """
+
+  @torch.inference_mode
+  def forward(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      input_embeds: torch.Tensor = None,
+      rope: Tuple[torch.Tensor, torch.Tensor] = None,
+      mask: Optional[torch.Tensor] = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    if input_embeds is None:
+      _, seq_len = tokens.size()
+      assert self.config.max_seq_len >= seq_len, (
+          f"Cannot forward sequence of length {seq_len}, max seq length is only"
+          f" {self.config.max_seq_len}"
+      )
+      # token embeddings of shape (b, t, n_embd)
+      input_embeds = self.tok_embedding(tokens)
+
+    if rope is None:
+      # 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(input_pos, n_elem, attn_config.rotary_base)
+
+    if mask is None:
+      mask = self.mask_cache.index_select(2, input_pos)
+      mask = mask[:, :, :, : self.config.kv_cache_max]
+
+    return self._forward_with_embeds(
+        input_embeds,
+        rope,
+        mask,
+        input_pos,
+        kv_cache,
+        export_config=export_config,
+    )
 
 
 def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
@@ -82,7 +127,7 @@ def get_fake_decoder_config(**kwargs) -> cfg.ModelConfig:
   return config
 
 
-def build_decoder(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_decoder(checkpoint_path: str, **kwargs) -> torch.nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_decoder_config(**kwargs),

ai_edge_torch/generative/examples/qwen_vl/image_encoder.py

@@ -356,6 +356,12 @@ def get_fake_image_encoder_config() -> QwenVLImageConfig:
 def build_image_encoder(checkpoint_path: str) -> QwenVLImageEncoder:
   config = get_image_encoder_config()
   encoder = QwenVLImageEncoder(config)
+  load_image_encoder(checkpoint_path, encoder)
+  encoder.eval()
+  return encoder
+
+
+def load_image_encoder(checkpoint_path: str, encoder: QwenVLImageEncoder):
   loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
   # Loose the strictness because only image encoder is being loaded.
   loader.load(encoder, strict=False)
@@ -365,15 +371,12 @@ def build_image_encoder(checkpoint_path: str) -> QwenVLImageEncoder:
   state = merger_loader.get_state()
   w1_state = dict()
   w1_state["weight"] = state.pop(f"{MERGER_TENSOR_NAMES.ff_up_proj}.weight")
-  if config.merger_config.use_bias:
+  if encoder.config.merger_config.use_bias:
     w1_state["bias"] = state.pop(f"{MERGER_TENSOR_NAMES.ff_up_proj}.bias")
   encoder.merger.w1.load_state_dict(w1_state)
 
   w2_state = dict()
   w2_state["weight"] = state.pop(f"{MERGER_TENSOR_NAMES.ff_down_proj}.weight")
-  if config.merger_config.use_bias:
+  if encoder.config.merger_config.use_bias:
     w2_state["bias"] = state.pop(f"{MERGER_TENSOR_NAMES.ff_down_proj}.bias")
   encoder.merger.w2.load_state_dict(w2_state)
-
-  encoder.eval()
-  return encoder

ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py

@@ -0,0 +1,211 @@
+# Copyright 2025 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Example of building a full-stack of Qwen 2.5 VL model."""
+
+import dataclasses
+from typing import List, Optional, Tuple
+
+from ai_edge_torch.generative.examples.qwen_vl import decoder
+from ai_edge_torch.generative.examples.qwen_vl import image_encoder
+import ai_edge_torch.generative.layers.kv_cache as kv_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import model_builder
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+from torch import nn
+
+
+@dataclasses.dataclass
+class QwenVLConfig:
+  """Qwen VL model configurations."""
+
+  image_encoder_config: image_encoder.QwenVLImageConfig
+  decoder_config: cfg.ModelConfig
+  image_token_id: int
+  mrope_section: List[int]
+
+
+class QwenVL(nn.Module):
+  """Qwen VL model from the Edge Generative API."""
+
+  def __init__(self, config: QwenVLConfig):
+    super().__init__()
+
+    self.image_encoder = image_encoder.QwenVLImageEncoder(
+        config.image_encoder_config
+    )
+    self.decoder = decoder.Decoder(config.decoder_config)
+    # The amount of adjustment in input_pos to calculate RoPE properly in
+    # forward() calls after image is handled.
+    self.rope_pos_adjust = 0
+    self.config = config
+
+  @torch.inference_mode
+  def forward(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      mask: Optional[torch.Tensor] = None,
+      pixel_values: torch.Tensor = None,
+      grid_thw: torch.Tensor = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    if pixel_values is None:
+      return self.decoder(
+          tokens=tokens,
+          input_pos=input_pos,
+          kv_cache=kv_cache,
+          mask=mask,
+          rope=self._build_text_rope(input_pos),
+          input_embeds=None,
+          export_config=export_config,
+      )
+
+    input_embeds = self.decoder.tok_embedding(tokens)
+    image_embeds = self.image_encoder(pixel_values, grid_thw).unsqueeze(0)
+
+    # Merging image_embeds into text_embeds as PaliGemmaForConditionalGeneration
+    # can be done like:
+    #
+    #   image_mask = tokens == self.config.image_token_id
+    #   image_mask = image_mask.unsqueeze(-1).expand_as(input_embeds)
+    #   input_embeds = input_embeds.masked_scatter(image_mask, image_embeds)
+    #
+    # Unfortunately, torch.Tensor.masked_scatter can't be lowered on CPU.
+    # Assume that image is put at the beginning of the input sequence wrapped
+    # with vision_start and vision_end tokens.
+    input_embeds = torch.cat(
+        (
+            input_embeds[:, :1, :],
+            image_embeds,
+            input_embeds[:, image_embeds.shape[1] + 1 :, :],
+        ),
+        dim=1,
+    )
+
+    return self.decoder(
+        tokens=None,
+        input_pos=input_pos,
+        kv_cache=kv_cache,
+        mask=mask,
+        input_embeds=input_embeds,
+        rope=self._build_multimodal_rope(input_pos, grid_thw),
+        export_config=export_config,
+    )
+
+  def _build_rope(
+      self, rope_pos: torch.Tensor
+  ) -> Tuple[torch.Tensor, torch.Tensor]:
+    # ROPE parameters for all attn_configs are the same. Take the first one.
+    attn_config = self.config.decoder_config.block_config(0).attn_config
+    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
+    return self.config.decoder_config.build_rope(
+        rope_pos, n_elem, attn_config.rotary_base
+    )
+
+  def _build_text_rope(
+      self, input_pos: torch.Tensor
+  ) -> Tuple[torch.Tensor, torch.Tensor]:
+    # Reset rope_pos_adjust to 0 when input sequence starts from scratch, i.e.
+    # input_pos[0] = 0.
+    if input_pos[0] == 0:
+      self.rope_pos_adjust = 0
+    return self._build_rope(input_pos + self.rope_pos_adjust)
+
+  def _build_multimodal_rope(
+      self, input_pos: torch.Tensor, grid_thw: torch.Tensor
+  ) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Builds RoPE of multimodal input for the Qwen VL model.
+
+    It's copied from Qwen2_5_VLForConditionalGeneration.get_rope_index() and
+    simplified based on the assumption that an image is put at the beginning of
+    the input sequence with vision start and vision end tokens.
+    """
+    spatial_merge_size = self.config.image_encoder_config.spatial_merge_size
+    height = grid_thw[0][1] // spatial_merge_size
+    width = grid_thw[0][2] // spatial_merge_size
+    image_pos_max = max(height, width)
+    image_pos_count = height * width
+
+    # The position of vision end tokek and text tokens and after the image.
+    text_pos_start = image_pos_max + 1
+    text_pos_count = len(input_pos) - image_pos_count - 1
+    text_pos = torch.arange(text_pos_start, text_pos_start + text_pos_count)
+    # Set input_pos_adjust since text_pos_start has changed.
+    self.rope_pos_adjust = image_pos_max - image_pos_count
+
+    temporal_rope = self._build_image_text_rope(
+        torch.ones(image_pos_count, dtype=torch.int), text_pos
+    )
+    height_rope = self._build_image_text_rope(
+        torch.arange(1, height + 1).view(-1, 1).expand(-1, width).flatten(),
+        text_pos,
+    )
+    width_rope = self._build_image_text_rope(
+        torch.arange(1, width + 1).view(1, -1).expand(height, -1).flatten(),
+        text_pos,
+    )
+
+    return (
+        self._merge_ropes(temporal_rope[0], height_rope[0], width_rope[0]),
+        self._merge_ropes(temporal_rope[1], height_rope[1], width_rope[1]),
+    )
+
+  def _build_image_text_rope(
+      self, image_pos: torch.Tensor, text_pos: torch.Tensor
+  ) -> Tuple[torch.Tensor, torch.Tensor]:
+    return self._build_rope(
+        torch.cat((torch.zeros(1, dtype=torch.int), image_pos, text_pos))
+    )
+
+  def _merge_ropes(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor):
+    """Merges RoPE tensors based on apply_multimodal_rotary_pos_emb()."""
+    split = torch.stack([a, b, c]).split(self.config.mrope_section, dim=-1)
+    return torch.cat([m[i % 3] for i, m in enumerate(split)], dim=-1)
+
+
+def get_model_config(**kwargs) -> QwenVLConfig:
+  """Returns the model config for a PaliGemma 3B-224 model.
+
+  Returns:
+    The model config for a PaliGemma 3B model.
+  """
+  return QwenVLConfig(
+      image_encoder_config=image_encoder.get_image_encoder_config(),
+      decoder_config=decoder.get_decoder_config(**kwargs),
+      image_token_id=151655,
+      mrope_section=[16, 24, 24],
+  )
+
+
+def get_fake_model_config(**kwargs) -> QwenVLConfig:
+  return QwenVLConfig(
+      image_encoder_config=image_encoder.get_fake_image_encoder_config(),
+      decoder_config=decoder.get_fake_decoder_config(**kwargs),
+      image_token_id=127,
+  )
+
+
+def build_model(checkpoint_path: str, **kwargs) -> QwenVL:
+  config = get_model_config(**kwargs)
+  model = QwenVL(config)
+  image_encoder.load_image_encoder(checkpoint_path, model.image_encoder)
+  # Load the parameters of decoder.
+  loader = loading_utils.ModelLoader(checkpoint_path, decoder.TENSOR_NAMES)
+  loader.load(model.decoder, strict=False)
+  model.eval()
+  return model

ai_edge_torch/generative/examples/qwen_vl/verify.py

@@ -0,0 +1,143 @@
+# Copyright 2025 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Verifies the reauthored Qwen 2.5 VL model."""
+
+import logging
+import pathlib
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.qwen_vl import qwen_vl
+from ai_edge_torch.generative.layers import kv_cache
+from ai_edge_torch.generative.utilities import verifier
+from PIL import Image
+import requests
+import torch
+import transformers
+
+_IMAGE_URL = flags.DEFINE_string(
+    "image_url",
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true",
+    "The image URI to encode.",
+)
+_PROMPTS = flags.DEFINE_string(
+    "prompts",
+    "<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+
+
+class ReauthoredQwenVLWrapper(verifier.ReauthoredModelWrapper):
+  """Reauthored Qwen VL model wrapper."""
+
+  def __init__(self, model: torch.nn.Module):
+    super().__init__(model)
+    self.grid_thw = None
+
+  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 {"grid_thw": self.grid_thw}
+
+
+def main(_):
+  checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_model = (
+      transformers.Qwen2_5_VLForConditionalGeneration.from_pretrained(
+          checkpoint
+      )
+  )
+
+  # Locate the cached dir.
+  cached_config_file = transformers.utils.cached_file(
+      checkpoint, transformers.utils.CONFIG_NAME
+  )
+  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = qwen_vl.build_model(reauthored_checkpoint)
+
+  logging.info("Loading the processor from: %s", checkpoint)
+  processor = transformers.AutoProcessor.from_pretrained(checkpoint)
+
+  logging.info("Loading the image from: %s", _IMAGE_URL.value)
+  image = Image.open(requests.get(_IMAGE_URL.value, stream=True).raw)
+  inputs = processor(text=_PROMPTS.value, images=image, return_tensors="pt")
+
+  logging.info("Verifying the reauthored model with model.forward()...")
+  logging.info("Forwarding the original model...")
+  outputs_original = original_model.forward(
+      input_ids=inputs["input_ids"],
+      pixel_values=inputs["pixel_values"],
+      image_grid_thw=inputs["image_grid_thw"],
+  )
+  outputs_original = outputs_original.logits
+  logging.info("outputs_original: %s", outputs_original)
+
+  logging.info("Forwarding the reauthored model...")
+  wrapped_reauthored_model = ReauthoredQwenVLWrapper(reauthored_model)
+  wrapped_reauthored_model.grid_thw = inputs["image_grid_thw"]
+  outputs_reauthored = wrapped_reauthored_model.forward(
+      tokens=inputs["input_ids"],
+      pixel_values=inputs["pixel_values"],
+  )
+  logging.info("outputs_reauthored: %s", outputs_reauthored)
+
+  try:
+    assert torch.allclose(outputs_original, outputs_reauthored, atol=1e-01)
+  except AssertionError as e:
+    logging.error("*** FAILED *** verify with forward()")
+    raise e
+  else:
+    logging.info("*** PASSED *** verify with forward()")
+
+  logging.info("Verifying the reauthored model with model.generate()...")
+  logging.info("Generating answer with the original model...")
+  outputs_original = original_model.generate(
+      **inputs, max_new_tokens=_MAX_NEW_TOKENS.value
+  )
+  response_original = processor.decode(
+      outputs_original[0], skip_special_tokens=True
+  )
+  logging.info("outputs_from_original_model: [[%s]]", response_original)
+
+  logging.info("Generating answer with the reauthored model...")
+  outputs_reauthored = wrapped_reauthored_model.generate(
+      prompts=inputs["input_ids"],
+      pixel_values=inputs["pixel_values"],
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+  )
+  response_reauthored = processor.decode(
+      outputs_reauthored[0], skip_special_tokens=True
+  )
+  logging.info("outputs from reauthored model: [[%s]]", response_reauthored)
+
+  try:
+    assert response_original == response_reauthored
+  except AssertionError as e:
+    logging.error("*** FAILED *** verify with generate()")
+    raise e
+  else:
+    logging.info("*** PASSED *** verify with generate()")
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/qwen_vl/verify_decoder.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The AI Edge Torch Authors.
+# Copyright 2025 The AI Edge Torch Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

ai_edge_torch/odml_torch/debuginfo/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from ._build import build_mlir_debuginfo
+from ._build import build_mlir_debuginfo, build_mlir_file_debuginfo
 from ._op_polyfill import write_mlir_debuginfo_op

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import torch
+import re
 
 
 def _class_fullname(cls):
@@ -34,6 +35,29 @@ def _get_hierarchy(node: torch.fx.Node):
   return hierachy_str
 
 
+def _get_canonical_filename(filename):
+  """Remove unnecessary path prefix to make the filename more readable.
+
+  This should be factored out so that pattern is a global option that a user
+  can override.
+  """
+
+  # TODO: We should add a config option to provide a regex to strip from the
+  # debug info. Currently absolute path is used.
+  return filename
+
+
+def build_mlir_file_debuginfo(node: torch.fx.Node):
+  """Build the file and line info for the given node's lowerings in MLIR."""
+
+  if not node.stack_trace:
+    return None, None
+
+  # Note: This uses internal APIs and may break in the future.
+  pt_trace = torch.fx.graph._parse_stack_trace(node.stack_trace)
+  return _get_canonical_filename(pt_trace.file), int(pt_trace.lineno)
+
+
 def build_mlir_debuginfo(node: torch.fx.Node):
   """Build the debuginfo string for the given node's lowerings in MLIR."""
 

ai_edge_torch/odml_torch/export.py

@@ -93,7 +93,12 @@ class LoweringInterpreter(torch.fx.Interpreter):
     if info is None:
       return ir.Location.unknown()
 
-    return ir.Location.name(name=info)
+    (file, line) = debuginfo.build_mlir_file_debuginfo(node)
+    fileinfo = None
+    if file is not None:
+      fileinfo = ir.Location.file(filename=file, line=line, col=0)
+
+    return ir.Location.name(name=info, childLoc=fileinfo)
 
   def run_node(self, node: torch.fx.Node):
     loc = self._build_loc(node)

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20250204"
+__version__ = "0.3.0.dev20250207"

{ai_edge_torch_nightly-0.3.0.dev20250204.dist-info → ai_edge_torch_nightly-0.3.0.dev20250207.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250204
+Version: 0.3.0.dev20250207
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.3.0.dev20250204.dist-info → ai_edge_torch_nightly-0.3.0.dev20250207.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=4XOGz1x6yfOnkOtBndF7qE1L3Ma12ZMJNwQ7wIWkyEs,706
+ai_edge_torch/version.py,sha256=9V9FbxtqLT70Tzmv_G0qlbqixmVc0pPPJs22C_iBlHE,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=mckvxznKLXdF2HuJg_IxQaT5Ty-iWl_iXElHEugH3VI,5452
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -73,12 +73,12 @@ ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=pyxRGgMxrn
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2ukiJbQOTIUGuMEZvmwZbt3n0,4556
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=VkigoqhAr8ew95neb3TifYv-SLOSheaWKv2AH0iKDrc,2441
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=scLsguzzuHfKYDWUd2uZkKYVRzdAbQHLd-kPam8QwvM,3004
-ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=S_W-0ojRu2Vd5SLNPs1kC-70xHB8AdSWslm-yPxyezk,5478
-ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=W009ky-yobueTzdaybSCqBAvNyArLXW3jDyp5MarzZU,6376
+ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=a6ISb96xhEJc1TtaFGCUiA4msKedPTAeMvkWrfIklx4,2792
+ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=z658dW_D0Iqvo6xnh4vG7_o17-Fufndyis8Rq5yafJY,5439
+ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=GZa0Ou_DvOijB2nTL_jRvGbn0_dvJPosQAPf47yqicw,5988
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=7K1xl64UvoHaYmqWjIbahwXHfppwTQ8sN7JrpGKX1XQ,5771
-ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=CEMG9gh51ev1KXPew927a6nfampiXX9bL6m-25tNYN8,6340
-ai_edge_torch/generative/examples/paligemma/verify.py,sha256=KT3Ruy40tSESxQuy-Sw01NAI3zId1BZr6Bp7FZj1wZk,5622
+ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=x1mgRtVLxkCTvlkPow3y7ADoGTjUh5uc5pF46mxatLw,6099
+ai_edge_torch/generative/examples/paligemma/verify.py,sha256=HLcu1fWMtFFFONAqVW94rOBqq4XvFHtatX3JFGOsfZw,5345
 ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
 ai_edge_torch/generative/examples/paligemma/verify_decoder2.py,sha256=tm-UfLr0YeBRVcQsWLBOMWI9JUzHmtPEbYK2vpITpqY,2534
 ai_edge_torch/generative/examples/paligemma/verify_image_encoder.py,sha256=vNm-wTT8BD6zbX6GocfP1QrVoHl0zSvuVxoXN36eeiU,3540
@@ -94,9 +94,11 @@ ai_edge_torch/generative/examples/qwen/convert_to_tflite.py,sha256=tqvXVGNdDehda
 ai_edge_torch/generative/examples/qwen/qwen.py,sha256=Zi_qiQ1JPokXZ95jgSEnQp3F-LKzFCvWvFLKhJjnASo,4199
 ai_edge_torch/generative/examples/qwen/verify.py,sha256=9_AyEJTeUfvhhID64Rto2bflFPyXMFokdQLsseLUMiI,2775
 ai_edge_torch/generative/examples/qwen_vl/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=rD_Ch5CzuXeatqv0C3z8vU-zou1z9QDUhoB6V4YTPIg,2829
-ai_edge_torch/generative/examples/qwen_vl/image_encoder.py,sha256=QIPbcturxn5OaVsF5zkRRsyAvCM2Bojyz9XBekHOaro,13405
-ai_edge_torch/generative/examples/qwen_vl/verify_decoder.py,sha256=FEY_PifD9fQGnERzSOljFLraRIbUVF3XTnCv95A30Cs,2602
+ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=0x4iDg2cBe3PFnjVce3nj7g2rjagGHcKqRCfbASNxA8,4402
+ai_edge_torch/generative/examples/qwen_vl/image_encoder.py,sha256=OYyF0bLVYJno9azmKDqX3gT8ojYYWEyp_F8nLtltPWs,13544
+ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py,sha256=Uzl1ZPkdYIaHN9QxezqxNwagZiGOHf1VreWnqgRQwf8,7627
+ai_edge_torch/generative/examples/qwen_vl/verify.py,sha256=2GPi0Vay4a69EwBSOfPMCMjE9PTwPOQus5j2KN7HE7I,5031
+ai_edge_torch/generative/examples/qwen_vl/verify_decoder.py,sha256=xPWoOBLh2eK12KEhELLYymfL7xvc0chmYC98c6x37oo,2602
 ai_edge_torch/generative/examples/qwen_vl/verify_image_encoder.py,sha256=lQR8p6Zp7PxDN_erMf-FKLIn_Rv4BGyQHjDbModFkeY,2946
 ai_edge_torch/generative/examples/smollm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=megskv1oiPhwHSnguoG7zV-esXp1Ns_FPeMLAYKhDb0,2522
@@ -195,14 +197,14 @@ ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
 ai_edge_torch/odml_torch/_torch_future.py,sha256=jSYHf1CMTJzMizPMbu2b39hAt0ZTR6gQLq67GMe9KTo,2336
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=YN7QPrQ8W6T3YVOdyIGadfSQuBroMjIqAMB9FeUa7Ho,13447
+ai_edge_torch/odml_torch/export.py,sha256=LDyZUehM1lmT3y2bGeA94rMGRUTLxzIUm4DTlCA8tQc,13640
 ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
 ai_edge_torch/odml_torch/tf_integration.py,sha256=NN29WeXmHZ0S1RPDFHUnBi2DEjMvAtwczStPYIsQ1w8,4849
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
 ai_edge_torch/odml_torch/composite/mark_tensor.py,sha256=U--rwl-XkWKgkdXCXDn6yySug8FR66o1YFUAIoSaWW4,3523
 ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=2Y52E_gLeoXpMcPpV-svXsgN3JbEIjnPVjm0xkpTUdQ,3319
-ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=9ag6-WWRG50rPCtIV7OpIokEKu2YRyGlMZZqVPWUH6g,762
-ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=1xCXOs3-9UcsOyLFH0uyQwLu7c06iYFTo0NQ7Ckbl2I,1465
+ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=DoE3HgAtV_GNKGBDGzH2Lb7JUHvyH7TUqWbDZIObr34,789
+ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=sjpYeqgdbDmD7lhp80yc8jfWq-HxX3xuQ58ND8ZeU-I,2213
 ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
 ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNitEeg-IoBUGNfUxsDSA,798
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
@@ -227,8 +229,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.3.0.dev20250204.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250204.dist-info/METADATA,sha256=Rf4w5EMQlNWOoFIuVlXUZPU9vmXlOJW7oB4yPrtgK0c,1966
-ai_edge_torch_nightly-0.3.0.dev20250204.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250204.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250204.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250207.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250207.dist-info/METADATA,sha256=pvcJfgIOezx3rNegfvMIVrkFXmZuqnnE_zMzC9Wt37k,1966
+ai_edge_torch_nightly-0.3.0.dev20250207.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250207.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250207.dist-info/RECORD,,