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

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -18,8 +18,6 @@ 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
@@ -28,10 +26,14 @@ 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
+import torch
 
 arg_parser = argparse.ArgumentParser()
 arg_parser.add_argument(
-    '--clip_ckpt', type=str, help='Path to source CLIP model checkpoint', required=True
+    '--clip_ckpt',
+    type=str,
+    help='Path to source CLIP model checkpoint',
+    required=True,
 )
 arg_parser.add_argument(
     '--diffusion_ckpt',
@@ -93,9 +95,13 @@ def convert_stable_diffusion_to_tflite(
   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)
+  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
+      (len_prompt, 4, image_height // 8, image_width // 8),
+      0,
+      dtype=torch.float32,
   )
 
   input_latents = torch.zeros_like(noise)
@@ -107,7 +113,9 @@ 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
+  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/generative/examples/stable_diffusion/decoder.py

@@ -13,14 +13,13 @@
 # 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
+import torch
+from torch import nn
 
 TENSOR_NAMES = stable_diffusion_loader.AutoEncoderModelLoader.TensorNames(
     post_quant_conv="first_stage_model.post_quant_conv",
@@ -104,7 +103,9 @@ TENSOR_NAMES = stable_diffusion_loader.AutoEncoderModelLoader.TensorNames(
                     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",
+                    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",
@@ -128,7 +129,9 @@ TENSOR_NAMES = stable_diffusion_loader.AutoEncoderModelLoader.TensorNames(
                     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",
+                    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",
@@ -299,7 +302,9 @@ def get_model_config() -> unet_cfg.AutoEncoderConfig:
   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),
+      activation_config=layers_cfg.ActivationConfig(
+          layers_cfg.ActivationType.SILU
+      ),
       num_layers=1,
       attention_block_config=att_config,
   )
@@ -308,7 +313,9 @@ def get_model_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),
+      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,

ai_edge_torch/generative/examples/stable_diffusion/diffusion.py

@@ -13,14 +13,13 @@
 # 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
+import torch
+from torch import nn
 
 _down_encoder_blocks_tensor_names = [
     stable_diffusion_loader.DownEncoderBlockTensorNames(
@@ -39,9 +38,15 @@ _down_encoder_blocks_tensor_names = [
         ],
         transformer_block_tensor_names=[
             stable_diffusion_loader.TransformerBlockTensorNames(
-                pre_conv_norm=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.norm",
-                conv_in=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.proj_in",
-                conv_out=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.proj_out",
+                pre_conv_norm=(
+                    f"model.diffusion_model.input_blocks.{i*3+j+1}.1.norm"
+                ),
+                conv_in=(
+                    f"model.diffusion_model.input_blocks.{i*3+j+1}.1.proj_in"
+                ),
+                conv_out=(
+                    f"model.diffusion_model.input_blocks.{i*3+j+1}.1.proj_out"
+                ),
                 self_attention=stable_diffusion_loader.AttentionBlockTensorNames(
                     norm=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.norm1",
                     q_proj=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.attn1.to_q",
@@ -80,7 +85,9 @@ _mid_block_tensor_names = stable_diffusion_loader.MidBlockTensorNames(
             conv_1=f"model.diffusion_model.middle_block.{i}.in_layers.2",
             norm_2=f"model.diffusion_model.middle_block.{i}.out_layers.0",
             conv_2=f"model.diffusion_model.middle_block.{i}.out_layers.3",
-            time_embedding=f"model.diffusion_model.middle_block.{i}.emb_layers.1",
+            time_embedding=(
+                f"model.diffusion_model.middle_block.{i}.emb_layers.1"
+            ),
         )
         for i in [0, 2]
     ],
@@ -117,8 +124,12 @@ _up_decoder_blocks_tensor_names = [
     stable_diffusion_loader.SkipUpDecoderBlockTensorNames(
         residual_block_tensor_names=[
             stable_diffusion_loader.ResidualBlockTensorNames(
-                norm_1=f"model.diffusion_model.output_blocks.{i*3+j}.0.in_layers.0",
-                conv_1=f"model.diffusion_model.output_blocks.{i*3+j}.0.in_layers.2",
+                norm_1=(
+                    f"model.diffusion_model.output_blocks.{i*3+j}.0.in_layers.0"
+                ),
+                conv_1=(
+                    f"model.diffusion_model.output_blocks.{i*3+j}.0.in_layers.2"
+                ),
                 norm_2=f"model.diffusion_model.output_blocks.{i*3+j}.0.out_layers.0",
                 conv_2=f"model.diffusion_model.output_blocks.{i*3+j}.0.out_layers.3",
                 time_embedding=f"model.diffusion_model.output_blocks.{i*3+j}.0.emb_layers.1",
@@ -128,9 +139,15 @@ _up_decoder_blocks_tensor_names = [
         ],
         transformer_block_tensor_names=[
             stable_diffusion_loader.TransformerBlockTensorNames(
-                pre_conv_norm=f"model.diffusion_model.output_blocks.{i*3+j}.1.norm",
-                conv_in=f"model.diffusion_model.output_blocks.{i*3+j}.1.proj_in",
-                conv_out=f"model.diffusion_model.output_blocks.{i*3+j}.1.proj_out",
+                pre_conv_norm=(
+                    f"model.diffusion_model.output_blocks.{i*3+j}.1.norm"
+                ),
+                conv_in=(
+                    f"model.diffusion_model.output_blocks.{i*3+j}.1.proj_in"
+                ),
+                conv_out=(
+                    f"model.diffusion_model.output_blocks.{i*3+j}.1.proj_out"
+                ),
                 self_attention=stable_diffusion_loader.AttentionBlockTensorNames(
                     norm=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.norm1",
                     q_proj=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.attn1.to_q",
@@ -157,7 +174,9 @@ _up_decoder_blocks_tensor_names = [
         else None,
         upsample_conv=f"model.diffusion_model.output_blocks.{i*3+2}.2.conv"
         if 0 < i < 3
-        else (f"model.diffusion_model.output_blocks.2.1.conv" if i == 0 else None),
+        else (
+            f"model.diffusion_model.output_blocks.2.1.conv" if i == 0 else None
+        ),
     )
     for i in range(4)
 ]
@@ -475,7 +494,10 @@ class Diffusion(nn.Module):
         layers_cfg.ActivationConfig(config.final_activation_type)
     )
     self.conv_out = nn.Conv2d(
-        reversed_block_out_channels[-1], config.out_channels, kernel_size=3, padding=1
+        reversed_block_out_channels[-1],
+        config.out_channels,
+        kernel_size=3,
+        padding=1,
     )
 
   @torch.inference_mode
@@ -496,12 +518,15 @@ class Diffusion(nn.Module):
     x = self.conv_in(latents)
     skip_connection_tensors = [x]
     for encoder in self.down_encoders:
-      x, hidden_states = encoder(x, time_emb, context, output_hidden_states=True)
+      x, hidden_states = encoder(
+          x, time_emb, context, output_hidden_states=True
+      )
       skip_connection_tensors.extend(hidden_states)
     x = self.mid_block(x, time_emb, context)
     for decoder in self.up_decoders:
       encoder_tensors = [
-          skip_connection_tensors.pop() for i in range(self.config.layers_per_block + 1)
+          skip_connection_tensors.pop()
+          for i in range(self.config.layers_per_block + 1)
       ]
       x = decoder(x, encoder_tensors, time_emb, context)
     x = self.final_norm(x)

ai_edge_torch/generative/examples/stable_diffusion/encoder.py

@@ -13,12 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
+from ai_edge_torch.generative.examples.stable_diffusion.attention import SelfAttention  # NOQA
 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
-
 
 class AttentionBlock(nn.Module):
 
@@ -50,7 +49,9 @@ class ResidualBlock(nn.Module):
     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)
+    self.conv_2 = nn.Conv2d(
+        out_channels, out_channels, kernel_size=3, padding=1
+    )
 
     if in_channels == out_channels:
       self.residual_layer = nn.Identity()

ai_edge_torch/generative/examples/stable_diffusion/pipeline.py

@@ -18,30 +18,41 @@ import os
 from pathlib import Path
 from typing import Dict, Optional
 
-import numpy as np
-from PIL import Image
-from tqdm import tqdm
-
 import ai_edge_torch.generative.examples.stable_diffusion.samplers as samplers
 from ai_edge_torch.generative.examples.stable_diffusion.tokenizer import Tokenizer  # NOQA
 import ai_edge_torch.generative.examples.stable_diffusion.util as util
 from ai_edge_torch.model import TfLiteModel
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
 
 arg_parser = argparse.ArgumentParser()
 arg_parser.add_argument(
     '--tokenizer_vocab_dir',
     type=str,
-    help='Directory to the tokenizer vocabulary files, which include `merges.txt` and `vocab.json`',
+    help=(
+        'Directory to the tokenizer vocabulary files, which include'
+        ' `merges.txt` and `vocab.json`'
+    ),
     required=True,
 )
 arg_parser.add_argument(
-    '--clip_ckpt', type=str, help='Path to CLIP TFLite tflite file', required=True
+    '--clip_ckpt',
+    type=str,
+    help='Path to CLIP TFLite tflite file',
+    required=True,
 )
 arg_parser.add_argument(
-    '--diffusion_ckpt', type=str, help='Path to diffusion tflite file', required=True
+    '--diffusion_ckpt',
+    type=str,
+    help='Path to diffusion tflite file',
+    required=True,
 )
 arg_parser.add_argument(
-    '--decoder_ckpt', type=str, help='Path to decoder tflite file', required=True
+    '--decoder_ckpt',
+    type=str,
+    help='Path to decoder tflite file',
+    required=True,
 )
 arg_parser.add_argument(
     '--output_path',
@@ -56,20 +67,29 @@ arg_parser.add_argument(
     help='The prompt to guide the image generation.',
 )
 arg_parser.add_argument(
-    '--n_inference_steps', default=20, type=int, help='The number of denoising steps.'
+    '--n_inference_steps',
+    default=20,
+    type=int,
+    help='The number of denoising steps.',
 )
 arg_parser.add_argument(
     '--sampler',
     default='k_euler',
     type=str,
     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`.',
+    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.',
+    help=(
+        'A seed to make generation deterministic. A random number is used if'
+        ' unspecified.'
+    ),
 )
 
 
@@ -154,7 +174,9 @@ def run_tflite_pipeline(
   elif sampler == 'k_euler':
     sampler = samplers.KEulerSampler(n_inference_steps=n_inference_steps)
   elif sampler == 'k_euler_ancestral':
-    sampler = samplers.KEulerAncestralSampler(n_inference_steps=n_inference_steps)
+    sampler = samplers.KEulerAncestralSampler(
+        n_inference_steps=n_inference_steps
+    )
   else:
     raise ValueError(
         'Unknown sampler value %s. '
@@ -173,7 +195,8 @@ def run_tflite_pipeline(
   if input_image:
     if not hasattr(model, 'encoder'):
       raise AttributeError(
-          'Stable Diffusion must be initialized with encoder to accept input_image.'
+          'Stable Diffusion must be initialized with encoder to accept'
+          ' input_image.'
       )
     input_image = input_image.resize((width, height))
     input_image_np = np.array(input_image).astype(np.float32)

ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 
-import numpy as np
-
 from ai_edge_torch.generative.examples.stable_diffusion import util
 from ai_edge_torch.generative.examples.stable_diffusion.samplers.sampler import SamplerInterface  # NOQA
+import numpy as np
 
 
 class KEulerSampler(SamplerInterface):
@@ -46,7 +45,9 @@ class KEulerSampler(SamplerInterface):
 
   def set_strength(self, strength=1):
     start_step = self.n_inference_steps - int(self.n_inference_steps * strength)
-    self.timesteps = np.linspace(self.n_training_steps - 1, 0, self.n_inference_steps)
+    self.timesteps = np.linspace(
+        self.n_training_steps - 1, 0, self.n_inference_steps
+    )
     self.timesteps = self.timesteps[start_step:]
     self.initial_scale = self.sigmas[start_step]
     self.step_count = start_step

ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler_ancestral.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 
-import numpy as np
-
 from ai_edge_torch.generative.examples.stable_diffusion import util
 from ai_edge_torch.generative.examples.stable_diffusion.samplers.sampler import SamplerInterface  # NOQA
+import numpy as np
 
 
 class KEulerAncestralSampler(SamplerInterface):
@@ -46,7 +45,9 @@ class KEulerAncestralSampler(SamplerInterface):
 
   def set_strength(self, strength=1):
     start_step = self.n_inference_steps - int(self.n_inference_steps * strength)
-    self.timesteps = np.linspace(self.n_training_steps - 1, 0, self.n_inference_steps)
+    self.timesteps = np.linspace(
+        self.n_training_steps - 1, 0, self.n_inference_steps
+    )
     self.timesteps = self.timesteps[start_step:]
     self.initial_scale = self.sigmas[start_step]
     self.step_count = start_step

ai_edge_torch/generative/examples/stable_diffusion/samplers/k_lms.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 
-import numpy as np
-
 from ai_edge_torch.generative.examples.stable_diffusion import util
 from ai_edge_torch.generative.examples.stable_diffusion.samplers.sampler import SamplerInterface  # NOQA
+import numpy as np
 
 
 class KLMSSampler(SamplerInterface):
@@ -48,7 +47,9 @@ class KLMSSampler(SamplerInterface):
 
   def set_strength(self, strength=1):
     start_step = self.n_inference_steps - int(self.n_inference_steps * strength)
-    self.timesteps = np.linspace(self.n_training_steps - 1, 0, self.n_inference_steps)
+    self.timesteps = np.linspace(
+        self.n_training_steps - 1, 0, self.n_inference_steps
+    )
     self.timesteps = self.timesteps[start_step:]
     self.initial_scale = self.sigmas[start_step]
     self.step_count = start_step

ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py

@@ -27,7 +27,10 @@ def create_bytes_table() -> dict:
   special_count = 0
   for byte in range(256):
     category = unicodedata.category(chr(byte))
-    if category[0] not in ['C', 'Z']:  # ith character is NOT control char or space
+    if category[0] not in [
+        'C',
+        'Z',
+    ]:  # ith character is NOT control char or space
       table[byte] = chr(byte)
     else:  # ith character IS control char or space
       table[byte] = chr(special_count + 256)

ai_edge_torch/generative/examples/stable_diffusion/util.py

@@ -20,14 +20,20 @@ import torch
 
 
 def get_time_embedding(timestep):
-  freqs = torch.pow(10000, -torch.arange(start=0, end=160, dtype=torch.float32) / 160)
+  freqs = torch.pow(
+      10000, -torch.arange(start=0, end=160, dtype=torch.float32) / 160
+  )
   x = torch.tensor([timestep], dtype=torch.float32)[:, None] * freqs[None]
   return torch.cat([torch.cos(x), torch.sin(x)], dim=-1)
 
 
-def get_alphas_cumprod(beta_start=0.00085, beta_end=0.0120, n_training_steps=1000):
+def get_alphas_cumprod(
+    beta_start=0.00085, beta_end=0.0120, n_training_steps=1000
+):
   betas = (
-      np.linspace(beta_start**0.5, beta_end**0.5, n_training_steps, dtype=np.float32)
+      np.linspace(
+          beta_start**0.5, beta_end**0.5, n_training_steps, dtype=np.float32
+      )
       ** 2
   )
   alphas = 1.0 - betas

ai_edge_torch/generative/examples/t5/convert_to_tflite.py

@@ -16,12 +16,11 @@
 import os
 from pathlib import Path
 
-import numpy as np
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.generative.examples.t5 import t5
 from ai_edge_torch.generative.quantize import quant_recipes
+import numpy as np
+import torch
 
 
 # TODO(haoliang): clean this up untile 2-sig model is validated e2e.
@@ -73,8 +72,12 @@ def convert_t5_to_tflite_multisig(checkpoint_path: str):
   embedding_layer = torch.nn.Embedding(
       config.vocab_size, config.embedding_dim, padding_idx=0
   )
-  t5_encoder_model = t5.build_t5_encoder_model(config, embedding_layer, checkpoint_path)
-  t5_decoder_model = t5.build_t5_decoder_model(config, embedding_layer, checkpoint_path)
+  t5_encoder_model = t5.build_t5_encoder_model(
+      config, embedding_layer, checkpoint_path
+  )
+  t5_decoder_model = t5.build_t5_decoder_model(
+      config, embedding_layer, checkpoint_path
+  )
 
   # encoder
   seq_len = 512