ai-edge-torch-nightly 0.2.0.dev20240731__py3-none-any.whl → 0.2.0.dev20240801__py3-none-any.whl

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -26,6 +26,7 @@ import ai_edge_torch.generative.examples.stable_diffusion.decoder as decoder
 import ai_edge_torch.generative.examples.stable_diffusion.diffusion as diffusion
 from ai_edge_torch.generative.examples.stable_diffusion.encoder import Encoder
 import ai_edge_torch.generative.examples.stable_diffusion.util as util
+from ai_edge_torch.generative.quantize import quant_recipes
 import ai_edge_torch.generative.utilities.stable_diffusion_loader as stable_diffusion_loader
 
 arg_parser = argparse.ArgumentParser()
@@ -60,6 +61,7 @@ def convert_stable_diffusion_to_tflite(
     decoder_ckpt_path: str,
     image_height: int = 512,
     image_width: int = 512,
+    quantize: bool = True,
 ):
 
   clip_model = clip.CLIP(clip.get_model_config())
@@ -105,15 +107,17 @@ def convert_stable_diffusion_to_tflite(
   if not os.path.exists(output_dir):
     Path(output_dir).mkdir(parents=True, exist_ok=True)
 
+  quant_config = quant_recipes.full_int8_weight_only_recipe() if quantize else None
+
   # TODO(yichunk): convert to multi signature tflite model.
   # CLIP text encoder
-  ai_edge_torch.signature('encode', clip_model, (prompt_tokens,)).convert().export(
-      f'{output_dir}/clip.tflite'
-  )
+  ai_edge_torch.signature('encode', clip_model, (prompt_tokens,)).convert(
+      quant_config=quant_config
+  ).export(f'{output_dir}/clip.tflite')
 
   # TODO(yichunk): enable image encoder conversion
   # Image encoder
-  # ai_edge_torch.signature('encode', encoder, (input_image, noise)).convert().export(
+  # ai_edge_torch.signature('encode', encoder, (input_image, noise)).convert(quant_config=quant_config).export(
   #     f'{output_dir}/encoder.tflite'
   # )
 
@@ -122,12 +126,12 @@ def convert_stable_diffusion_to_tflite(
       'diffusion',
       diffusion_model,
       (torch.repeat_interleave(input_latents, 2, 0), context, time_embedding),
-  ).convert().export(f'{output_dir}/diffusion.tflite')
+  ).convert(quant_config=quant_config).export(f'{output_dir}/diffusion.tflite')
 
   # Image decoder
-  ai_edge_torch.signature('decode', decoder_model, (input_latents,)).convert().export(
-      f'{output_dir}/decoder.tflite'
-  )
+  ai_edge_torch.signature('decode', decoder_model, (input_latents,)).convert(
+      quant_config=quant_config
+  ).export(f'{output_dir}/decoder.tflite')
 
 
 if __name__ == '__main__':
@@ -139,4 +143,5 @@ if __name__ == '__main__':
       decoder_ckpt_path=args.decoder_ckpt,
       image_height=512,
       image_width=512,
+      quantize=True,
   )

ai_edge_torch/generative/examples/stable_diffusion/decoder.py

@@ -293,6 +293,7 @@ def get_model_config() -> unet_cfg.AutoEncoderConfig:
           qkv_fused_interleaved=False,
           rotary_percentage=0.0,
       ),
+      enable_hlfb=False,
   )
 
   mid_block_config = unet_cfg.MidBlock2DConfig(

ai_edge_torch/generative/examples/stable_diffusion/diffusion.py

@@ -294,6 +294,7 @@ class Diffusion(nn.Module):
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=attention_config,
+                            enable_hlfb=False,
                         ),
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
@@ -301,6 +302,7 @@ class Diffusion(nn.Module):
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=attention_config,
+                            enable_hlfb=False,
                         ),
                         pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                         feed_forward_block_config=unet_cfg.FeedForwardBlock2DConfig(
@@ -354,6 +356,7 @@ class Diffusion(nn.Module):
                     attention_batch_size=config.transformer_batch_size,
                     normalization_config=config.transformer_norm_config,
                     attention_config=attention_config,
+                    enable_hlfb=False,
                 ),
                 cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                     query_dim=mid_block_channels,
@@ -361,6 +364,7 @@ class Diffusion(nn.Module):
                     attention_batch_size=config.transformer_batch_size,
                     normalization_config=config.transformer_norm_config,
                     attention_config=attention_config,
+                    enable_hlfb=False,
                 ),
                 pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                 feed_forward_block_config=unet_cfg.FeedForwardBlock2DConfig(
@@ -415,6 +419,7 @@ class Diffusion(nn.Module):
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=attention_config,
+                            enable_hlfb=False,
                         ),
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
@@ -422,6 +427,7 @@ class Diffusion(nn.Module):
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=attention_config,
+                            enable_hlfb=False,
                         ),
                         pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                         feed_forward_block_config=unet_cfg.FeedForwardBlock2DConfig(

ai_edge_torch/generative/examples/stable_diffusion/pipeline.py

@@ -65,6 +65,12 @@ arg_parser.add_argument(
     choices=['k_euler', 'k_euler_ancestral', 'k_lms'],
     help='A sampler to be used to denoise the encoded image latents. Can be one of `k_lms, `k_euler`, or `k_euler_ancestral`.',
 )
+arg_parser.add_argument(
+    '--seed',
+    default=None,
+    type=int,
+    help='A seed to make generation deterministic. A random number is used if unspecified.',
+)
 
 
 class StableDiffusion:
@@ -219,4 +225,5 @@ if __name__ == '__main__':
       output_path=args.output_path,
       sampler=args.sampler,
       n_inference_steps=args.n_inference_steps,
+      seed=args.seed,
   )

ai_edge_torch/generative/quantize/quant_recipe_utils.py

@@ -41,6 +41,16 @@ def create_layer_quant_int8_dynamic() -> quant_recipe.LayerQuantRecipe:
   )
 
 
+def create_layer_quant_int8_weight_only() -> quant_recipe.LayerQuantRecipe:
+  return quant_recipe.LayerQuantRecipe(
+      activation_dtype=quant_attrs.Dtype.FP32,
+      weight_dtype=quant_attrs.Dtype.INT8,
+      mode=quant_attrs.Mode.WEIGHT_ONLY,
+      algorithm=quant_attrs.Algorithm.MIN_MAX,
+      granularity=quant_attrs.Granularity.CHANNELWISE,
+  )
+
+
 def create_layer_quant_fp16() -> quant_recipe.LayerQuantRecipe:
   return quant_recipe.LayerQuantRecipe(
       activation_dtype=quant_attrs.Dtype.FP32,

ai_edge_torch/generative/quantize/quant_recipes.py

@@ -40,6 +40,14 @@ def full_int8_dynamic_recipe() -> quant_config.QuantConfig:
   )
 
 
+def full_int8_weight_only_recipe() -> quant_config.QuantConfig:
+  return quant_config.QuantConfig(
+      generative_recipe=quant_recipe.GenerativeQuantRecipe(
+          default=quant_recipe_utils.create_layer_quant_int8_weight_only(),
+      )
+  )
+
+
 def full_fp16_recipe() -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
       generative_recipe=quant_recipe.GenerativeQuantRecipe(

ai_edge_torch/generative/test/test_quantize.py

@@ -111,6 +111,7 @@ class TestQuantizeConvert(unittest.TestCase):
       [
           (quant_recipes.full_fp16_recipe()),
           (quant_recipes.full_int8_dynamic_recipe()),
+          (quant_recipes.full_int8_weight_only_recipe()),
           (_attention_int8_dynamic_recipe()),
           (_feedforward_int8_dynamic_recipe()),
       ]

{ai_edge_torch_nightly-0.2.0.dev20240731.dist-info → ai_edge_torch_nightly-0.2.0.dev20240801.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240731
+Version: 0.2.0.dev20240801
 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.2.0.dev20240731.dist-info → ai_edge_torch_nightly-0.2.0.dev20240801.dist-info}/RECORD

@@ -54,11 +54,11 @@ ai_edge_torch/generative/examples/phi2/phi2.py,sha256=KjfTrD2OBzOfq83-XvJ6ZhmXLu
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=Lo4Dq7a3Kg-lyH56iqGtqCo5UaClQHRCTDdNagXGTo8,3535
 ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=P-cUUQaQKGKV2p-7hvLJ--RpCIA7gk8WCDRgg0pNtd0,4331
-ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=XwV1z7cVkQ947k_ERftEeL8n0NUFCJAltLtqDVfzYGI,4704
-ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=xHcmOZaW7hoWlEEEqtB4FWoHMw5AsGHPHXMNiXEfviY,13814
-ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=G-MgiEM_PpegNMePBPuNQDeUfjk42EYrVZAyJHC54AY,28468
+ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=UmKqiUbgte8PR-uslaYln-Z_TNrVWgubq_2nSyy8lQ4,4997
+ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=o-FprmF_LSxte62p0Ud1wZGE9_sC_ClX9PKnDNfJR9E,13839
+ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=MEiVP1x8kDJkvYqimtVVZt_UCTxEjcSd208Lwp8qPvc,28734
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=mgbxkeFDMkNIGmnbcFTIFPu8EWKokghiviYIOB2lE3Q,3437
-ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=AopJ-KE74lzq4QJUP_hYeiXvGth7uWv7nNKqkhtcoF8,8277
+ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=MzK0wMrpOcMhU4wLFwZnmn3eLMy8BjU-mHC_85SKP70,8465
 ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py,sha256=r9RqbyNvuvXOGu3ojtl7ZmbC7o4Pt8aUKAhN1yCdtEc,3397
 ai_edge_torch/generative/examples/stable_diffusion/util.py,sha256=NFpOfA4KN0JpShm5QvuYbQYZ844NzexWD8nV3WjMOZM,2397
 ai_edge_torch/generative/examples/stable_diffusion/samplers/__init__.py,sha256=uQWKzCD_49ackNFrt50H04dkDXxfAwUCtMWWQre5SVE,830
@@ -97,8 +97,8 @@ ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQe
 ai_edge_torch/generative/quantize/example.py,sha256=Oy-Ss1oKXMu5RVOGt8QiUwKtrHEfhbVjTXXjxPcOqDA,1536
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=n1Fm8BFC8gJa_oiwwAOOghJyHtOXYZ4q-5ZRy4pHrIw,1957
 ai_edge_torch/generative/quantize/quant_recipe.py,sha256=Y8zahKw7b_h7ajPaJZVef4jG-MoqImRCpVSbFtV_i24,5139
-ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=-vd6Qp0BdXJVKg4f0_hhwbKOi3QPIAPVqyXnJ-ZnISQ,1915
-ai_edge_torch/generative/quantize/quant_recipes.py,sha256=4OdKES9BhofzFoHut4qPVh-3ndVL9fu-BNOEEZc_2xE,1781
+ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=5yCOwHTUA-SgWqP27pvCLPBj1z_AcjXCqyPwQFo15O8,2270
+ai_edge_torch/generative/quantize/quant_recipes.py,sha256=0Kvr_o7pbMnE8VMe6Ml0FBxkHM6RJ3C14B2I1mjItjc,2030
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=FjdycEOvxRgBmQdZVufetPvkDoD7rUowIOSKV9oV5Kk,1418
 ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/translate_recipe.py,sha256=iTNPrlubmq9ia7C3zHl50J2YEMsc4o33GwL5tr5VkkE,5229
@@ -106,7 +106,7 @@ ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudj
 ai_edge_torch/generative/test/loader_test.py,sha256=N88CbrLW7Q2x1EyurwdXQ6YjsA-ySQcPxpZH3QOGp-M,3317
 ai_edge_torch/generative/test/test_experimental_ekv.py,sha256=qMR0r7Pr_t2bn-cyeA7Qw_Rl94H1NmFcqM2ua8gpDDw,4230
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=LsPTrLC1I4JW2GowTS3V9Eu257vLHr2Yj5f_qaFUX84,7589
-ai_edge_torch/generative/test/test_quantize.py,sha256=QbF7LC9olJFGXqlAVGciac7xXc4rDtCSr71tTIYuqPk,5230
+ai_edge_torch/generative/test/test_quantize.py,sha256=nHzhthe_zcXpdAC6ZyYSW_B-UYuvEHx-5cUMHXyG5Uc,5288
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/loader.py,sha256=NTaCrU2qmeJpqdAau13ZgyeOpwATqhZB68GY0LZjU6A,11438
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=zixjZryUaCSDKmfPkQvYwbPJhUyTmZ4AK_lWN8iFo68,33324
@@ -125,8 +125,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=eO9Ra160ITjQSyRBEGy6nNIVH3gYacSWDd
 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=kzIulTldq8R9E-lAZsvfSTvLu3FYEX7b9DyYM3qisXM,4485
-ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/METADATA,sha256=B2Nf7g2PWOU-bYTAByfDNV_FAKy3ah88O-Plsk-uW_M,1889
-ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240801.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240801.dist-info/METADATA,sha256=r_k99TD1FMN2U-8-xG1j24NbC6Ynph8lHBqXcY315BI,1889
+ai_edge_torch_nightly-0.2.0.dev20240801.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240801.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240801.dist-info/RECORD,,