ai-edge-torch-nightly 0.2.0.dev20240611__py3-none-any.whl → 0.2.0.dev20240617__py3-none-any.whl

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py

@@ -25,6 +25,25 @@ from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layo
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
 
 
+def can_partition(graph_module: torch.fx.GraphModule):
+  """Returns true if the input graph_module can be partitioned by min cut solver
+  in a reasonable time.
+
+  The min cut solver implements O(|V|^2|E|) Dinic's algorithm, which may
+  take a long time to complete for large graph module. This function determines
+  whether the graph module can be partitioned by the graph module size.
+
+  See go/pytorch-layout-transpose-optimization for more details.
+  """
+  graph = graph_module.graph
+  n_nodes = len(graph.nodes)
+  n_edges = sum(len(n.users) for n in graph.nodes)
+
+  # According to the experiments our model set, |V| < 2000 can
+  # be partitioned generally in a reasonable time.
+  return n_nodes**2 * n_edges < 2000**3
+
+
 class MinCutSolver:
   # A number that is large enough but can fit into int32 with all computations
   # in the maximum flow.

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -261,10 +261,17 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
     self.mark_const_nodes(exported_program)
 
     graph_module = exported_program.graph_module
-    if os.environ.get("AIEDGETORCH_LAYOUT_OPTIMIZE_USE_MINCUT_PARTITIONER"):
+    partitioner = os.environ.get("AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", None)
+    if partitioner == "MINCUT":
       graph_module = layout_partitioners.min_cut.partition(graph_module)
-    else:
+    elif partitioner == "GREEDY":
       graph_module = layout_partitioners.greedy.partition(graph_module)
+    else:
+      # By default use min cut partitioner if possible
+      if layout_partitioners.min_cut.can_partition(graph_module):
+        graph_module = layout_partitioners.min_cut.partition(graph_module)
+      else:
+        graph_module = layout_partitioners.greedy.partition(graph_module)
 
     graph = graph_module.graph
     for node in list(graph.nodes):

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -21,11 +21,11 @@ import torch
 import ai_edge_torch
 import ai_edge_torch.generative.examples.stable_diffusion.clip as clip
 import ai_edge_torch.generative.examples.stable_diffusion.decoder as decoder
-from ai_edge_torch.generative.examples.stable_diffusion.diffusion import Diffusion  # NOQA
+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
-import ai_edge_torch.generative.utilities.autoencoder_loader as autoencoder_loader
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import ai_edge_torch.generative.utilities.stable_diffusion_loader as stable_diffusion_loader
 
 
 @torch.inference_mode
@@ -45,11 +45,14 @@ def convert_stable_diffusion_to_tflite(
   encoder = Encoder()
   encoder.load_state_dict(torch.load(encoder_ckpt_path))
 
-  diffusion = Diffusion()
-  diffusion.load_state_dict(torch.load(diffusion_ckpt_path))
+  diffusion_model = diffusion.Diffusion(diffusion.get_model_config(2))
+  diffusion_loader = stable_diffusion_loader.DiffusionModelLoader(
+      diffusion_ckpt_path, diffusion.TENSORS_NAMES
+  )
+  diffusion_loader.load(diffusion_model)
 
   decoder_model = decoder.Decoder(decoder.get_model_config())
-  decoder_loader = autoencoder_loader.AutoEncoderModelLoader(
+  decoder_loader = stable_diffusion_loader.AutoEncoderModelLoader(
       decoder_ckpt_path, decoder.TENSORS_NAMES
   )
   decoder_loader.load(decoder_model)
@@ -84,7 +87,7 @@ def convert_stable_diffusion_to_tflite(
   # Diffusion
   ai_edge_torch.signature(
       'diffusion',
-      diffusion,
+      diffusion_model,
       (torch.repeat_interleave(input_latents, 2, 0), context, time_embedding),
   ).convert().export('/tmp/stable_diffusion/diffusion.tflite')
 

ai_edge_torch/generative/examples/stable_diffusion/decoder.py

@@ -20,20 +20,20 @@ import ai_edge_torch.generative.layers.builder as layers_builder
 import ai_edge_torch.generative.layers.model_config as layers_cfg
 import ai_edge_torch.generative.layers.unet.blocks_2d as blocks_2d
 import ai_edge_torch.generative.layers.unet.model_config as unet_cfg
-import ai_edge_torch.generative.utilities.autoencoder_loader as autoencoder_loader
+import ai_edge_torch.generative.utilities.stable_diffusion_loader as stable_diffusion_loader
 
-TENSORS_NAMES = autoencoder_loader.AutoEncoderModelLoader.TensorNames(
+TENSORS_NAMES = stable_diffusion_loader.AutoEncoderModelLoader.TensorNames(
     post_quant_conv="0",
     conv_in="1",
-    mid_block_tensor_names=autoencoder_loader.MidBlockTensorNames(
+    mid_block_tensor_names=stable_diffusion_loader.MidBlockTensorNames(
         residual_block_tensor_names=[
-            autoencoder_loader.ResidualBlockTensorNames(
+            stable_diffusion_loader.ResidualBlockTensorNames(
                 norm_1="2.groupnorm_1",
                 norm_2="2.groupnorm_2",
                 conv_1="2.conv_1",
                 conv_2="2.conv_2",
             ),
-            autoencoder_loader.ResidualBlockTensorNames(
+            stable_diffusion_loader.ResidualBlockTensorNames(
                 norm_1="4.groupnorm_1",
                 norm_2="4.groupnorm_2",
                 conv_1="4.conv_1",
@@ -41,7 +41,7 @@ TENSORS_NAMES = autoencoder_loader.AutoEncoderModelLoader.TensorNames(
             ),
         ],
         attention_block_tensor_names=[
-            autoencoder_loader.AttnetionBlockTensorNames(
+            stable_diffusion_loader.AttentionBlockTensorNames(
                 norm="3.groupnorm",
                 fused_qkv_proj="3.attention.in_proj",
                 output_proj="3.attention.out_proj",
@@ -49,21 +49,21 @@ TENSORS_NAMES = autoencoder_loader.AutoEncoderModelLoader.TensorNames(
         ],
     ),
     up_decoder_blocks_tensor_names=[
-        autoencoder_loader.UpDecoderBlockTensorNames(
+        stable_diffusion_loader.UpDecoderBlockTensorNames(
             residual_block_tensor_names=[
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="5.groupnorm_1",
                     norm_2="5.groupnorm_2",
                     conv_1="5.conv_1",
                     conv_2="5.conv_2",
                 ),
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="6.groupnorm_1",
                     norm_2="6.groupnorm_2",
                     conv_1="6.conv_1",
                     conv_2="6.conv_2",
                 ),
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="7.groupnorm_1",
                     norm_2="7.groupnorm_2",
                     conv_1="7.conv_1",
@@ -72,21 +72,21 @@ TENSORS_NAMES = autoencoder_loader.AutoEncoderModelLoader.TensorNames(
             ],
             upsample_conv="9",
         ),
-        autoencoder_loader.UpDecoderBlockTensorNames(
+        stable_diffusion_loader.UpDecoderBlockTensorNames(
             residual_block_tensor_names=[
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="10.groupnorm_1",
                     norm_2="10.groupnorm_2",
                     conv_1="10.conv_1",
                     conv_2="10.conv_2",
                 ),
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="11.groupnorm_1",
                     norm_2="11.groupnorm_2",
                     conv_1="11.conv_1",
                     conv_2="11.conv_2",
                 ),
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="12.groupnorm_1",
                     norm_2="12.groupnorm_2",
                     conv_1="12.conv_1",
@@ -95,22 +95,22 @@ TENSORS_NAMES = autoencoder_loader.AutoEncoderModelLoader.TensorNames(
             ],
             upsample_conv="14",
         ),
-        autoencoder_loader.UpDecoderBlockTensorNames(
+        stable_diffusion_loader.UpDecoderBlockTensorNames(
             residual_block_tensor_names=[
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="15.groupnorm_1",
                     norm_2="15.groupnorm_2",
                     conv_1="15.conv_1",
                     conv_2="15.conv_2",
                     residual_layer="15.residual_layer",
                 ),
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="16.groupnorm_1",
                     norm_2="16.groupnorm_2",
                     conv_1="16.conv_1",
                     conv_2="16.conv_2",
                 ),
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="17.groupnorm_1",
                     norm_2="17.groupnorm_2",
                     conv_1="17.conv_1",
@@ -119,22 +119,22 @@ TENSORS_NAMES = autoencoder_loader.AutoEncoderModelLoader.TensorNames(
             ],
             upsample_conv="19",
         ),
-        autoencoder_loader.UpDecoderBlockTensorNames(
+        stable_diffusion_loader.UpDecoderBlockTensorNames(
             residual_block_tensor_names=[
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="20.groupnorm_1",
                     norm_2="20.groupnorm_2",
                     conv_1="20.conv_1",
                     conv_2="20.conv_2",
                     residual_layer="20.residual_layer",
                 ),
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="21.groupnorm_1",
                     norm_2="21.groupnorm_2",
                     conv_1="21.conv_1",
                     conv_2="21.conv_2",
                 ),
-                autoencoder_loader.ResidualBlockTensorNames(
+                stable_diffusion_loader.ResidualBlockTensorNames(
                     norm_1="22.groupnorm_1",
                     norm_2="22.groupnorm_2",
                     conv_1="22.conv_1",
@@ -225,8 +225,8 @@ class Decoder(nn.Module):
                   num_layers=config.layers_per_block,
                   add_upsample=not_final_block,
                   upsample_conv=True,
-                  sampling_config=unet_cfg.SamplingConfig(
-                      2, unet_cfg.SamplingType.NEAREST
+                  sampling_config=unet_cfg.UpSamplingConfig(
+                      mode=unet_cfg.SamplingType.NEAREST, scale_factor=2
                   ),
               )
           )
@@ -245,6 +245,14 @@ class Decoder(nn.Module):
     )
 
   def forward(self, latents_tensor: torch.Tensor) -> torch.Tensor:
+    """Forward function of decoder model.
+
+    Args:
+        latents (torch.Tensor): latents space tensor.
+
+    Returns:
+        output decoded image tensor from decoder model.
+    """
     x = latents_tensor / self.config.scaling_factor
     x = self.post_quant_conv(x)
     x = self.conv_in(x)
@@ -271,7 +279,7 @@ def get_model_config() -> unet_cfg.AutoEncoderConfig:
   )
 
   att_config = unet_cfg.AttentionBlock2DConfig(
-      dims=block_out_channels[-1],
+      dim=block_out_channels[-1],
       normalization_config=norm_config,
       attention_config=layers_cfg.AttentionConfig(
           num_heads=1,