ai-edge-torch-nightly 0.2.0.dev20240606__py3-none-any.whl → 0.2.0.dev20240609__py3-none-any.whl

ai_edge_torch/convert/conversion.py

@@ -25,7 +25,7 @@ from torch_xla import stablehlo
 from ai_edge_torch import model
 from ai_edge_torch.convert import conversion_utils as cutils
 from ai_edge_torch.convert.fx_passes import BuildAtenCompositePass
-from ai_edge_torch.convert.fx_passes import BuildUpsampleBilinear2DCompositePass  # NOQA
+from ai_edge_torch.convert.fx_passes import BuildInterpolateCompositePass  # NOQA
 from ai_edge_torch.convert.fx_passes import CanonicalizePass
 from ai_edge_torch.convert.fx_passes import InjectMlirDebuginfoPass
 from ai_edge_torch.convert.fx_passes import OptimizeLayoutTransposesPass
@@ -41,7 +41,7 @@ def _run_convert_passes(
   return run_passes(
       exported_program,
       [
-          BuildUpsampleBilinear2DCompositePass(),
+          BuildInterpolateCompositePass(),
           CanonicalizePass(),
           OptimizeLayoutTransposesPass(),
           CanonicalizePass(),

ai_edge_torch/convert/fx_passes/__init__.py

@@ -24,7 +24,7 @@ from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult
 from ai_edge_torch.convert.fx_passes._pass_base import FxPassBase
 from ai_edge_torch.convert.fx_passes._pass_base import FxPassResult
 from ai_edge_torch.convert.fx_passes.build_aten_composite_pass import BuildAtenCompositePass  # NOQA
-from ai_edge_torch.convert.fx_passes.build_upsample_bilinear2d_composite_pass import BuildUpsampleBilinear2DCompositePass  # NOQA
+from ai_edge_torch.convert.fx_passes.build_interpolate_composite_pass import BuildInterpolateCompositePass  # NOQA
 from ai_edge_torch.convert.fx_passes.canonicalize_pass import CanonicalizePass
 from ai_edge_torch.convert.fx_passes.inject_mlir_debuginfo_pass import InjectMlirDebuginfoPass  # NOQA
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import OptimizeLayoutTransposesPass  # NOQA

ai_edge_torch/convert/fx_passes/{build_upsample_bilinear2d_composite_pass.py → build_interpolate_composite_pass.py}

@@ -66,13 +66,34 @@ def _get_upsample_bilinear2d_align_corners_pattern():
   return pattern
 
 
-class BuildUpsampleBilinear2DCompositePass(FxPassBase):
+@functools.cache
+def _get_interpolate_nearest2d_pattern():
+  pattern = mark_pattern.Pattern(
+      "tfl.resize_nearest_neighbor",
+      lambda x: torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest"),
+      export_args=(torch.rand(1, 3, 100, 100),),
+  )
+
+  @pattern.register_attr_builder
+  def attr_builder(pattern, graph_module, internal_match):
+    output = internal_match.returning_nodes[0]
+    output_h, output_w = output.meta["val"].shape[-2:]
+    return {
+        "size": (int(output_h), int(output_w)),
+        "is_nchw_op": True,
+    }
+
+  return pattern
+
+
+class BuildInterpolateCompositePass(FxPassBase):
 
   def __init__(self):
     super().__init__()
     self._patterns = [
         _get_upsample_bilinear2d_pattern(),
         _get_upsample_bilinear2d_align_corners_pattern(),
+        _get_interpolate_nearest2d_pattern(),
     ]
 
   def call(self, graph_module: torch.fx.GraphModule):

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -20,10 +20,11 @@ import torch
 
 import ai_edge_torch
 import ai_edge_torch.generative.examples.stable_diffusion.clip as clip
-from ai_edge_torch.generative.examples.stable_diffusion.decoder import Decoder
+import ai_edge_torch.generative.examples.stable_diffusion.decoder as decoder
 from ai_edge_torch.generative.examples.stable_diffusion.diffusion import Diffusion  # NOQA
 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
 
 
@@ -47,8 +48,11 @@ def convert_stable_diffusion_to_tflite(
   diffusion = Diffusion()
   diffusion.load_state_dict(torch.load(diffusion_ckpt_path))
 
-  decoder = Decoder()
-  decoder.load_state_dict(torch.load(decoder_ckpt_path))
+  decoder_model = decoder.Decoder(decoder.get_model_config())
+  decoder_loader = autoencoder_loader.AutoEncoderModelLoader(
+      decoder_ckpt_path, decoder.TENSORS_NAMES
+  )
+  decoder_loader.load(decoder_model)
 
   # Tensors used to trace the model graph during conversion.
   n_tokens = 77
@@ -85,7 +89,7 @@ def convert_stable_diffusion_to_tflite(
   ).convert().export('/tmp/stable_diffusion/diffusion.tflite')
 
   # Image decoder
-  ai_edge_torch.signature('decode', decoder, (input_latents,)).convert().export(
+  ai_edge_torch.signature('decode', decoder_model, (input_latents,)).convert().export(
       '/tmp/stable_diffusion/decoder.tflite'
   )
 

ai_edge_torch/generative/examples/stable_diffusion/decoder.py

@@ -15,99 +15,292 @@
 
 import torch
 from torch import nn
-from torch.nn import functional as F
 
-from ai_edge_torch.generative.examples.stable_diffusion.attention import SelfAttention  # NOQA
+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
 
+TENSORS_NAMES = autoencoder_loader.AutoEncoderModelLoader.TensorNames(
+    post_quant_conv="0",
+    conv_in="1",
+    mid_block_tensor_names=autoencoder_loader.MidBlockTensorNames(
+        residual_block_tensor_names=[
+            autoencoder_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(
+                norm_1="4.groupnorm_1",
+                norm_2="4.groupnorm_2",
+                conv_1="4.conv_1",
+                conv_2="4.conv_2",
+            ),
+        ],
+        attention_block_tensor_names=[
+            autoencoder_loader.AttnetionBlockTensorNames(
+                norm="3.groupnorm",
+                fused_qkv_proj="3.attention.in_proj",
+                output_proj="3.attention.out_proj",
+            )
+        ],
+    ),
+    up_decoder_blocks_tensor_names=[
+        autoencoder_loader.UpDecoderBlockTensorNames(
+            residual_block_tensor_names=[
+                autoencoder_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(
+                    norm_1="6.groupnorm_1",
+                    norm_2="6.groupnorm_2",
+                    conv_1="6.conv_1",
+                    conv_2="6.conv_2",
+                ),
+                autoencoder_loader.ResidualBlockTensorNames(
+                    norm_1="7.groupnorm_1",
+                    norm_2="7.groupnorm_2",
+                    conv_1="7.conv_1",
+                    conv_2="7.conv_2",
+                ),
+            ],
+            upsample_conv="9",
+        ),
+        autoencoder_loader.UpDecoderBlockTensorNames(
+            residual_block_tensor_names=[
+                autoencoder_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(
+                    norm_1="11.groupnorm_1",
+                    norm_2="11.groupnorm_2",
+                    conv_1="11.conv_1",
+                    conv_2="11.conv_2",
+                ),
+                autoencoder_loader.ResidualBlockTensorNames(
+                    norm_1="12.groupnorm_1",
+                    norm_2="12.groupnorm_2",
+                    conv_1="12.conv_1",
+                    conv_2="12.conv_2",
+                ),
+            ],
+            upsample_conv="14",
+        ),
+        autoencoder_loader.UpDecoderBlockTensorNames(
+            residual_block_tensor_names=[
+                autoencoder_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(
+                    norm_1="16.groupnorm_1",
+                    norm_2="16.groupnorm_2",
+                    conv_1="16.conv_1",
+                    conv_2="16.conv_2",
+                ),
+                autoencoder_loader.ResidualBlockTensorNames(
+                    norm_1="17.groupnorm_1",
+                    norm_2="17.groupnorm_2",
+                    conv_1="17.conv_1",
+                    conv_2="17.conv_2",
+                ),
+            ],
+            upsample_conv="19",
+        ),
+        autoencoder_loader.UpDecoderBlockTensorNames(
+            residual_block_tensor_names=[
+                autoencoder_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(
+                    norm_1="21.groupnorm_1",
+                    norm_2="21.groupnorm_2",
+                    conv_1="21.conv_1",
+                    conv_2="21.conv_2",
+                ),
+                autoencoder_loader.ResidualBlockTensorNames(
+                    norm_1="22.groupnorm_1",
+                    norm_2="22.groupnorm_2",
+                    conv_1="22.conv_1",
+                    conv_2="22.conv_2",
+                ),
+            ],
+        ),
+    ],
+    final_norm="23",
+    conv_out="25",
+)
 
-class AttentionBlock(nn.Module):
 
-  def __init__(self, channels):
-    super().__init__()
-    self.groupnorm = nn.GroupNorm(32, channels)
-    self.attention = SelfAttention(1, channels)
+class Decoder(nn.Module):
+  """The Decoder model used in Stable Diffusion.
 
-  def forward(self, x):
-    residue = x
-    x = self.groupnorm(x)
+  For details, see https://arxiv.org/abs/2103.00020
 
-    n, c, h, w = x.shape
-    x = x.view((n, c, h * w))
-    x = x.transpose(-1, -2)
-    x = self.attention(x)
-    x = x.transpose(-1, -2)
-    x = x.view((n, c, h, w))
+  Sturcture of the Decoder:
 
-    x += residue
-    return x
+      latents tensor
+            |
+            ▼
+  ┌───────────────────┐
+  │  Post Quant Conv  │
+  └─────────┬─────────┘
+            │
+  ┌─────────▼─────────┐
+  │      ConvIn       │
+  └─────────┬─────────┘
+            │
+  ┌─────────▼─────────┐
+  │     MidBlock2D    │
+  └─────────┬─────────┘
+            │
+  ┌─────────▼─────────┐
+  │    UpDecoder2D    │ x 4
+  └─────────┬─────────┘
+            │
+  ┌─────────▼─────────┐
+  │     FinalNorm     │
+  └─────────┬─────────┘
+            |
+  ┌─────────▼─────────┐
+  │    Activation     │
+  └─────────┬─────────┘
+            |
+  ┌─────────▼─────────┐
+  │      ConvOut      │
+  └─────────┬─────────┘
+            |
+            ▼
+      Output Image
+  """
+
+  def __init__(self, config: unet_cfg.AutoEncoderConfig):
+    super().__init__()
+    self.config = config
+    self.post_quant_conv = nn.Conv2d(
+        config.latent_channels,
+        config.latent_channels,
+        kernel_size=1,
+        stride=1,
+        padding=0,
+    )
+    reversed_block_out_channels = list(reversed(config.block_out_channels))
+    self.conv_in = nn.Conv2d(
+        config.latent_channels,
+        reversed_block_out_channels[0],
+        kernel_size=3,
+        stride=1,
+        padding=1,
+    )
+    self.mid_block = blocks_2d.MidBlock2D(config.mid_block_config)
+    up_decoder_blocks = []
+    block_out_channels = reversed_block_out_channels[0]
+    for i, out_channels in enumerate(reversed_block_out_channels):
+      prev_output_channel = block_out_channels
+      block_out_channels = out_channels
+      not_final_block = i < len(reversed_block_out_channels) - 1
+      up_decoder_blocks.append(
+          blocks_2d.UpDecoderBlock2D(
+              unet_cfg.UpDecoderBlock2DConfig(
+                  in_channels=prev_output_channel,
+                  out_channels=block_out_channels,
+                  normalization_config=config.normalization_config,
+                  activation_type=config.activation_type,
+                  num_layers=config.layers_per_block,
+                  add_upsample=not_final_block,
+                  upsample_conv=True,
+                  sampling_config=unet_cfg.SamplingConfig(
+                      2, unet_cfg.SamplingType.NEAREST
+                  ),
+              )
+          )
+      )
+    self.up_decoder_blocks = nn.ModuleList(up_decoder_blocks)
+    self.final_norm = layers_builder.build_norm(
+        block_out_channels, config.normalization_config
+    )
+    self.act_fn = layers_builder.get_activation(config.activation_type)
+    self.conv_out = nn.Conv2d(
+        block_out_channels,
+        config.out_channels,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+    )
 
+  def forward(self, latents_tensor: torch.Tensor) -> torch.Tensor:
+    x = latents_tensor / self.config.scaling_factor
+    x = self.post_quant_conv(x)
+    x = self.conv_in(x)
+    x = self.mid_block(x)
+    for up_decoder_block in self.up_decoder_blocks:
+      x = up_decoder_block(x)
+    x = self.final_norm(x)
+    x = self.act_fn(x)
+    x = self.conv_out(x)
+    return x
 
-class ResidualBlock(nn.Module):
 
-  def __init__(self, in_channels, out_channels):
-    super().__init__()
-    self.groupnorm_1 = nn.GroupNorm(32, in_channels)
-    self.conv_1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+def get_model_config() -> unet_cfg.AutoEncoderConfig:
+  """Get configs for the Decoder of Stable Diffusion v1.5"""
+  in_channels = 3
+  latent_channels = 4
+  out_channels = 3
+  block_out_channels = [128, 256, 512, 512]
+  scaling_factor = 0.18215
+  layers_per_block = 3
 
-    self.groupnorm_2 = nn.GroupNorm(32, out_channels)
-    self.conv_2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+  norm_config = layers_cfg.NormalizationConfig(
+      layers_cfg.NormalizationType.GROUP_NORM, group_num=32
+  )
 
-    if in_channels == out_channels:
-      self.residual_layer = nn.Identity()
-    else:
-      self.residual_layer = nn.Conv2d(
-          in_channels, out_channels, kernel_size=1, padding=0
-      )
+  att_config = unet_cfg.AttentionBlock2DConfig(
+      dims=block_out_channels[-1],
+      normalization_config=norm_config,
+      attention_config=layers_cfg.AttentionConfig(
+          num_heads=1,
+          num_query_groups=1,
+          qkv_use_bias=True,
+          output_proj_use_bias=True,
+          enable_kv_cache=False,
+          qkv_transpose_before_split=True,
+          rotary_percentage=0.0,
+      ),
+  )
 
-  def forward(self, x):
-    residue = x
-
-    x = self.groupnorm_1(x)
-    x = F.silu(x)
-    x = self.conv_1(x)
-
-    x = self.groupnorm_2(x)
-    x = F.silu(x)
-    x = self.conv_2(x)
-
-    return x + self.residual_layer(residue)
-
-
-class Decoder(nn.Sequential):
-
-  def __init__(self):
-    super().__init__(
-        nn.Conv2d(4, 4, kernel_size=1, padding=0),
-        nn.Conv2d(4, 512, kernel_size=3, padding=1),
-        ResidualBlock(512, 512),
-        AttentionBlock(512),
-        ResidualBlock(512, 512),
-        ResidualBlock(512, 512),
-        ResidualBlock(512, 512),
-        ResidualBlock(512, 512),
-        nn.Upsample(scale_factor=2),
-        nn.Conv2d(512, 512, kernel_size=3, padding=1),
-        ResidualBlock(512, 512),
-        ResidualBlock(512, 512),
-        ResidualBlock(512, 512),
-        nn.Upsample(scale_factor=2),
-        nn.Conv2d(512, 512, kernel_size=3, padding=1),
-        ResidualBlock(512, 256),
-        ResidualBlock(256, 256),
-        ResidualBlock(256, 256),
-        nn.Upsample(scale_factor=2),
-        nn.Conv2d(256, 256, kernel_size=3, padding=1),
-        ResidualBlock(256, 128),
-        ResidualBlock(128, 128),
-        ResidualBlock(128, 128),
-        nn.GroupNorm(32, 128),
-        nn.SiLU(),
-        nn.Conv2d(128, 3, kernel_size=3, padding=1),
-    )
+  mid_block_config = unet_cfg.MidBlock2DConfig(
+      in_channels=block_out_channels[-1],
+      normalization_config=norm_config,
+      activation_type=layers_cfg.ActivationType.SILU,
+      num_layers=1,
+      attention_block_config=att_config,
+  )
 
-  @torch.inference_mode
-  def forward(self, x):
-    x = x / 0.18215
-    for module in self:
-      x = module(x)
-    return x
+  config = unet_cfg.AutoEncoderConfig(
+      in_channels=in_channels,
+      latent_channels=latent_channels,
+      out_channels=out_channels,
+      activation_type=layers_cfg.ActivationType.SILU,
+      block_out_channels=block_out_channels,
+      scaling_factor=scaling_factor,
+      layers_per_block=layers_per_block,
+      normalization_config=norm_config,
+      mid_block_config=mid_block_config,
+  )
+  return config

ai_edge_torch/generative/examples/stable_diffusion/encoder.py

@@ -17,9 +17,60 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 
-from ai_edge_torch.generative.examples.stable_diffusion.decoder import AttentionBlock  # NOQA
-from ai_edge_torch.generative.examples.stable_diffusion.decoder import ResidualBlock  # NOQA
-import ai_edge_torch.generative.utilities.loader as loading_utils
+from ai_edge_torch.generative.examples.stable_diffusion.attention import SelfAttention  # NOQA
+
+
+class AttentionBlock(nn.Module):
+
+  def __init__(self, channels):
+    super().__init__()
+    self.groupnorm = nn.GroupNorm(32, channels)
+    self.attention = SelfAttention(1, channels)
+
+  def forward(self, x):
+    residue = x
+    x = self.groupnorm(x)
+
+    n, c, h, w = x.shape
+    x = x.view((n, c, h * w))
+    x = x.transpose(-1, -2)
+    x = self.attention(x)
+    x = x.transpose(-1, -2)
+    x = x.view((n, c, h, w))
+
+    x += residue
+    return x
+
+
+class ResidualBlock(nn.Module):
+
+  def __init__(self, in_channels, out_channels):
+    super().__init__()
+    self.groupnorm_1 = nn.GroupNorm(32, in_channels)
+    self.conv_1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+
+    self.groupnorm_2 = nn.GroupNorm(32, out_channels)
+    self.conv_2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+
+    if in_channels == out_channels:
+      self.residual_layer = nn.Identity()
+    else:
+      self.residual_layer = nn.Conv2d(
+          in_channels, out_channels, kernel_size=1, padding=0
+      )
+
+  def forward(self, x):
+    residue = x
+
+    x = self.groupnorm_1(x)
+    x = F.silu(x)
+    x = self.conv_1(x)
+
+    x = self.groupnorm_2(x)
+    x = F.silu(x)
+    x = self.conv_2(x)
+
+    return x + self.residual_layer(residue)
 
 
 class Encoder(nn.Sequential):

ai_edge_torch/generative/layers/attention.py

@@ -199,3 +199,28 @@ class CausalSelfAttention(nn.Module):
     # Compute the output projection.
     y = self.output_projection(y)
     return y
+
+
+class SelfAttention(CausalSelfAttention):
+  """Non-causal Self Attention module, which is equivalent to CausalSelfAttention without mask."""
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+      input_pos: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
+    """Forward function of the SelfAttention layer, which can support MQA, GQA and MHA.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+
+    Returns:
+      output activation from this self attention layer.
+    """
+    B, T, _ = x.size()
+    return super().forward(
+        x, rope=rope, mask=torch.zeros((B, T), dtype=torch.float32), input_pos=input_pos
+    )

ai_edge_torch/generative/layers/builder.py

@@ -44,6 +44,8 @@ def build_norm(dim: int, config: cfg.NormalizationConfig):
     )
   elif config.type == cfg.NormalizationType.LAYER_NORM:
     return nn.LayerNorm(dim, eps=config.epsilon)
+  elif config.type == cfg.NormalizationType.GROUP_NORM:
+    return nn.GroupNorm(config.group_num, dim, config.epsilon)
   else:
     raise ValueError("Unsupported norm type.")
 
@@ -69,7 +71,7 @@ def build_ff(dim: int, config: cfg.FeedForwardConfig):
   else:
     raise ValueError("Unsupported feedforward type.")
 
-  activation = _get_activation(config.activation)
+  activation = get_activation(config.activation)
 
   return ff_module(
       dim=dim,
@@ -79,7 +81,7 @@ def build_ff(dim: int, config: cfg.FeedForwardConfig):
   )
 
 
-def _get_activation(type_: cfg.ActivationType):
+def get_activation(type_: cfg.ActivationType):
   """Get pytorch callable activation from the name.
 
   Args:

ai_edge_torch/generative/layers/model_config.py

@@ -39,6 +39,7 @@ class NormalizationType(enum.Enum):
   NONE = enum.auto()
   RMS_NORM = enum.auto()
   LAYER_NORM = enum.auto()
+  GROUP_NORM = enum.auto()
 
 
 @enum.unique
@@ -90,6 +91,8 @@ class NormalizationConfig:
   type: NormalizationType = NormalizationType.NONE
   epsilon: float = 1e-5
   zero_centered: bool = False
+  # Number of groups used in group normalization.
+  group_num: Optional[float] = None
 
 
 @dataclass

ai_edge_torch/generative/layers/unet/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 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.
+# ==============================================================================