ai-edge-torch-nightly 0.5.0.dev20250417__py3-none-any.whl → 0.5.0.dev20250418__py3-none-any.whl

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -48,6 +48,7 @@ def main(_):
       pixel_values_size=torch.Size(
           [1, config.channels, config.image_size, config.image_size]
       ),
+      pixel_seq_len=(config.image_size // config.patch_size) ** 2,
       quantize=flags.FLAGS.quantize,
       config=pytorch_model.config.decoder_config,
       export_config=ExportConfig(),

ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py

@@ -43,6 +43,9 @@ def main(_):
   )
 
   grid_thw = pytorch_model.image_encoder.get_grid_thw()
+  spatial_merge_size = (
+      pytorch_model.config.image_encoder_config.spatial_merge_size
+  )
   converter.convert_to_tflite(
       pytorch_model,
       output_path=flags.FLAGS.output_path,
@@ -51,6 +54,10 @@ def main(_):
       pixel_values_size=(
           pytorch_model.image_encoder.get_pixel_values_size(grid_thw)
       ),
+      pixel_seq_len=(
+          (grid_thw[0][1] // spatial_merge_size)
+          * (grid_thw[0][2] // spatial_merge_size)
+      ),
       quantize=flags.FLAGS.quantize,
       config=pytorch_model.config.decoder_config,
       export_config=ExportConfig(),

ai_edge_torch/generative/utilities/converter.py

@@ -57,7 +57,7 @@ def define_conversion_flags(model_name: str):
   )
   flags.DEFINE_string(
       'output_name_prefix',
-      f'{model_name}',
+      model_name,
       'The prefix of the output tflite model name.',
   )
   flags.DEFINE_multi_integer(
@@ -91,6 +91,7 @@ def convert_to_tflite(
     output_name_prefix: str,
     prefill_seq_len: Union[int, list[int]],
     pixel_values_size: torch.Size = None,
+    pixel_seq_len: int = 0,
     quantize: bool = True,
     config: cfg.ModelConfig = None,
     lora_ranks: Optional[list[int]] = None,
@@ -133,12 +134,18 @@ def convert_to_tflite(
         use. If a list, the model will have multiple prefill signatures.
       pixel_values_size (torch.Size, optional): The size of pixel values to pass
         to the model. If None, the model is not expected to take pixel values.
+      pixel_seq_len (int, optional): The length of pixel tokens, or pixel
+        embeddings generated by the image encoder with pixel values. The actual
+        length of prefill_seq_len will be added by pixel_seq_len when pixel
+        values are passed.
       quantize (bool, optional): Whether the model should be quanized. Defaults
         to True.
       config (cfg.ModelConfig, optional): The model config used to configure KV
         cache. If None, it uses the config of the pytorch_model.
       lora_ranks (list[int], optional): The ranks of the LORA layers. If None,
         no LoRA signatures will be added.
+      export_config (ExportConfig, optional): The export configuration. If None,
+        it uses the default export configuration.
   """
   # pylint: disable=protected-access
   torch._dynamo.config.cache_size_limit = 64
@@ -173,6 +180,7 @@ def convert_to_tflite(
       output_file,
       prefill_seq_lens,
       pixel_values_size,
+      pixel_seq_len,
       quantize,
       config,
       loras,
@@ -185,6 +193,7 @@ def _export_helper(
     output_file: str,
     prefill_seq_lens: list[int],
     pixel_values_size: torch.Size,
+    pixel_seq_len: int,
     quantize: bool,
     config: cfg.ModelConfig,
     loras: list[None | lora_utils.LoRA],
@@ -197,11 +206,18 @@ def _export_helper(
     prefill_tokens_list.append(torch.full((1, seq_len), 0, dtype=torch.int))
     prefill_input_pos_list.append(torch.arange(0, seq_len, dtype=torch.int))
 
-  prefill_pixel_values = (
-      torch.full(pixel_values_size, 0, dtype=torch.float32)
-      if pixel_values_size
-      else None
-  )
+  prefill_pixel_values = None
+  prefill_tokens_list_with_pixel = []
+  prefill_input_pos_list_with_pixel = []
+  if pixel_values_size is not None:
+    prefill_pixel_values = torch.full(pixel_values_size, 0, dtype=torch.float32)
+    for seq_len in prefill_seq_lens:
+      prefill_tokens_list_with_pixel.append(
+          torch.full((1, seq_len + pixel_seq_len), 0, dtype=torch.int)
+      )
+      prefill_input_pos_list_with_pixel.append(
+          torch.arange(0, seq_len + pixel_seq_len, dtype=torch.int)
+      )
 
   if export_config.prefill_mask is None:
     prefill_masks = None
@@ -238,13 +254,11 @@ def _export_helper(
   for lora in loras:
     for i in range(len(prefill_seq_lens)):
       prefill_seq_len = prefill_seq_lens[i]
-      prefill_tokens = prefill_tokens_list[i]
-      prefill_input_pos = prefill_input_pos_list[i]
       prefill_signature_name = f'prefill_{prefill_seq_len}'
 
       sample_kwargs = {
-          'tokens': prefill_tokens,
-          'input_pos': prefill_input_pos,
+          'tokens': prefill_tokens_list[i],
+          'input_pos': prefill_input_pos_list[i],
           'kv_cache': prefill_kv,
       }
       if prefill_masks is not None:
@@ -261,13 +275,13 @@ def _export_helper(
       )
 
       if prefill_pixel_values is not None:
+        sample_kwargs['tokens'] = prefill_tokens_list_with_pixel[i]
+        sample_kwargs['input_pos'] = prefill_input_pos_list_with_pixel[i]
+        sample_kwargs['pixel_values'] = prefill_pixel_values
         converter.add_signature(
             prefill_signature_name + '_pixel',
             mod,
-            sample_kwargs={
-                **sample_kwargs,
-                'pixel_values': prefill_pixel_values,
-            },
+            sample_kwargs=sample_kwargs,
         )
 
     sample_kwargs = {

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250417"
+__version__ = "0.5.0.dev20250418"

{ai_edge_torch_nightly-0.5.0.dev20250417.dist-info → ai_edge_torch_nightly-0.5.0.dev20250418.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250417
+Version: 0.5.0.dev20250418
 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.5.0.dev20250417.dist-info → ai_edge_torch_nightly-0.5.0.dev20250418.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=kwb6M7GEr85K6sLrsbI9sNCggXojl_5TX9GeVCyP9OI,706
+ai_edge_torch/version.py,sha256=RijgVzFc62dFf_A0Pr5jwDo3wtofAoxUcfBieTnHvIw,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=GPDsXhfECjDzOut4vh_d9qWcyfpxobFMBTsC7MyJbM0,5557
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -84,7 +84,7 @@ ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=Hgp31zIQdJ
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2ukiJbQOTIUGuMEZvmwZbt3n0,4556
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=4W26ZtPF5Cb9mpHYuRM4b2QB_4W76zf4WV36KzexVjs,2446
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=2GLE4empjc8IANssR02ECFUqhdUNJV_OVHCf1UXKL8Y,1956
+ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=7HHXkC-IIu7ieBvBI4RlXs_oITz7R8a6YVYQskAs_Uk,2023
 ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=G1dwtWp_v77AI3uyIY-8g6qRP2tRH3CIKjJTeYNqFPU,5511
 ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=Z-SKdb0dd8uWT1d-FRwFx5-tJEqpdrQwiIZnFRhOtVo,6060
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=SvuR97sjkBtfkerH7Hu1UXB8kCFLpEATNbPfCbNAyfo,5614
@@ -108,7 +108,7 @@ ai_edge_torch/generative/examples/qwen/convert_to_tflite.py,sha256=om3lXL1RnA87P
 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=OcE2-8lqAukoK5hM1sqdgfXU37kxWQ84racweNAdjyk,1995
+ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=yVebRatt2SLCsGvrYTBXOM-0S2REhkpikHTyy5MCjUw,2222
 ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=7RFM25tDj_b0FkpSv8RUWir8K8v9p2jMtwZmP4VAUhw,4474
 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=mfLFrT8NPEPh9CqlJYHwh-I2y6ST7hH_vEmbZYartHQ,7764
@@ -186,7 +186,7 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=-v2Vj7Qdd3Gy
 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/converter.py,sha256=swtz69oyMOxSaCEYST_Gzd5sjGZ1qOBAfd_0xl207Nk,9766
+ai_edge_torch/generative/utilities/converter.py,sha256=LtBHjnslhL-uf4sDRoC8JIbbUD73g0QW3FiWsHUdV1g,10631
 ai_edge_torch/generative/utilities/export_config.py,sha256=8-795nyd3M34LkGhgW7hwHlJyTc2Oz1iipHK8yBhdFs,1633
 ai_edge_torch/generative/utilities/loader.py,sha256=7p__m2JryWphGlYOuRxdoT4id4_tWJEVOV7y2X4H-Ak,13737
 ai_edge_torch/generative/utilities/model_builder.py,sha256=ZYX1TxpFdj573du2QCyHJlFjx4q1m12R74fp4Gwl92A,6343
@@ -245,8 +245,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=dguMa-aEi-WDPnmGBUs2IPdEmt2IVmHOELH19uiJ1uU,3014
 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.5.0.dev20250417.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250417.dist-info/METADATA,sha256=ovMriaKRgveZtN2i-cTOM2_8BuNvgf-SYNITAte1wjs,2051
-ai_edge_torch_nightly-0.5.0.dev20250417.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250417.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250417.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250418.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250418.dist-info/METADATA,sha256=cmJPI-Zen5YAjRWQPoVwvcAwVZJGz3Jz0OuNAHs5498,2051
+ai_edge_torch_nightly-0.5.0.dev20250418.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250418.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250418.dist-info/RECORD,,