ai-edge-torch-nightly 0.3.0.dev20250203__py3-none-any.whl → 0.3.0.dev20250205__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

@@ -19,7 +19,6 @@ from typing import Optional
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -56,28 +55,34 @@ 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(tokens, input_pos, kv_cache)
+      return super().forward(
+          tokens, input_pos, kv_cache, mask, export_config=export_config
+      )
 
     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 = rotary_pos_emb.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.
-    embeds_len = input_embeds.shape[1]
     if mask is None:
+      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
+        input_embeds,
+        rope,
+        mask,
+        input_pos,
+        kv_cache,
+        export_config=export_config,
     )
 
 

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -20,7 +20,6 @@ from typing import Optional
 from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -59,33 +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)
+      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 = rotary_pos_emb.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() use 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/image_encoder.py

@@ -60,6 +60,7 @@ class SiglipVisionEncoder(nn.Module):
         kernel_size=config.image_embedding.patch_size,
         stride=config.image_embedding.patch_size,
         padding=0,
+        bias=config.embedding_use_bias,
     )
     num_patches = (
         config.image_embedding.image_size // config.image_embedding.patch_size

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/image_encoder.py

@@ -0,0 +1,379 @@
+# 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 an image encoder of Qwen 2.5 VL model."""
+
+import dataclasses
+from typing import Optional
+
+from ai_edge_torch.generative.layers import attention
+from ai_edge_torch.generative.layers import attention_utils
+from ai_edge_torch.generative.layers import builder
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="visual.blocks.{}.mlp.up_proj",
+    ff_down_proj="visual.blocks.{}.mlp.down_proj",
+    ff_gate_proj="visual.blocks.{}.mlp.gate_proj",
+    attn_fused_qkv_proj="visual.blocks.{}.attn.qkv",
+    attn_output_proj="visual.blocks.{}.attn.proj",
+    pre_attn_norm="visual.blocks.{}.norm1",
+    post_attn_norm="visual.blocks.{}.norm2",
+    embedding="visual.patch_embed.proj",
+    final_norm="visual.merger.ln_q",
+)
+
+MERGER_TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="visual.merger.mlp.0",
+    ff_down_proj="visual.merger.mlp.2",
+)
+
+
+@dataclasses.dataclass
+class QwenVLMergerConfig:
+  """Merger parameters."""
+
+  activation: cfg.ActivationConfig
+  intermediate_size: int
+  out_embedding_dim: int
+  use_bias: bool = False
+
+
+@dataclasses.dataclass
+class QwenVLImageConfig(cfg.ModelConfig):
+  """model config for Qwen 2.5 VL model."""
+
+  merger_config: Optional[QwenVLMergerConfig] = None
+  window_size: Optional[int] = None
+  spatial_merge_size: Optional[int] = None
+  full_atten_block_indexes: Optional[list[int]] = None
+
+
+class QwenVLMerger(nn.Module):
+  """Merger of Qwen 2.5 VL models from the Edge Generative API.
+
+  It's based on Qwen2_5_VLPatchMerger.
+  """
+
+  def __init__(self, config: QwenVLImageConfig):
+    super().__init__()
+    self.intermediate_size = config.merger_config.intermediate_size
+    self.w1 = nn.Linear(self.intermediate_size, self.intermediate_size)
+    self.act = builder.get_activation(config.merger_config.activation)
+    self.w2 = nn.Linear(
+        self.intermediate_size, config.merger_config.out_embedding_dim
+    )
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    x_reshaped = x.view(-1, self.intermediate_size)
+    return self.w2(self.act(self.w1(x_reshaped)))
+
+
+class QwenVLImageEncoder(nn.Module):
+  """Image encoder of Qwen 2.5 VL models from the Edge Generative API."""
+
+  def __init__(self, config: QwenVLImageConfig):
+    super().__init__()
+
+    # Tensor shape used to reshape pixel_values in forward() and various places.
+    self.kernel_size = (
+        -1,  # batch size
+        config.image_embedding.channels,
+        config.image_embedding.temporal_patch_size,
+        config.image_embedding.patch_size,
+        config.image_embedding.patch_size,
+    )
+    self.tok_embedding = nn.Conv3d(
+        in_channels=self.kernel_size[1],
+        out_channels=config.embedding_dim,
+        kernel_size=self.kernel_size[2:],
+        stride=self.kernel_size[2:],
+        padding=0,
+        bias=config.embedding_use_bias,
+    )
+
+    self.transformer_blocks = nn.ModuleList(
+        attention.TransformerBlock(config.block_config(idx), config)
+        for idx in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.merger = QwenVLMerger(config)
+    self.config = config
+
+  @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)
+
+    # Embed the image and rearrange the embedding tensor.
+    pixel_reshaped = pixel_values.view(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)
+
+    # 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),
+    )
+
+    mask = self._get_mask(x.shape[1], cu_seqlens)
+    full_mask = torch.zeros(x.shape[:2])
+    for i, block in enumerate(self.transformer_blocks):
+      x = block(
+          x,
+          rope=rope,
+          mask=full_mask if i in self.config.full_atten_block_indexes else mask,
+      )
+
+    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:
+    """Get RoPE for Qwen VL model based on image grid information.
+
+    It's copied from Qwen2_5_VisionTransformerPretrainedModel.rot_pos_emb() and
+    modified accordingly.
+    """
+    pos_ids = []
+    for t, h, w in grid_thw:
+      hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+      hpos_ids = hpos_ids.reshape(
+          h // self.config.spatial_merge_size,
+          self.config.spatial_merge_size,
+          w // self.config.spatial_merge_size,
+          self.config.spatial_merge_size,
+      )
+      hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+      hpos_ids = hpos_ids.flatten()
+
+      wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+      wpos_ids = wpos_ids.reshape(
+          h // self.config.spatial_merge_size,
+          self.config.spatial_merge_size,
+          w // self.config.spatial_merge_size,
+          self.config.spatial_merge_size,
+      )
+      wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+      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()
+
+    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)
+
+  def _get_window_index(self, grid_thw: torch.Tensor):
+    """Get window index for Qwen VL model to rearrange the input tensor.
+
+    It's copied from Qwen2_5_VisionTransformerPretrainedModel.get_window_index()
+    and modified accordingly.
+    """
+    window_index: list = []
+    cu_window_seqlens: list = [0]
+    window_index_id = 0
+    vit_merger_window_size = (
+        self.config.window_size
+        // self.config.spatial_merge_size
+        // self.config.image_embedding.patch_size
+    )
+
+    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
+      )
+      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
+      num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size
+      index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100)
+      index_padded = index_padded.reshape(
+          grid_t,
+          num_windows_h,
+          vit_merger_window_size,
+          num_windows_w,
+          vit_merger_window_size,
+      )
+      index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape(
+          grid_t,
+          num_windows_h * num_windows_w,
+          vit_merger_window_size,
+          vit_merger_window_size,
+      )
+      seqlens = (index_padded != -100).sum([2, 3]).reshape(-1)
+      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
+      )
+      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 = 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(
+      self, x: torch.Tensor, window_index: torch.Tensor
+  ) -> torch.Tensor:
+    """Rearrange the tensor according to window_index.
+
+    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)
+    x_rearranged = x_reshaped[window_index, ...]
+    return x_rearranged.view(size, -1)
+
+  def _get_mask(self, seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor:
+    """Get attention mask for Qwen VL model.
+
+    It's copied from Qwen2_5_VLVisionAttention.forward() and modified
+    accordingly.
+    """
+    mask = torch.full([1, 1, seqlen, seqlen], float("-inf"))
+    for i in range(1, len(cu_seqlens)):
+      mask[
+          ...,
+          cu_seqlens[i - 1] : cu_seqlens[i],
+          cu_seqlens[i - 1] : cu_seqlens[i],
+      ] = 0
+    return mask
+
+
+def get_image_encoder_config() -> QwenVLImageConfig:
+  """Returns the model config for the image encoder of a Qwen 2.5 VL model.
+
+  Returns:
+    The model config for the image encoder of a Qwen 2.5 VL model.
+  """
+  image_embedding_config = cfg.ImageEmbeddingConfig(
+      channels=3,
+      image_size=0,  # Not used in image encoder.
+      patch_size=14,
+      temporal_patch_size=2,
+  )
+  attn_config = cfg.AttentionConfig(
+      num_heads=16,
+      head_dim=80,
+      num_query_groups=16,
+      qkv_transpose_before_split=True,
+      qkv_use_bias=True,
+      output_proj_use_bias=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+      intermediate_size=3420,
+      use_bias=True,
+  )
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-6,
+  )
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+  merger_config = QwenVLMergerConfig(
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU),
+      intermediate_size=5120,  # embedding_dim(1280) * spatial_merge_size(2)^2
+      out_embedding_dim=2048,  # embedding_dim of decoder config.
+      use_bias=True,
+  )
+  config = QwenVLImageConfig(
+      vocab_size=0,  # Not used in image encoder.
+      num_layers=32,
+      max_seq_len=0,  # Not used in image encoder.
+      embedding_dim=1280,
+      image_embedding=image_embedding_config,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      merger_config=merger_config,
+      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,
+  )
+  return config
+
+
+def get_fake_image_encoder_config() -> QwenVLImageConfig:
+  config = get_image_encoder_config()
+  # PaliGemma image encoder has only one block config.
+  config.block_config(0).ff_config.intermediate_size = 128
+  config.image_embedding.patch_size = 2
+  config.num_layers = 2
+  config.merger_config.intermediate_size = 128
+  return config
+
+
+def build_image_encoder(checkpoint_path: str) -> QwenVLImageEncoder:
+  config = get_image_encoder_config()
+  encoder = QwenVLImageEncoder(config)
+  loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
+  # Loose the strictness because only image encoder is being loaded.
+  loader.load(encoder, strict=False)
+
+  # Load merger weights.
+  merger_loader = loading_utils.ModelLoader(checkpoint_path, None)
+  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:
+    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:
+    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/verify_image_encoder.py

@@ -0,0 +1,84 @@
+# 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 image encoder of 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 image_encoder
+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.",
+)
+
+
+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
+      )
+  )
+  original_vision_model = original_model.eval().visual
+
+  # 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 = image_encoder.build_image_encoder(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)
+  image_input = processor(images=image, text="", return_tensors="pt")
+
+  logging.info("Forwarding the original model...")
+  outputs_original = original_vision_model.forward(
+      image_input["pixel_values"], image_input["image_grid_thw"]
+  )
+  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"]
+  )
+  logging.info("outputs_reauthored: %s", outputs_reauthored)
+
+  try:
+    assert torch.allclose(
+        outputs_original, outputs_reauthored, atol=1e-03, rtol=1e-05
+    )
+  except AssertionError as e:
+    logging.error("*** FAILED *** verify with an image")
+    raise e
+  else:
+    logging.info("*** PASSED *** verify with an image")
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/layers/attention.py

@@ -224,7 +224,6 @@ class CausalSelfAttention(nn.Module):
 
     if rope is not None:
       # Compute rotary positional embedding for query and key.
-      n_elem = int(self.config.rotary_percentage * self.config.head_dim)
       cos, sin = rope
       q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
 

ai_edge_torch/generative/layers/model_config.py

@@ -177,6 +177,8 @@ class ImageEmbeddingConfig:
   # All images should be normalized to the size of [image_size * image_size].
   image_size: int
   patch_size: int
+  # Meaningful only when image embedding is Conv3d.
+  temporal_patch_size: Optional[int] = None
 
 
 @dataclasses.dataclass

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20250203.dist-info → ai_edge_torch_nightly-0.3.0.dev20250205.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250203
+Version: 0.3.0.dev20250205
 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.dev20250203.dist-info → ai_edge_torch_nightly-0.3.0.dev20250205.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=L1vAuA2I33V8wn9Ar4IzXsPvJ7RWvbgm1MNqp1h-H0E,706
+ai_edge_torch/version.py,sha256=3qCqU6b85lrBJn0A7eFSW9dGx1TkEsCXhffIwwFwUv4,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=g0Fbtf9WigOzQij7W1ksUca4eZTwVdCO2RcuFO2GD3M,5439
-ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=craPUFxlBniBz9a0Jc7VjK01jROMg5a47xJiEA1brnw,6430
-ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=u4hEHjvLaMu-UnRrISOFXKMEJIMSLa9CfpjjmSIrlSY,5731
-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/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=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
@@ -95,7 +95,9 @@ ai_edge_torch/generative/examples/qwen/qwen.py,sha256=Zi_qiQ1JPokXZ95jgSEnQp3F-L
 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/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
 ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py,sha256=CjY1i0iCYxFSjhCpQZwxkmVxILgeo0zu1m0oBrHqyDU,2311
@@ -131,13 +133,13 @@ ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=7Bk8z033M-BCXJ299f
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=4rFrppMRKlTwwZeX1ON_cdp4yUqoTOES161IZQkJF6c,1143
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=myGjal5A8yIBoqgArd2k40rZmCgD1Ya369KR7182bhI,2129
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=GrAy8CT1pEsgRoB8JQP6PlnNYk8kQ4U3YANfSiTJKn8,13776
+ai_edge_torch/generative/layers/attention.py,sha256=Pm8FLKh-NnOvUjqQC9oX5oghPbdivZvlPVkgOVTShoU,13703
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbXLFchGEKYwOOUA,5090
 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=ZVRWEGw1BnLbLCuoR71kWGqQteKp-UM1YvMbbWYlkNw,7999
+ai_edge_torch/generative/layers/model_config.py,sha256=Yqa3wqZLBe0Lj4PPTIaVFaZ--sV6NJ6k8KPjRguDvCc,8095
 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
@@ -225,8 +227,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.dev20250203.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250203.dist-info/METADATA,sha256=Jybn0dpOrId6u1ZcmYrWnjHnjLE3tk7Opt4XZ2nvGYg,1966
-ai_edge_torch_nightly-0.3.0.dev20250203.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250203.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250203.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/METADATA,sha256=F9YG6dtQw7Vh9T4m0C2z4JAiddvpobcdY-Rxjmh4WX4,1966
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/RECORD,,