ai-edge-torch-nightly 0.3.0.dev20250214__py3-none-any.whl → 0.3.0.dev20250216__py3-none-any.whl

ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py

@@ -0,0 +1,92 @@
+# 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 converting a Qwen 2.5 VL model to multi-signature tflite model."""
+
+import os
+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.utilities import converter
+from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+
+_CHECKPOINT_PATH = flags.DEFINE_string(
+    'checkpoint_path',
+    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/qwen-vl'),
+    'The path to the model checkpoint, or directory holding the checkpoint.',
+)
+_OUTPUT_PATH = flags.DEFINE_string(
+    'output_path',
+    '/tmp/',
+    'The path to export the tflite model.',
+)
+_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
+    'output_name_prefix',
+    'qwen_vl',
+    'The prefix of the output tflite model name.',
+)
+_PREFILL_SEQ_LEN = flags.DEFINE_integer(
+    'prefill_seq_len',
+    1024,
+    'The maximum size of prefill input tensor.',
+)
+_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
+    'kv_cache_max_len',
+    1280,
+    'The maximum size of KV cache buffer, including both prefill and decode.',
+)
+_IMAGE_HEIGHT = flags.DEFINE_integer(
+    'image_height',
+    34 * 14,
+    'The height of image.',
+)
+_IMAGE_WIDTH = flags.DEFINE_integer(
+    'image_width',
+    46 * 14,
+    'The width of image.',
+)
+_QUANTIZE = flags.DEFINE_bool(
+    'quantize',
+    True,
+    'Whether the model should be quantized.',
+)
+
+
+def main(_):
+  pytorch_model = qwen_vl.build_model(
+      _CHECKPOINT_PATH.value,
+      kv_cache_max_len=_KV_CACHE_MAX_LEN.value,
+      image_size=(_IMAGE_HEIGHT.value, _IMAGE_WIDTH.value),
+  )
+
+  grid_thw = pytorch_model.image_encoder.get_grid_thw()
+  converter.convert_to_tflite(
+      pytorch_model,
+      output_path=_OUTPUT_PATH.value,
+      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
+      prefill_seq_len=_PREFILL_SEQ_LEN.value,
+      pixel_values_size=(
+          pytorch_model.image_encoder.get_pixel_values_size(grid_thw)
+      ),
+      quantize=_QUANTIZE.value,
+      config=pytorch_model.config.decoder_config,
+      export_config=ExportConfig(),
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

ai_edge_torch/generative/examples/qwen_vl/image_encoder.py

@@ -16,7 +16,7 @@
 """Example of building an image encoder of Qwen 2.5 VL model."""
 
 import dataclasses
-from typing import Optional
+from typing import List, Optional, Tuple
 
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import attention_utils
@@ -93,7 +93,7 @@ class QwenVLImageEncoder(nn.Module):
 
     # Tensor shape used to reshape pixel_values in forward() and various places.
     self.kernel_size = (
-        -1,  # batch size
+        -1,  # pixel_values.size(0)
         config.image_embedding.channels,
         config.image_embedding.temporal_patch_size,
         config.image_embedding.patch_size,
@@ -118,28 +118,22 @@ class QwenVLImageEncoder(nn.Module):
     )
     self.merger = QwenVLMerger(config)
     self.config = config
+    self.set_image_size(config.image_embedding.image_size)
 
   @torch.inference_mode
-  def forward(
-      self, pixel_values: torch.Tensor, grid_thw: torch.Tensor
-  ) -> torch.Tensor:
-    # Get window index and sequence lengths to rearrange the input tensor.
-    window_index, cu_seqlens = self._get_window_index(grid_thw)
+  def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+    # Check if the pixel value size matches with grid size and image config.
+    assert pixel_values.size() == self.get_pixel_values_size(self.grid_thw)
 
     # Embed the image and rearrange the embedding tensor.
-    pixel_reshaped = pixel_values.view(self.kernel_size)
+    pixel_reshaped = pixel_values.reshape(self.kernel_size)
     x = self.tok_embedding(pixel_reshaped)
     x = x.view(-1, self.config.embedding_dim)
-    x = self._rearrange(x, window_index).unsqueeze(0)
+    x = self._rearrange(x, self.window_index).unsqueeze(0)
 
-    # Get RoPE and attention mask arranged according to the window index.
-    cos, sin = self._get_rope(grid_thw)
-    rope = (
-        self._rearrange(cos, window_index),
-        self._rearrange(sin, window_index),
-    )
+    rope = self._get_rope(self.grid_thw, self.window_index)
 
-    mask = self._get_mask(x.shape[1], cu_seqlens)
+    mask = self._get_mask(self.grid_thw, self.cu_seqlens)
     full_mask = torch.zeros(x.shape[:2])
     for i, block in enumerate(self.transformer_blocks):
       x = block(
@@ -150,10 +144,42 @@ class QwenVLImageEncoder(nn.Module):
 
     y = self.merger.forward(self.final_norm(x))
     # Arrange the output back to the original order.
-    reverse_index = torch.argsort(window_index)
-    return y[reverse_index, ...]
-
-  def _get_rope(self, grid_thw: torch.Tensor) -> torch.Tensor:
+    return y[self.reverse_index, ...]
+
+  def set_image_size(self, image_size: Tuple[int, int]):
+    """Set the image size and pre-calculate some values including mask."""
+    self.config.image_embedding.image_size = image_size
+    self.grid_thw = self.get_grid_thw()
+
+    # Precalculate the window index which can't be lowered to HLO because of
+    # inconcrete index in:
+    #     index_new = index_padded[index_padded != -100]
+    self.window_index, self.cu_seqlens = self._get_window_index(self.grid_thw)
+
+    # Precalculate the reverse index of window_index until "vhlo.sort_v1" op is
+    # supported.
+    self.reverse_index = torch.argsort(self.window_index)
+
+  def get_grid_thw(self, num_images: int = 1) -> List[Tuple[int, int, int]]:
+    """Calculate the grid size of the input images based on the image config."""
+    height, width = self.config.image_embedding.image_size
+    patch_height = height // self.config.image_embedding.patch_size
+    patch_width = width // self.config.image_embedding.patch_size
+    # Support only image, i.e. temporal step size is always 1.
+    return [(1, patch_height, patch_width)] * num_images
+
+  def get_pixel_values_size(
+      self, grid_thw: List[Tuple[int, int, int]]
+  ) -> torch.Size:
+    """Calculate the size of pixel values tensor."""
+    dim_0 = sum(t * h * w for t, h, w in grid_thw)
+    config = self.config.image_embedding
+    dim_1 = config.channels * config.temporal_patch_size * config.patch_size**2
+    return torch.Size((dim_0, dim_1))
+
+  def _get_rope(
+      self, grid_thw: List[Tuple[int, int, int]], window_index: torch.Tensor
+  ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Get RoPE for Qwen VL model based on image grid information.
 
     It's copied from Qwen2_5_VisionTransformerPretrainedModel.rot_pos_emb() and
@@ -182,16 +208,20 @@ class QwenVLImageEncoder(nn.Module):
       wpos_ids = wpos_ids.flatten()
       pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
     pos_ids = torch.cat(pos_ids, dim=0)
-    max_grid_size = grid_thw[:, 1:].max()
+    # Assume all the heights and widths are the same for all images.
+    max_grid_size = max(grid_thw[0][1], grid_thw[0][2])
 
     cos, sin = attention_utils.build_rope_cache(
         max_grid_size,
         # ROPE parameters for all attn_configs are the same. Take the first one.
         self.config.block_config(0).attn_config.head_dim // 2,
     )
-    return cos[pos_ids].flatten(1), sin[pos_ids].flatten(1)
+    return (
+        self._rearrange(cos[pos_ids].flatten(1), window_index),
+        self._rearrange(sin[pos_ids].flatten(1), window_index),
+    )
 
-  def _get_window_index(self, grid_thw: torch.Tensor):
+  def _get_window_index(self, grid_thw: List[Tuple[int, int, int]]):
     """Get window index for Qwen VL model to rearrange the input tensor.
 
     It's copied from Qwen2_5_VisionTransformerPretrainedModel.get_window_index()
@@ -207,13 +237,10 @@ class QwenVLImageEncoder(nn.Module):
     )
 
     for grid_t, grid_h, grid_w in grid_thw:
-      llm_grid_h, llm_grid_w = (
-          grid_h // self.config.spatial_merge_size,
-          grid_w // self.config.spatial_merge_size,
-      )
-      index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(
-          grid_t, llm_grid_h, llm_grid_w
-      )
+      llm_grid_h = grid_h // self.config.spatial_merge_size
+      llm_grid_w = grid_w // self.config.spatial_merge_size
+      index = torch.arange(grid_t * llm_grid_h * llm_grid_w)
+      index = index.reshape((grid_t, llm_grid_h, llm_grid_w))
       pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size
       pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size
       num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size
@@ -236,18 +263,14 @@ class QwenVLImageEncoder(nn.Module):
       index_padded = index_padded.reshape(-1)
       index_new = index_padded[index_padded != -100]
       window_index.append(index_new + window_index_id)
-      spatial_merge_unit = (
-          self.config.spatial_merge_size * self.config.spatial_merge_size
-      )
+      spatial_merge_unit = self.config.spatial_merge_size**2
       cu_seqlens_tmp = (
           seqlens.cumsum(0) * spatial_merge_unit + cu_window_seqlens[-1]
       )
       cu_window_seqlens.extend(cu_seqlens_tmp.tolist())
-      window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
+      window_index_id += grid_t * llm_grid_h * llm_grid_w
 
     window_index = torch.cat(window_index, dim=0)
-    cu_window_seqlens = torch.tensor(cu_window_seqlens)
-    cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
     return window_index, cu_window_seqlens
 
   def _rearrange(
@@ -258,20 +281,20 @@ class QwenVLImageEncoder(nn.Module):
     It's copied from Qwen2_5_VisionTransformerPretrainedModel.forward() and
     modified accordingly.
     """
-    size = x.shape[0]
-    spatial_merge_unit = (
-        self.config.spatial_merge_size * self.config.spatial_merge_size
-    )
-    x_reshaped = x.view(size // spatial_merge_unit, spatial_merge_unit, -1)
+    spatial_merge_unit = self.config.spatial_merge_size**2
+    x_reshaped = x.view(x.size(0) // spatial_merge_unit, spatial_merge_unit, -1)
     x_rearranged = x_reshaped[window_index, ...]
-    return x_rearranged.view(size, -1)
+    return x_rearranged.view(x.shape)
 
-  def _get_mask(self, seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor:
+  def _get_mask(
+      self, grid_thw: List[Tuple[int, int, int]], cu_seqlens: List[int]
+  ) -> torch.Tensor:
     """Get attention mask for Qwen VL model.
 
     It's copied from Qwen2_5_VLVisionAttention.forward() and modified
     accordingly.
     """
+    seqlen = self.get_pixel_values_size(grid_thw)[0]
     mask = torch.full([1, 1, seqlen, seqlen], float("-inf"))
     for i in range(1, len(cu_seqlens)):
       mask[
@@ -282,7 +305,7 @@ class QwenVLImageEncoder(nn.Module):
     return mask
 
 
-def get_image_encoder_config() -> QwenVLImageConfig:
+def get_image_encoder_config(image_size: Tuple[int, int]) -> QwenVLImageConfig:
   """Returns the model config for the image encoder of a Qwen 2.5 VL model.
 
   Returns:
@@ -290,7 +313,7 @@ def get_image_encoder_config() -> QwenVLImageConfig:
   """
   image_embedding_config = cfg.ImageEmbeddingConfig(
       channels=3,
-      image_size=0,  # Not used in image encoder.
+      image_size=image_size,
       patch_size=14,
       temporal_patch_size=2,
   )
@@ -336,15 +359,13 @@ def get_image_encoder_config() -> QwenVLImageConfig:
       window_size=112,
       spatial_merge_size=2,
       full_atten_block_indexes=[7, 15, 23, 31],
-      # TODO: b/377051577 - Once RemoveSDPACompositeZeroMaskPass is removed,
-      # enable_hlfb can be set to True. See b/383865404#comment3 for details.
-      # enable_hlfb=True,
+      enable_hlfb=True,
   )
   return config
 
 
 def get_fake_image_encoder_config() -> QwenVLImageConfig:
-  config = get_image_encoder_config()
+  config = get_image_encoder_config((8, 12))
   # PaliGemma image encoder has only one block config.
   config.block_config(0).ff_config.intermediate_size = 128
   config.image_embedding.patch_size = 2
@@ -353,8 +374,11 @@ def get_fake_image_encoder_config() -> QwenVLImageConfig:
   return config
 
 
-def build_image_encoder(checkpoint_path: str) -> QwenVLImageEncoder:
-  config = get_image_encoder_config()
+def build_image_encoder(
+    checkpoint_path: str,
+    image_size: Tuple[int, int] = (34 * 14, 46 * 14),
+) -> QwenVLImageEncoder:
+  config = get_image_encoder_config(image_size)
   encoder = QwenVLImageEncoder(config)
   load_image_encoder(checkpoint_path, encoder)
   encoder.eval()

ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py

@@ -61,7 +61,6 @@ class QwenVL(nn.Module):
       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:
@@ -69,14 +68,14 @@ class QwenVL(nn.Module):
           tokens=tokens,
           input_pos=input_pos,
           kv_cache=kv_cache,
-          mask=mask,
-          rope=self._build_text_rope(input_pos),
           input_embeds=None,
+          rope=self._build_text_rope(input_pos),
+          mask=mask,
           export_config=export_config,
       )
 
     input_embeds = self.decoder.tok_embedding(tokens)
-    image_embeds = self.image_encoder(pixel_values, grid_thw).unsqueeze(0)
+    image_embeds = self.image_encoder(pixel_values).unsqueeze(0)
 
     # Merging image_embeds into text_embeds as PaliGemmaForConditionalGeneration
     # can be done like:
@@ -92,18 +91,19 @@ class QwenVL(nn.Module):
         (
             input_embeds[:, :1, :],
             image_embeds,
-            input_embeds[:, image_embeds.shape[1] + 1 :, :],
+            input_embeds[:, image_embeds.size(1) + 1 :, :],
         ),
         dim=1,
     )
 
+    grid_thw = self.image_encoder.get_grid_thw()
     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),
+        mask=mask,
         export_config=export_config,
     )
 
@@ -120,9 +120,9 @@ class QwenVL(nn.Module):
   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:
+    # Reset rope_pos_adjust to 0 when it's prefill, i.e. input has 2 or more
+    # tokens.
+    if input_pos.numel() > 1:
       self.rope_pos_adjust = 0
     return self._build_rope(input_pos + self.rope_pos_adjust)
 
@@ -178,15 +178,18 @@ class QwenVL(nn.Module):
     return torch.cat([m[i % 3] for i, m in enumerate(split)], dim=-1)
 
 
-def get_model_config(**kwargs) -> QwenVLConfig:
+def get_model_config(
+    kv_cache_max_len: int = 1024,
+    image_size: Tuple[int, int] = (34 * 14, 46 * 14),
+) -> 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_encoder_config=image_encoder.get_image_encoder_config(image_size),
+      decoder_config=decoder.get_decoder_config(kv_cache_max_len),
       image_token_id=151655,
       mrope_section=[16, 24, 24],
   )
@@ -197,6 +200,7 @@ def get_fake_model_config(**kwargs) -> QwenVLConfig:
       image_encoder_config=image_encoder.get_fake_image_encoder_config(),
       decoder_config=decoder.get_fake_decoder_config(**kwargs),
       image_token_id=127,
+      mrope_section=[16, 24, 24],
   )
 
 

ai_edge_torch/generative/examples/qwen_vl/verify.py

@@ -17,6 +17,7 @@
 
 import logging
 import pathlib
+
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.qwen_vl import qwen_vl
@@ -47,16 +48,9 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
 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"
@@ -94,7 +88,11 @@ def main(_):
 
   logging.info("Forwarding the reauthored model...")
   wrapped_reauthored_model = ReauthoredQwenVLWrapper(reauthored_model)
-  wrapped_reauthored_model.grid_thw = inputs["image_grid_thw"]
+  grid_thw = inputs["image_grid_thw"].tolist()
+  config = reauthored_model.config.image_encoder_config.image_embedding
+  reauthored_model.image_encoder.set_image_size(
+      (grid_thw[0][1] * config.patch_size, grid_thw[0][2] * config.patch_size)
+  )
   outputs_reauthored = wrapped_reauthored_model.forward(
       tokens=inputs["input_ids"],
       pixel_values=inputs["pixel_values"],

ai_edge_torch/generative/examples/qwen_vl/verify_image_encoder.py

@@ -64,9 +64,12 @@ def main(_):
   logging.info("outputs_original: %s", outputs_original)
 
   logging.info("Forwarding the reauthored model...")
-  outputs_reauthored = reauthored_model.forward(
-      image_input["pixel_values"], image_input["image_grid_thw"]
+  grid_thw = image_input["image_grid_thw"].tolist()
+  config = reauthored_model.config.image_embedding
+  reauthored_model.set_image_size(
+      (grid_thw[0][1] * config.patch_size, grid_thw[0][2] * config.patch_size)
   )
+  outputs_reauthored = reauthored_model.forward(image_input["pixel_values"])
   logging.info("outputs_reauthored: %s", outputs_reauthored)
 
   try:

ai_edge_torch/generative/layers/model_config.py

@@ -17,7 +17,7 @@
 
 import dataclasses
 import enum
-from typing import Callable, Optional, Sequence, Union
+from typing import Callable, Optional, Sequence, Tuple, Union
 from ai_edge_torch.generative.layers import rotary_position_embedding
 
 @enum.unique
@@ -174,8 +174,10 @@ class ImageEmbeddingConfig:
   """Image embedding parameters."""
 
   channels: int
-  # All images should be normalized to the size of [image_size * image_size].
-  image_size: int
+  # All images should be normalized to image_size * image_size if image_size is
+  # a single integer, or image_size[0] (height) * image_size[1] (width) if
+  # image_size is a tuple of 2 integers.
+  image_size: Union[int | Tuple[int, int]]
   patch_size: int
   # Meaningful only when image embedding is Conv3d.
   temporal_patch_size: Optional[int] = None
@@ -205,7 +207,7 @@ class ModelConfig:
   embedding_use_bias: bool = False
   # Image embedding parameters.
   image_embedding: Optional[ImageEmbeddingConfig] = None
-  # Number of image tokens 
+  # Number of image tokens
   num_mm_tokens_per_image: Optional[int] = None
   # Use bias term within LLM's HEAD.
   lm_head_use_bias: bool = False

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -28,6 +28,7 @@ from ai_edge_torch.generative.examples.paligemma import paligemma
 from ai_edge_torch.generative.examples.phi import phi2
 from ai_edge_torch.generative.examples.phi import phi3
 from ai_edge_torch.generative.examples.qwen import qwen
+from ai_edge_torch.generative.examples.qwen_vl import qwen_vl
 from ai_edge_torch.generative.examples.smollm import smollm
 from ai_edge_torch.generative.examples.stable_diffusion import clip as sd_clip
 from ai_edge_torch.generative.examples.stable_diffusion import decoder as sd_decoder
@@ -196,17 +197,15 @@ class TestModelConversion(googletest.TestCase):
     config = paligemma.get_fake_model_config(decoder_config)
     pytorch_model = paligemma.PaliGemma(config, decoder_class).eval()
 
-    image_embedding_config = config.image_encoder_config.image_embedding
-    num_patches = (
-        image_embedding_config.image_size // image_embedding_config.patch_size
-    ) ** 2
+    image_config = config.image_encoder_config.image_embedding
+    num_patches = (image_config.image_size // image_config.patch_size) ** 2
 
     # Make sure the token size is longer than the number of image patches.
     seq_len = num_patches + 10
-    tokens = torch.zeros((1, seq_len), dtype=torch.int, device="cpu")
+    tokens = torch.zeros((1, seq_len), dtype=torch.int)
     input_pos = torch.arange(0, seq_len, dtype=torch.int)
     kv = kv_cache.KVCache.from_model_config(config.decoder_config)
-    pixel_values = torch.zeros((1, 3, 8, 8), dtype=torch.float32, device="cpu")
+    pixel_values = torch.zeros((1, 3, 8, 8), dtype=torch.float32)
 
     edge_model = ai_edge_torch.signature(
         "prefill_pixel",
@@ -258,6 +257,55 @@ class TestModelConversion(googletest.TestCase):
         rtol=1e-5,
     )
 
+  @googletest.skipIf(
+      ai_edge_torch.config.in_oss,
+      reason="tests with custom ops are not supported in oss",
+  )
+  def test_qwen_vl_model(self):
+    config = qwen_vl.get_fake_model_config()
+    pytorch_model = qwen_vl.QwenVL(config).eval()
+
+    grid_thw = pytorch_model.image_encoder.get_grid_thw()
+    pixel_values_size = pytorch_model.image_encoder.get_pixel_values_size(
+        grid_thw
+    )
+
+    # Make sure the token size is longer than the number of pixel values.
+    seq_len = pixel_values_size[0] + 10
+    tokens = torch.zeros((1, seq_len), dtype=torch.int)
+    input_pos = torch.arange(0, seq_len, dtype=torch.int)
+    kv = kv_cache.KVCache.from_model_config(config.decoder_config)
+    pixel_values = torch.zeros(pixel_values_size, dtype=torch.float32)
+
+    edge_model = ai_edge_torch.signature(
+        "prefill_pixel",
+        pytorch_model,
+        sample_kwargs={
+            "tokens": tokens,
+            "input_pos": input_pos,
+            "kv_cache": kv,
+            "pixel_values": pixel_values,
+        },
+    ).convert()
+    edge_model.set_interpreter_builder(
+        self._interpreter_builder(edge_model.tflite_model())
+    )
+
+    tokens = torch.arange(1, seq_len + 1, dtype=torch.int).unsqueeze(0)
+    self.assertTrue(
+        test_utils.compare_tflite_torch(
+            edge_model,
+            pytorch_model,
+            tokens,
+            input_pos,
+            kv,
+            pixel_values=pixel_values,
+            signature_name="prefill_pixel",
+            atol=1e-3,
+            rtol=1e-5,
+        )
+    )
+
   @googletest.skipIf(
       ai_edge_torch.config.in_oss,
       reason="tests with custom ops are not supported in oss",

ai_edge_torch/generative/utilities/converter.py

@@ -170,7 +170,7 @@ def _export_helper(
 
   # For export, we create a module that captures any non-exportable,
   # arugments, e.g. the generation config object.
-  mod = ExportableModule(pytorch_model, export_config=export_config)
+  mod = ExportableModule(pytorch_model, export_config=export_config).eval()
 
   converter = converter_utils.Converter()
   for lora in loras:

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20250214.dist-info → ai_edge_torch_nightly-0.3.0.dev20250216.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250214
+Version: 0.3.0.dev20250216
 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.dev20250214.dist-info → ai_edge_torch_nightly-0.3.0.dev20250216.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=Gg-N8K4Pfmxd2OLKzGJ1nVBowEkZcjrFj8TYG8TNnWI,706
+ai_edge_torch/version.py,sha256=vklXbqGLRDju4mlU9vpIceTDodbvgQCmd7eyCsV5ckM,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
@@ -94,12 +94,13 @@ 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/convert_to_tflite.py,sha256=MXK75-Upoq_RhCbiXJEl8SKJ-msmvpVivsgfqqy-cfg,2780
 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/image_encoder.py,sha256=nHzBe_YSPnUe1d5i09v4bePQomVifzJNeUjRfprmxC0,14878
+ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py,sha256=rcYHkpO-NbF4F1Da7q2xNiTng9NHiLx59HyuOgQX5W0,7753
+ai_edge_torch/generative/examples/qwen_vl/verify.py,sha256=cKinMEDXauR5yKxtNTQk1RvwIHUG8-FOkmAie18sukY,5039
 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/qwen_vl/verify_image_encoder.py,sha256=PZ392nDoJG2OmHZ_7Jet3Zu1JkN6QErxKcDc7a-PPds,3126
 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
 ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py,sha256=CjY1i0iCYxFSjhCpQZwxkmVxILgeo0zu1m0oBrHqyDU,2311
@@ -141,7 +142,7 @@ ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbX
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
 ai_edge_torch/generative/layers/kv_cache.py,sha256=sGGAZD0mWYuO4FukZfDbHXoxpBOBE9lTYICvZzDj5F8,6400
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
-ai_edge_torch/generative/layers/model_config.py,sha256=Yqa3wqZLBe0Lj4PPTIaVFaZ--sV6NJ6k8KPjRguDvCc,8095
+ai_edge_torch/generative/layers/model_config.py,sha256=EA1Ey5-c1IOLRNANSUnZ7gtNTA0o6OJxrz_I_mp8cjw,8244
 ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=vp8dVx6tOe99neJhpbrtIt5fvN5NFw19JVH1v0yi5Mg,4154
@@ -167,12 +168,12 @@ ai_edge_torch/generative/test/test_kv_cache.py,sha256=2AulHBS3hC4b_68PNNBkRVOryp
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_lora.py,sha256=6QIM6RLTc2HrodGpp_aS3OxM9Rco2KAzEnYgotkg41M,5310
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=jfqkECCX7XKHeBAuDXrkwQJf0vM72eG3LMc5rluha84,6191
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=AJs_ARfWUqwuFRwYtQQOLd87CiD4mUDwAhq885cqc4Q,12875
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=bXJwDxSPgxVKp-_6BsEmMA3TuMUaUNiZoYomNounxco,14416
 ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728FcH-i3SyE4JSZZUgU0A,6071
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/bmm_4d.py,sha256=2BMOYiFVUsl-bjxmLkrX4N7kpO0CnhB7eDYxm_iBCr8,2533
-ai_edge_torch/generative/utilities/converter.py,sha256=K9taR0KY59dvfU_jO1yBe_p7w8lDns1Q3U6oJTTKZzM,8058
+ai_edge_torch/generative/utilities/converter.py,sha256=_PO9lYCdNNYPVsAqh8QQVMG_8TUBshKwmaR1cdT6Ang,8065
 ai_edge_torch/generative/utilities/dynamic_update_slice.py,sha256=e2mhx-Vp8sUK4EXoPtpZLSx3TViqLAKs67EhKcXBjAQ,2121
 ai_edge_torch/generative/utilities/loader.py,sha256=A3SOjPXp--AsvoP1hqj5QKWE4sgxoFc3H5EBUz_Eogc,13531
 ai_edge_torch/generative/utilities/model_builder.py,sha256=5WqcxpeTdt51nVoUwt9g5kKB5wQKj2eYbiaz7k6Ofxg,6815
@@ -229,8 +230,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.dev20250214.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250214.dist-info/METADATA,sha256=u-x1rHrzHOUQBPLQbu8r3-HvX0EvMYP1RkZQ1ZZHEKY,1966
-ai_edge_torch_nightly-0.3.0.dev20250214.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250214.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250214.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250216.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250216.dist-info/METADATA,sha256=uej57gx3UQtqqGHydXxTTrLAlbzRg48u-YmLvdPxrIk,1966
+ai_edge_torch_nightly-0.3.0.dev20250216.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250216.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250216.dist-info/RECORD,,