ai-edge-torch-nightly 0.3.0.dev20241028__py3-none-any.whl → 0.3.0.dev20241030__py3-none-any.whl

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -61,6 +61,14 @@ _QUANTIZE = flags.DEFINE_bool(
     default=True,
 )
 
+_DEVICE_TYPE = flags.DEFINE_string(
+    'device_type',
+    None,
+    help='The device type of the model. Currently supported: cpu, gpu.',
+    default='cpu',
+    required=True,
+)
+
 
 @torch.inference_mode
 def convert_stable_diffusion_to_tflite(
@@ -80,13 +88,17 @@ def convert_stable_diffusion_to_tflite(
   )
   loader.load(clip_model, strict=False)
 
-  diffusion_model = diffusion.Diffusion(diffusion.get_model_config(2))
+  diffusion_model = diffusion.Diffusion(
+      diffusion.get_model_config(batch_size=2, device_type=_DEVICE_TYPE.value)
+  )
   diffusion_loader = stable_diffusion_loader.DiffusionModelLoader(
       diffusion_ckpt_path, diffusion.TENSOR_NAMES
   )
   diffusion_loader.load(diffusion_model, strict=False)
 
-  decoder_model = decoder.Decoder(decoder.get_model_config())
+  decoder_model = decoder.Decoder(
+      decoder.get_model_config(device_type=_DEVICE_TYPE.value)
+  )
   decoder_loader = stable_diffusion_loader.AutoEncoderModelLoader(
       decoder_ckpt_path, decoder.TENSOR_NAMES
   )

ai_edge_torch/generative/examples/stable_diffusion/decoder.py

@@ -270,8 +270,8 @@ class Decoder(nn.Module):
     return x
 
 
-def get_model_config() -> unet_cfg.AutoEncoderConfig:
-  """Get configs for the Decoder of Stable Diffusion v1.5"""
+def get_model_config(device_type: str = "cpu") -> unet_cfg.AutoEncoderConfig:
+  """Get configs for the Decoder of Stable Diffusion v1.5."""
   in_channels = 3
   latent_channels = 4
   out_channels = 3
@@ -279,8 +279,14 @@ def get_model_config() -> unet_cfg.AutoEncoderConfig:
   scaling_factor = 0.18215
   layers_per_block = 3
 
+  # For now, only turns on StableHLO composite ops on GPU backend for better
+  # performance. CPU should also switch to it once the support is done.
+  enable_hlfb = True if device_type == "gpu" else False
+
   norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.GROUP_NORM, group_num=32
+      layers_cfg.NormalizationType.GROUP_NORM,
+      group_num=32,
+      enable_hlfb=enable_hlfb,
   )
 
   att_config = unet_cfg.AttentionBlock2DConfig(
@@ -298,7 +304,7 @@ def get_model_config() -> unet_cfg.AutoEncoderConfig:
           rotary_base=0,
           rotary_percentage=0.0,
       ),
-      enable_hlfb=False,
+      enable_hlfb=enable_hlfb,
   )
 
   mid_block_config = unet_cfg.MidBlock2DConfig(
@@ -327,7 +333,9 @@ def get_model_config() -> unet_cfg.AutoEncoderConfig:
   return config
 
 
-def get_fake_model_config() -> unet_cfg.AutoEncoderConfig:
+def get_fake_model_config(
+    device_type: str = "cpu",
+) -> unet_cfg.AutoEncoderConfig:
   """Get fake configs for the Decoder of Stable Diffusion v1.5 for testing."""
   in_channels = 3
   latent_channels = 4
@@ -336,8 +344,14 @@ def get_fake_model_config() -> unet_cfg.AutoEncoderConfig:
   scaling_factor = 0.18215
   layers_per_block = 2
 
+  # For now, only turns on StableHLO composite ops on GPU backend for better
+  # performance. CPU should also switch to it once the support is done.
+  enable_hlfb = True if device_type == "gpu" else False
+
   norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.GROUP_NORM, group_num=2
+      layers_cfg.NormalizationType.GROUP_NORM,
+      group_num=2,
+      enable_hlfb=enable_hlfb,
   )
 
   att_config = unet_cfg.AttentionBlock2DConfig(
@@ -355,7 +369,7 @@ def get_fake_model_config() -> unet_cfg.AutoEncoderConfig:
           rotary_base=0,
           rotary_percentage=0.0,
       ),
-      enable_hlfb=False,
+      enable_hlfb=enable_hlfb,
   )
 
   mid_block_config = unet_cfg.MidBlock2DConfig(

ai_edge_torch/generative/examples/stable_diffusion/diffusion.py

@@ -333,7 +333,7 @@ class Diffusion(nn.Module):
                                 dim=output_channel,
                                 num_query_groups=config.transformer_num_attention_heads,
                             ),
-                            enable_hlfb=False,
+                            enable_hlfb=config.enable_hlfb,
                         ),
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
@@ -347,7 +347,7 @@ class Diffusion(nn.Module):
                                 dim=output_channel,
                                 num_query_groups=config.transformer_num_attention_heads,
                             ),
-                            enable_hlfb=False,
+                            enable_hlfb=config.enable_hlfb,
                         ),
                         pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                         feed_forward_block_config=unet_cfg.FeedForwardBlock2DConfig(
@@ -405,7 +405,7 @@ class Diffusion(nn.Module):
                         dim=mid_block_channels,
                         num_query_groups=config.transformer_num_attention_heads,
                     ),
-                    enable_hlfb=False,
+                    enable_hlfb=config.enable_hlfb,
                 ),
                 cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                     query_dim=mid_block_channels,
@@ -419,7 +419,7 @@ class Diffusion(nn.Module):
                         dim=mid_block_channels,
                         num_query_groups=config.transformer_num_attention_heads,
                     ),
-                    enable_hlfb=False,
+                    enable_hlfb=config.enable_hlfb,
                 ),
                 pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                 feed_forward_block_config=unet_cfg.FeedForwardBlock2DConfig(
@@ -478,7 +478,7 @@ class Diffusion(nn.Module):
                                 dim=output_channel,
                                 num_query_groups=config.transformer_num_attention_heads,
                             ),
-                            enable_hlfb=False,
+                            enable_hlfb=config.enable_hlfb,
                         ),
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
@@ -492,7 +492,7 @@ class Diffusion(nn.Module):
                                 dim=output_channel,
                                 num_query_groups=config.transformer_num_attention_heads,
                             ),
-                            enable_hlfb=False,
+                            enable_hlfb=config.enable_hlfb,
                         ),
                         pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                         feed_forward_block_config=unet_cfg.FeedForwardBlock2DConfig(
@@ -581,13 +581,16 @@ class Diffusion(nn.Module):
     return x
 
 
-def get_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
-  """Get configs for the Diffusion model of Stable Diffusion v1.5
+def get_model_config(
+    batch_size: int, device_type: str = "cpu"
+) -> unet_cfg.DiffusionModelConfig:
+  """Get configs for the Diffusion model of Stable Diffusion v1.5.
 
   Args:
     batch_size (int): the batch size of input.
+    device_type (str): the device type of the model. Default to "cpu".
 
-  Retruns:
+  Returns:
     The configuration of diffusion model of Stable Diffusion v1.5.
   """
   in_channels = 4
@@ -596,9 +599,15 @@ def get_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
   layers_per_block = 2
   downsample_padding = 1
 
+  # For now, only turns on StableHLO composite ops on GPU backend for better
+  # performance. CPU should also switch to it once the support is done.
+  enable_hlfb = True if device_type == "gpu" else False
+
   # Residual configs.
   residual_norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.GROUP_NORM, group_num=32
+      layers_cfg.NormalizationType.GROUP_NORM,
+      group_num=32,
+      enable_hlfb=enable_hlfb,
   )
   residual_activation_type = layers_cfg.ActivationType.SILU
 
@@ -607,10 +616,14 @@ def get_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
   transformer_batch_size = batch_size
   transformer_cross_attention_dim = 768  # Embedding from CLIP model
   transformer_pre_conv_norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.GROUP_NORM, epsilon=1e-6, group_num=32
+      layers_cfg.NormalizationType.GROUP_NORM,
+      epsilon=1e-6,
+      group_num=32,
+      enable_hlfb=enable_hlfb,
   )
   transformer_norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.LAYER_NORM
+      layers_cfg.NormalizationType.LAYER_NORM,
+      enable_hlfb=enable_hlfb,
   )
   transformer_ff_activation_type = layers_cfg.ActivationType.GE_GLU
 
@@ -623,7 +636,9 @@ def get_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
 
   # Finaly layer configs.
   final_norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.GROUP_NORM, group_num=32
+      layers_cfg.NormalizationType.GROUP_NORM,
+      group_num=32,
+      enable_hlfb=enable_hlfb,
   )
   final_activation_type = layers_cfg.ActivationType.SILU
 
@@ -646,16 +661,20 @@ def get_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
       time_embedding_blocks_dim=time_embedding_blocks_dim,
       final_norm_config=final_norm_config,
       final_activation_type=final_activation_type,
+      enable_hlfb=enable_hlfb,
   )
 
 
-def get_fake_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
+def get_fake_model_config(
+    batch_size: int, device_type: str = "cpu"
+) -> unet_cfg.DiffusionModelConfig:
   """Get fake configs for the Diffusion model of Stable Diffusion v1.5 for testing.
 
   Args:
     batch_size (int): the batch size of input.
+    device_type (str): the device type of the model. Default to "cpu".
 
-  Retruns:
+  Returns:
     The configuration of diffusion model of Stable Diffusion v1.5.
   """
   in_channels = 4
@@ -664,9 +683,15 @@ def get_fake_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
   layers_per_block = 1
   downsample_padding = 1
 
+  # For now, only turns on StableHLO composite ops on GPU backend for better
+  # performance. CPU should also switch to it once the support is done.
+  enable_hlfb = True if device_type == "gpu" else False
+
   # Residual configs.
   residual_norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.GROUP_NORM, group_num=2
+      layers_cfg.NormalizationType.GROUP_NORM,
+      group_num=2,
+      enable_hlfb=enable_hlfb,
   )
   residual_activation_type = layers_cfg.ActivationType.SILU
 
@@ -675,10 +700,14 @@ def get_fake_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
   transformer_batch_size = batch_size
   transformer_cross_attention_dim = 4  # Embedding from CLIP model
   transformer_pre_conv_norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.GROUP_NORM, epsilon=1e-6, group_num=2
+      layers_cfg.NormalizationType.GROUP_NORM,
+      epsilon=1e-6,
+      group_num=2,
+      enable_hlfb=enable_hlfb,
   )
   transformer_norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.LAYER_NORM
+      layers_cfg.NormalizationType.LAYER_NORM,
+      enable_hlfb=enable_hlfb,
   )
   transformer_ff_activation_type = layers_cfg.ActivationType.GE_GLU
 
@@ -691,7 +720,9 @@ def get_fake_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
 
   # Finaly layer configs.
   final_norm_config = layers_cfg.NormalizationConfig(
-      layers_cfg.NormalizationType.GROUP_NORM, group_num=2
+      layers_cfg.NormalizationType.GROUP_NORM,
+      group_num=2,
+      enable_hlfb=enable_hlfb,
   )
   final_activation_type = layers_cfg.ActivationType.SILU
 
@@ -714,4 +745,5 @@ def get_fake_model_config(batch_size: int) -> unet_cfg.DiffusionModelConfig:
       time_embedding_blocks_dim=time_embedding_blocks_dim,
       final_norm_config=final_norm_config,
       final_activation_type=final_activation_type,
+      enable_hlfb=enable_hlfb,
   )

ai_edge_torch/generative/layers/unet/model_config.py

@@ -269,3 +269,6 @@ class DiffusionModelConfig:
 
   # Activation type used in final layer
   final_activation_type: layers_cfg.ActivationType
+
+  # Whether to enable StableHLO composite ops in the model.
+  enable_hlfb: bool = False

ai_edge_torch/generative/quantize/example.py

@@ -15,7 +15,9 @@
 
 import ai_edge_torch
 from ai_edge_torch.generative.examples.gemma import gemma1
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.quantize import quant_recipes
+from ai_edge_torch.generative.utilities import model_builder
 import numpy as np
 import torch
 
@@ -23,11 +25,12 @@ import torch
 def main():
   # Build a PyTorch model as usual
   config = gemma1.get_fake_model_config()
-  model = gemma1.Gemma(config)
+  model = model_builder.DecoderOnlyModel(config).eval()
   idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
   tokens = torch.full((1, 10), 0, dtype=torch.int, device="cpu")
   tokens[0, :4] = idx
   input_pos = torch.arange(0, 10, dtype=torch.int)
+  kv = kv_utils.KVCache.from_model_config(config)
 
   # Create a quantization recipe to be applied to the model
   quant_config = quant_recipes.full_int8_dynamic_recipe()
@@ -35,7 +38,7 @@ def main():
 
   # Convert with quantization
   edge_model = ai_edge_torch.convert(
-      model, (tokens, input_pos), quant_config=quant_config
+      model, (tokens, input_pos, kv), quant_config=quant_config
   )
   edge_model.export("/tmp/gemma_2b_quantized.tflite")
 

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241028.dist-info → ai_edge_torch_nightly-0.3.0.dev20241030.dist-info}/METADATA

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

@@ -3,7 +3,7 @@ ai_edge_torch/config.py,sha256=FMWeCH2b7HYILBvaI1iZNnYCO4WAhDOwBZBmIE-xrF0,909
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/fx_pass_base.py,sha256=SrYveglaiA_DXPoRBqSXClWM1q7853I5ujRorq_MV0M,4251
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=lZ9DKdslX4Io1UEJcL-gy2jhJ0jhO5yryL7wW_ibZG0,706
+ai_edge_torch/version.py,sha256=MlL0epetNoc10jxbXPvsj8gL4DSmIQTJq51OGNz2Qhc,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=HwzfRx_DX5TLtPqwEH1_NOm38_INvHzHl4_mX67KOdQ,5448
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -78,9 +78,9 @@ ai_edge_torch/generative/examples/smollm/verify.py,sha256=HXYcCjDJMylVL3Pc9HU-UX
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=kDWG6MlIGa89zC5KSRcJlw2c4ITuw8KcchtfmF55f4g,3545
 ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=5M4auM33SgCTODt0VT8TO-EVILruqGDRiNILBPeB83Y,6072
-ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=i9mcBITt4jJqKLA4Qdt3uFotCrglv14tPg8VnqsVnaI,5004
-ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=ClXNntmh0PF3s6U3C7SW3tyVrsSSrV2kyz-_RF4BcqA,15715
-ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=mBEAUYjV1qDJy9ZAsHtm9RGce0Mbzv0VoPZpdcQl1mk,33730
+ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=Fw0ZsJSG8fM-07mEi6QLCn6LpRveGGL8vt7OIn0Av5c,5276
+ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=sQKQ-k6H9kG2brgwLsktjCMeN2h0POyfMP6iNsPNKWc,16271
+ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=6W58LxmHHkz2ctgpknQkyoDANZAnE9Byp_svfqLpQf0,34793
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=CAPsW84A8f00nS6fLFeh_XUjCPsDCA5UxHOUsMrLfSU,3450
 ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=GnY3vPZ-obrWuJifuE5bUooKLqAI7v6q71oaTuLKeBE,8778
 ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py,sha256=xychak9hdLd6ieXBYEwrK2BkF8NRZWZSSCijIsESpBA,3420
@@ -117,9 +117,9 @@ ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=gXxh3papK
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=JwndhL3Z31TvkdGlAoTL5PQzmKfHdRWaaE1EbaMI4Gs,27540
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
-ai_edge_torch/generative/layers/unet/model_config.py,sha256=8ze9kVWMuyZVQcgK7hWYw9TM1W9lXD-2j0iMHlxoGX4,9267
+ai_edge_torch/generative/layers/unet/model_config.py,sha256=raYm8Ol-EFi0zs5vNqmj2ZJCFsnQW2TfwhgDcClfwFA,9356
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/quantize/example.py,sha256=tlACaRsz6lqOxakzpXVFJZYfFKOiFqetcYVJqWVRdPE,1542
+ai_edge_torch/generative/quantize/example.py,sha256=1lfVNUd2cEyRUnoZ7BLbRJ9IN-FTKiWBtZNPFUzAiWE,1747
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=n1Fm8BFC8gJa_oiwwAOOghJyHtOXYZ4q-5ZRy4pHrIw,1957
 ai_edge_torch/generative/quantize/quant_recipe.py,sha256=tKnuJq6hPD23JPCB9nPAlE1UHAwdbChkgPShiVaz4CE,5156
 ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=4fgmP_GgeiFUOkIaC9ZZXC12eO3DQZdrWDXRz5YXiwU,2270
@@ -186,8 +186,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.dev20241028.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241028.dist-info/METADATA,sha256=Dl4KrN2CLjNvdT1UsCCcLioNuAZbEwT_4x_pU5myJrw,1897
-ai_edge_torch_nightly-0.3.0.dev20241028.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20241028.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241028.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241030.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241030.dist-info/METADATA,sha256=OmF5Xcg12MFdHCNcW_0A4l8s7mU0yUXSsjkGRHiRnBI,1897
+ai_edge_torch_nightly-0.3.0.dev20241030.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20241030.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241030.dist-info/RECORD,,