ai-edge-torch-nightly 0.4.0.dev20250310__py3-none-any.whl → 0.4.0.dev20250312__py3-none-any.whl

ai_edge_torch/generative/examples/gemma3/gemma3.py

@@ -0,0 +1,176 @@
+# 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 Gemma3 gpu model."""
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+import xmlrpc
+
+from ai_edge_torch.generative.examples.gemma3 import decoder
+from ai_edge_torch.generative.examples.gemma3.cpu_only import image_encoder
+from ai_edge_torch.generative.layers import builder
+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
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+from torch import nn
+
+
+PROJECTION_TENSOR_NAME = "multi_modal_projector.linear"
+
+
+@dataclass
+class Gemma3MMConfig:
+  """Gemma3 model configurations."""
+
+  image_encoder_config: cfg.ModelConfig
+  decoder_config: cfg.ModelConfig
+  mm_norm_config: cfg.NormalizationConfig
+  mm_extra_tokens: int
+  image_token_id: int
+  image_projection_scale: float
+  image_projection_use_bias: bool = False
+
+
+class Gemma3MM(nn.Module):
+  """A Gemma3 multimodal model built from the Edge Generative API layers."""
+
+  def __init__(self, config: Gemma3MMConfig):
+    super().__init__()
+
+    self.image_encoder = image_encoder.SiglipVisionEncoderWithExit(
+        config.image_encoder_config
+    )
+    self.decoder = decoder.Decoder(config.decoder_config)
+    self.mm_norm = builder.build_norm(
+        config.image_encoder_config.embedding_dim,
+        config.mm_norm_config,
+    )
+    self.extra_embedding = nn.Embedding(
+        config.mm_extra_tokens, config.image_encoder_config.embedding_dim
+    )
+    self.image_projection = nn.Linear(
+        config.image_encoder_config.embedding_dim,
+        config.decoder_config.embedding_dim,
+        bias=config.image_projection_use_bias,
+    )
+    image_embedding_config = config.image_encoder_config.image_embedding
+    self.num_patches = (
+        image_embedding_config.image_size // image_embedding_config.patch_size
+    ) ** 2
+    self.config = config
+
+  @torch.inference_mode
+  def forward(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      image_indices: Optional[torch.Tensor] = None,
+      image_feat_indices: Optional[torch.Tensor] = None,
+      pixel_values: torch.Tensor = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    _, seq_len = tokens.size()
+    assert self.config.decoder_config.max_seq_len >= seq_len, (
+        f"Cannot forward sequence of length {seq_len}, max seq length is only"
+        f" {self.config.decoder_config.max_seq_len}"
+    )
+    if pixel_values is None:
+      return self.decoder(
+          tokens=tokens,
+          input_pos=input_pos,
+          kv_cache=kv_cache,
+          input_embeds=None,
+          export_config=export_config,
+      )
+    vocab_size = self.config.decoder_config.vocab_size
+    input_embeds = self.decoder.tok_embedding(
+        torch.clip(tokens, 0, vocab_size - 1)
+    )
+    if self.decoder.config.embedding_scale is not None:
+      input_embeds = input_embeds * self.decoder.config.embedding_scale
+
+    # TODO: Identify embedding path for hard tokens if required.
+    # extra_embeds = self.extra_embedding(
+    #     torch.clip(tokens - vocab_size, 0, self.config.mm_extra_tokens - 1)
+    # )
+    # extra_embeds = self.image_projection(extra_embeds)
+    # input_embeds = torch.where(tokens < self.config.decoder_config.vocab_size,
+    #                            input_embeds, extra_embeds)
+    # alternate method of implementation
+    # rows, cols = torch.where(tokens >= self.config.vocab_size)
+    # ext_embeds = self.ext_embedding(
+    #     tokens[rows, cols] - self.config.vocab_size
+    # )
+    # ext_embeds = self.mm_projection(extra_embeds)
+    # input_embeds[rows, cols, :] = extra_embeds
+
+    # Shape of pixel_values: (b, n, c, h, w)
+    batch_size, num_media, c, h, w = pixel_values.size()
+    pixel_values = pixel_values.view(-1, c, h, w)
+    image_encoded = self.image_encoder(pixel_values=pixel_values)
+    image_encoded = self.mm_norm(image_encoded)
+    image_encoded = self.image_projection(image_encoded)
+    _, num_patches, num_channels = image_encoded.size()
+    image_encoded = image_encoded.view(
+        batch_size, num_media, num_patches, num_channels
+    )
+
+    # Interleave the image soft embeddings with the text embeddings
+    for b in range(tokens.shape[0]):
+      unbatched_image_encoded = image_encoded[b]
+      image_features = unbatched_image_encoded[
+          image_indices[b], image_feat_indices[b]
+      ]
+      index_to_copy = torch.where(image_indices[b] >= 0)[0]
+      input_embeds[b] = torch.index_copy(
+          input_embeds[b], 0, index_to_copy, image_features[index_to_copy]
+      )
+    return self.decoder(
+        tokens=None,
+        input_pos=input_pos,
+        kv_cache=kv_cache,
+        input_embeds=input_embeds,
+        image_indices=image_indices,
+        export_config=export_config,
+    )
+
+def get_fake_model_config(**kwargs) -> Gemma3MMConfig:
+  return Gemma3MMConfig(
+      image_encoder_config=image_encoder.get_fake_image_encoder_config(),
+      decoder_config=decoder.get_fake_decoder_config_4b(**kwargs),
+      image_token_id=127,
+      image_projection_scale=128**0.5,
+      image_projection_use_bias=False,
+      mm_norm_config=cfg.NormalizationConfig(
+          type=cfg.NormalizationType.LAYER_NORM,
+          epsilon=1e-6,
+          enable_hlfb=True,
+      ),
+      mm_extra_tokens=32,
+  )
+
+def build_model_1b(checkpoint_path: str, **kwargs) -> decoder.Decoder:
+  if checkpoint_path:
+    model = decoder.build_model_1b(checkpoint_path, **kwargs)
+  else:
+    config = decoder.get_decoder_config_1b(**kwargs)
+    model = decoder.Decoder(config)
+  # TODO: Load the parameters of decoder from checkpoint.
+  model.eval()
+  return model

ai_edge_torch/generative/layers/unet/blocks_2d.py

@@ -122,7 +122,6 @@ class AttentionBlock2D(nn.Module):
         hidden_dim, config.normalization_config
     )
     self.attention = SelfAttention(
-        config.attention_batch_size,
         hidden_dim,
         config.attention_config,
         enable_hlfb=config.enable_hlfb,
@@ -178,7 +177,6 @@ class CrossAttentionBlock2D(nn.Module):
         config.output_dim, config.normalization_config
     )
     self.attention = CrossAttention(
-        config.attention_batch_size,
         config.query_dim,
         config.cross_dim,
         config.hidden_dim,
@@ -305,7 +303,8 @@ class TransformerBlock2D(nn.Module):
     Args:
       config (unet_cfg.TransformerBlock2Dconfig): the configuration of this
         block.
-      dim_override: in case specified, overrides config.attention_block_config.hidden_dim. Set to None by default.
+      dim_override: in case specified, overrides
+        config.attention_block_config.hidden_dim. Set to None by default.
     """
     super().__init__()
     self.config = config

ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py

@@ -41,8 +41,8 @@ class StableHLOCompositeBuilder:
     self.attr = attr
     self.name = name
     self.id = _get_uuid()
-    self._inputs = []
-    self._outputs = []
+    self._input_cnt = 0
+    self._output_cnt = 0
 
   def _mark_tensor(self, *tensors: torch.Tensor, is_input: bool):
     """Mark the input/output tensors of the StableHLO Composite."""
@@ -53,9 +53,20 @@ class StableHLOCompositeBuilder:
         else None
     )
 
-    for pos, tensor in enumerate(tensors):
+    def _pos() -> int:
+      if is_input:
+        self._input_cnt += 1
+        return self._input_cnt - 1
+      else:
+        self._output_cnt += 1
+        return self._output_cnt - 1
+
+    for tensor in tensors:
       if not isinstance(tensor, torch.Tensor):
         raise ValueError(f"input must be a torch tensor. Got {type(tensor)}.")
+
+      pos = _pos()
+
       marked_tensors.append(
           mark_tensor.mark_tensor_op(
               tensor,

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.4.0.dev20250310"
+__version__ = "0.4.0.dev20250312"

{ai_edge_torch_nightly-0.4.0.dev20250310.dist-info → ai_edge_torch_nightly-0.4.0.dev20250312.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.4.0.dev20250310
+Version: 0.4.0.dev20250312
 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.4.0.dev20250310.dist-info → ai_edge_torch_nightly-0.4.0.dev20250312.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=rkRy3DyD2DornHefvbWznzVRjj0A3967fcNBIlVQEz8,706
+ai_edge_torch/version.py,sha256=PlBamGX8JQKmFS2RJl0lWF-mxslb0_eASGwwnezOHuY,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -61,6 +61,15 @@ ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=BOQ4zYKMXEX8Adly9-Yt6FB
 ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSdAIdrectxQWJ0t328KCsA4nfHuGg,1736
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=IoBhEMwH07-tFm5-U6F2hpCsI8xynglhq1x9tIOdaPQ,1322
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=tR8RflXocDZqvuStyw9aFlzuiTllEC8rNnjrxms6_Is,5727
+ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
+ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=xAjMqhNrSv2srrBvrwCsnbLzdQXVpkZEOYImb3Mvw3w,3910
+ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=_7s_JrzwW4rX07f41VDuRLDZDJDshc3vqhXVY92K8q8,15423
+ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=n2EQVp5SrnMeb0csHrz46_gdNiHTpsApaRmcAc8xyj8,6482
+ai_edge_torch/generative/examples/gemma3/cpu_only/__init__.py,sha256=P11xO0F1MUbLMs8ySz6tu6qGDOOyK43q-HV_pqdsCUY,670
+ai_edge_torch/generative/examples/gemma3/cpu_only/convert_gemma3_to_tflite.py,sha256=4Ym4f8pvHu7dUSkTXfSToNuX8X3fhV5kKuhgEzOcyuw,3012
+ai_edge_torch/generative/examples/gemma3/cpu_only/decoder.py,sha256=fB2oYR08u7GcrWYjNbeADRZM5z1vTbE03mHXi497RRw,16140
+ai_edge_torch/generative/examples/gemma3/cpu_only/gemma3.py,sha256=NeMqW67uQEQl09R7nE3vSpT84KXmAHEg9oy4-7TVC5k,8104
+ai_edge_torch/generative/examples/gemma3/cpu_only/image_encoder.py,sha256=uRoLoBWzFtQz5wFZfPCxbkvZsgPAqSkUUsV3977GbYc,5184
 ai_edge_torch/generative/examples/llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/llama/convert_to_tflite.py,sha256=tMSsqg7LU3LR-PHtKvlWtLCqlk71mfcO9hANU4vnvDM,2734
 ai_edge_torch/generative/examples/llama/llama.py,sha256=UKvMO85_5z1vEY5MVu6QBW_vpQYA8LWHbJI4Yx6BrCc,6592
@@ -155,7 +164,7 @@ ai_edge_torch/generative/layers/experimental/kv_cache.py,sha256=0H-Rqtm6ArMxchHS
 ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py,sha256=1vMh1L3uYX4ptKQMWcAjxkL1v2-g0jmOiuai8ydp0dc,2879
 ai_edge_torch/generative/layers/experimental/types.py,sha256=bPPxw6TOCZVWdeDP3vCbOnjNP5-bdUMmfsfO-EtdazQ,2847
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=9jKzOfiBQ66bp1ZnVIAoREIifVNFx4aTlQeYMAx2_pA,29062
+ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=ZteHZXK6HKyxYji49DQ46sA9aIy7U3Jnz0HZp6hfevY,28996
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
 ai_edge_torch/generative/layers/unet/model_config.py,sha256=pPDwLawc23pfMaPVyMJlYmxVVusjMvx-l8wBwOYOH-c,9692
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -206,7 +215,7 @@ ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5x
 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/composite/stablehlo_composite_builder.py,sha256=h6DQkYVS4fnKMALIVdU6Q7J6Ehg3hMCV4C406SyIk3k,3513
 ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=3A_lMyj-B-DOhLJG6WmjKvZK5te2rXje8FrfqOhZsN0,959
 ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=6Ns2rlfOilLJEk5cUxlkRwm2uxOgEF2-0S2DMcOqr6A,3319
 ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
@@ -233,8 +242,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.4.0.dev20250310.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250310.dist-info/METADATA,sha256=GXkTBv2DzWbm9p52gwaLQNl3iyKQ8qGK72is7Wy0gMM,1966
-ai_edge_torch_nightly-0.4.0.dev20250310.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250310.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250310.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250312.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250312.dist-info/METADATA,sha256=DTFxRv9pdU_Uy4wUIU8th5ZpgLzkYVFmM81SgjJAzAo,1966
+ai_edge_torch_nightly-0.4.0.dev20250312.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250312.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250312.dist-info/RECORD,,