ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl

ai_edge_torch/generative/examples/stable_diffusion/clip.py

@@ -0,0 +1,115 @@
+# 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.
+# ==============================================================================
+
+import torch
+from torch import nn
+
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attention_utils
+import ai_edge_torch.generative.layers.builder as builder
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.utilities.loader as loading_utils
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.mlp.fc1",
+    ff_down_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.mlp.fc2",
+    attn_query_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.q_proj",
+    attn_key_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.k_proj",
+    attn_value_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.v_proj",
+    attn_output_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.out_proj",
+    pre_attn_norm="cond_stage_model.transformer.text_model.encoder.layers.{}.layer_norm1",
+    pre_ff_norm="cond_stage_model.transformer.text_model.encoder.layers.{}.layer_norm2",
+    embedding="cond_stage_model.transformer.text_model.embeddings.token_embedding",
+    embedding_position="cond_stage_model.transformer.text_model.embeddings.position_embedding.weight",
+    final_norm="cond_stage_model.transformer.text_model.final_layer_norm",
+    lm_head=None,
+)
+
+
+class CLIP(nn.Module):
+  """CLIP text encoder
+  For details, see https://arxiv.org/abs/2103.00020
+  """
+
+  def __init__(self, config: cfg.ModelConfig):
+    super().__init__()
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.tok_embedding_position = nn.Parameter(
+        torch.zeros((config.max_seq_len, config.embedding_dim))
+    )
+
+    self.config = config
+    self.transformer_blocks = nn.ModuleList(
+        TransformerBlock(config) for _ in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(config.embedding_dim, config.final_norm_config)
+
+    self.mask_cache = attention_utils.build_causal_mask_cache(
+        size=config.max_seq_len, dtype=torch.float32
+    )
+
+  @torch.inference_mode
+  def forward(self, tokens: torch.LongTensor) -> torch.FloatTensor:
+    tokens = tokens.type(torch.long)
+
+    state = self.tok_embedding(tokens) + self.tok_embedding_position
+    for layer in self.transformer_blocks:
+      state = layer(state, mask=self.mask_cache)
+    output = self.final_norm(state)
+    return output
+
+
+def get_model_config() -> cfg.ModelConfig:
+  max_seq_len = 77
+  vocab_size = 49408
+  num_layers = 12
+  num_heads = 12
+  num_query_groups = 12
+  embedding_dim = 768
+
+  attn_config = cfg.AttentionConfig(
+      num_heads=num_heads,
+      num_query_groups=num_query_groups,
+      rotary_percentage=0.0,
+      qkv_use_bias=True,
+      qkv_transpose_before_split=True,
+      qkv_fused_interleaved=False,
+      output_proj_use_bias=True,
+      enable_kv_cache=False,
+  )
+
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.SEQUENTIAL,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_QUICK),
+      intermediate_size=embedding_dim * 4,
+      use_bias=True,
+  )
+
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
+
+  config = cfg.ModelConfig(
+      vocab_size=vocab_size,
+      num_layers=num_layers,
+      max_seq_len=max_seq_len,
+      embedding_dim=embedding_dim,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      pre_ff_norm_config=norm_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+
+  return config

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -0,0 +1,142 @@
+# 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.
+# ==============================================================================
+
+import argparse
+import os
+from pathlib import Path
+from typing import Optional
+
+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
+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.stable_diffusion_loader as stable_diffusion_loader
+
+arg_parser = argparse.ArgumentParser()
+arg_parser.add_argument(
+    '--clip_ckpt', type=str, help='Path to source CLIP model checkpoint', required=True
+)
+arg_parser.add_argument(
+    '--diffusion_ckpt',
+    type=str,
+    help='Path to source diffusion model checkpoint',
+    required=True,
+)
+arg_parser.add_argument(
+    '--decoder_ckpt',
+    type=str,
+    help='Path to source image decoder model checkpoint',
+    required=True,
+)
+arg_parser.add_argument(
+    '--output_dir',
+    type=str,
+    help='Path to the converted TF Lite directory.',
+    required=True,
+)
+
+
+@torch.inference_mode
+def convert_stable_diffusion_to_tflite(
+    output_dir: str,
+    clip_ckpt_path: str,
+    diffusion_ckpt_path: str,
+    decoder_ckpt_path: str,
+    image_height: int = 512,
+    image_width: int = 512,
+):
+
+  clip_model = clip.CLIP(clip.get_model_config())
+  loader = stable_diffusion_loader.ClipModelLoader(
+      clip_ckpt_path,
+      clip.TENSOR_NAMES,
+  )
+  loader.load(clip_model, strict=False)
+
+  diffusion_model = diffusion.Diffusion(diffusion.get_model_config(2))
+  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_loader = stable_diffusion_loader.AutoEncoderModelLoader(
+      decoder_ckpt_path, decoder.TENSOR_NAMES
+  )
+  decoder_loader.load(decoder_model, strict=False)
+
+  # TODO(yichunk): enable image encoder conversion
+  # if encoder_ckpt_path is not None:
+  #   encoder = Encoder()
+  #   encoder.load_state_dict(torch.load(encoder_ckpt_path))
+
+  # Tensors used to trace the model graph during conversion.
+  n_tokens = 77
+  timestamp = 0
+  len_prompt = 1
+  prompt_tokens = torch.full((1, n_tokens), 0, dtype=torch.long)
+  input_image = torch.full((1, 3, image_height, image_width), 0, dtype=torch.float32)
+  noise = torch.full(
+      (len_prompt, 4, image_height // 8, image_width // 8), 0, dtype=torch.float32
+  )
+
+  input_latents = torch.zeros_like(noise)
+  context_cond = clip_model(prompt_tokens)
+  context_uncond = torch.zeros_like(context_cond)
+  context = torch.cat([context_cond, context_uncond], axis=0)
+  time_embedding = util.get_time_embedding(timestamp)
+
+  if not os.path.exists(output_dir):
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+  # 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'
+  )
+
+  # TODO(yichunk): enable image encoder conversion
+  # Image encoder
+  # ai_edge_torch.signature('encode', encoder, (input_image, noise)).convert().export(
+  #     f'{output_dir}/encoder.tflite'
+  # )
+
+  # Diffusion
+  ai_edge_torch.signature(
+      'diffusion',
+      diffusion_model,
+      (torch.repeat_interleave(input_latents, 2, 0), context, time_embedding),
+  ).convert().export(f'{output_dir}/diffusion.tflite')
+
+  # Image decoder
+  ai_edge_torch.signature('decode', decoder_model, (input_latents,)).convert().export(
+      f'{output_dir}/decoder.tflite'
+  )
+
+
+if __name__ == '__main__':
+  args = arg_parser.parse_args()
+  convert_stable_diffusion_to_tflite(
+      output_dir=args.output_dir,
+      clip_ckpt_path=args.clip_ckpt,
+      diffusion_ckpt_path=args.diffusion_ckpt,
+      decoder_ckpt_path=args.decoder_ckpt,
+      image_height=512,
+      image_width=512,
+  )

ai_edge_torch/generative/examples/stable_diffusion/decoder.py

@@ -0,0 +1,317 @@
+# 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.
+# ==============================================================================
+
+import torch
+from torch import nn
+
+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.stable_diffusion_loader as stable_diffusion_loader
+
+TENSOR_NAMES = stable_diffusion_loader.AutoEncoderModelLoader.TensorNames(
+    post_quant_conv="first_stage_model.post_quant_conv",
+    conv_in="first_stage_model.decoder.conv_in",
+    mid_block_tensor_names=stable_diffusion_loader.MidBlockTensorNames(
+        residual_block_tensor_names=[
+            stable_diffusion_loader.ResidualBlockTensorNames(
+                norm_1="first_stage_model.decoder.mid.block_1.norm1",
+                norm_2="first_stage_model.decoder.mid.block_1.norm2",
+                conv_1="first_stage_model.decoder.mid.block_1.conv1",
+                conv_2="first_stage_model.decoder.mid.block_1.conv2",
+            ),
+            stable_diffusion_loader.ResidualBlockTensorNames(
+                norm_1="first_stage_model.decoder.mid.block_2.norm1",
+                norm_2="first_stage_model.decoder.mid.block_2.norm2",
+                conv_1="first_stage_model.decoder.mid.block_2.conv1",
+                conv_2="first_stage_model.decoder.mid.block_2.conv2",
+            ),
+        ],
+        attention_block_tensor_names=[
+            stable_diffusion_loader.AttentionBlockTensorNames(
+                norm="first_stage_model.decoder.mid.attn_1.norm",
+                q_proj="first_stage_model.decoder.mid.attn_1.q",
+                k_proj="first_stage_model.decoder.mid.attn_1.k",
+                v_proj="first_stage_model.decoder.mid.attn_1.v",
+                output_proj="first_stage_model.decoder.mid.attn_1.proj_out",
+            )
+        ],
+    ),
+    up_decoder_blocks_tensor_names=[
+        stable_diffusion_loader.UpDecoderBlockTensorNames(
+            residual_block_tensor_names=[
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.3.block.0.norm1",
+                    norm_2="first_stage_model.decoder.up.3.block.0.norm2",
+                    conv_1="first_stage_model.decoder.up.3.block.0.conv1",
+                    conv_2="first_stage_model.decoder.up.3.block.0.conv2",
+                ),
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.3.block.1.norm1",
+                    norm_2="first_stage_model.decoder.up.3.block.1.norm2",
+                    conv_1="first_stage_model.decoder.up.3.block.1.conv1",
+                    conv_2="first_stage_model.decoder.up.3.block.1.conv2",
+                ),
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.3.block.2.norm1",
+                    norm_2="first_stage_model.decoder.up.3.block.2.norm2",
+                    conv_1="first_stage_model.decoder.up.3.block.2.conv1",
+                    conv_2="first_stage_model.decoder.up.3.block.2.conv2",
+                ),
+            ],
+            upsample_conv="first_stage_model.decoder.up.3.upsample.conv",
+        ),
+        stable_diffusion_loader.UpDecoderBlockTensorNames(
+            residual_block_tensor_names=[
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.2.block.0.norm1",
+                    norm_2="first_stage_model.decoder.up.2.block.0.norm2",
+                    conv_1="first_stage_model.decoder.up.2.block.0.conv1",
+                    conv_2="first_stage_model.decoder.up.2.block.0.conv2",
+                ),
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.2.block.1.norm1",
+                    norm_2="first_stage_model.decoder.up.2.block.1.norm2",
+                    conv_1="first_stage_model.decoder.up.2.block.1.conv1",
+                    conv_2="first_stage_model.decoder.up.2.block.1.conv2",
+                ),
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.2.block.2.norm1",
+                    norm_2="first_stage_model.decoder.up.2.block.2.norm2",
+                    conv_1="first_stage_model.decoder.up.2.block.2.conv1",
+                    conv_2="first_stage_model.decoder.up.2.block.2.conv2",
+                ),
+            ],
+            upsample_conv="first_stage_model.decoder.up.2.upsample.conv",
+        ),
+        stable_diffusion_loader.UpDecoderBlockTensorNames(
+            residual_block_tensor_names=[
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.1.block.0.norm1",
+                    norm_2="first_stage_model.decoder.up.1.block.0.norm2",
+                    conv_1="first_stage_model.decoder.up.1.block.0.conv1",
+                    conv_2="first_stage_model.decoder.up.1.block.0.conv2",
+                    residual_layer="first_stage_model.decoder.up.1.block.0.nin_shortcut",
+                ),
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.1.block.1.norm1",
+                    norm_2="first_stage_model.decoder.up.1.block.1.norm2",
+                    conv_1="first_stage_model.decoder.up.1.block.1.conv1",
+                    conv_2="first_stage_model.decoder.up.1.block.1.conv2",
+                ),
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.1.block.2.norm1",
+                    norm_2="first_stage_model.decoder.up.1.block.2.norm2",
+                    conv_1="first_stage_model.decoder.up.1.block.2.conv1",
+                    conv_2="first_stage_model.decoder.up.1.block.2.conv2",
+                ),
+            ],
+            upsample_conv="first_stage_model.decoder.up.1.upsample.conv",
+        ),
+        stable_diffusion_loader.UpDecoderBlockTensorNames(
+            residual_block_tensor_names=[
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.0.block.0.norm1",
+                    norm_2="first_stage_model.decoder.up.0.block.0.norm2",
+                    conv_1="first_stage_model.decoder.up.0.block.0.conv1",
+                    conv_2="first_stage_model.decoder.up.0.block.0.conv2",
+                    residual_layer="first_stage_model.decoder.up.0.block.0.nin_shortcut",
+                ),
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.0.block.1.norm1",
+                    norm_2="first_stage_model.decoder.up.0.block.1.norm2",
+                    conv_1="first_stage_model.decoder.up.0.block.1.conv1",
+                    conv_2="first_stage_model.decoder.up.0.block.1.conv2",
+                ),
+                stable_diffusion_loader.ResidualBlockTensorNames(
+                    norm_1="first_stage_model.decoder.up.0.block.2.norm1",
+                    norm_2="first_stage_model.decoder.up.0.block.2.norm2",
+                    conv_1="first_stage_model.decoder.up.0.block.2.conv1",
+                    conv_2="first_stage_model.decoder.up.0.block.2.conv2",
+                ),
+            ],
+        ),
+    ],
+    final_norm="first_stage_model.decoder.norm_out",
+    conv_out="first_stage_model.decoder.conv_out",
+)
+
+
+class Decoder(nn.Module):
+  """The Decoder model used in Stable Diffusion.
+
+  For details, see https://arxiv.org/abs/2103.00020
+
+  Sturcture of the Decoder:
+
+      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_config=config.activation_config,
+                  num_layers=config.layers_per_block,
+                  add_upsample=not_final_block,
+                  upsample_conv=True,
+                  sampling_config=unet_cfg.UpSamplingConfig(
+                      mode=unet_cfg.SamplingType.NEAREST, scale_factor=2
+                  ),
+              )
+          )
+      )
+    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_config)
+    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:
+    """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)
+    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
+
+
+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
+
+  norm_config = layers_cfg.NormalizationConfig(
+      layers_cfg.NormalizationType.GROUP_NORM, group_num=32
+  )
+
+  att_config = unet_cfg.AttentionBlock2DConfig(
+      dim=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,
+          qkv_fused_interleaved=False,
+          rotary_percentage=0.0,
+      ),
+  )
+
+  mid_block_config = unet_cfg.MidBlock2DConfig(
+      in_channels=block_out_channels[-1],
+      normalization_config=norm_config,
+      activation_config=layers_cfg.ActivationConfig(layers_cfg.ActivationType.SILU),
+      num_layers=1,
+      attention_block_config=att_config,
+  )
+
+  config = unet_cfg.AutoEncoderConfig(
+      in_channels=in_channels,
+      latent_channels=latent_channels,
+      out_channels=out_channels,
+      activation_config=layers_cfg.ActivationConfig(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