ai-edge-torch-nightly 0.2.0.dev20240730__py3-none-any.whl → 0.2.0.dev20240802__py3-none-any.whl

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -18,19 +18,22 @@ 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
+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',
@@ -60,6 +63,7 @@ def convert_stable_diffusion_to_tflite(
     decoder_ckpt_path: str,
     image_height: int = 512,
     image_width: int = 512,
+    quantize: bool = True,
 ):
 
   clip_model = clip.CLIP(clip.get_model_config())
@@ -91,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)
@@ -105,15 +113,19 @@ 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
+  )
+
   # 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'
-  )
+  ai_edge_torch.signature('encode', clip_model, (prompt_tokens,)).convert(
+      quant_config=quant_config
+  ).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(
+  # ai_edge_torch.signature('encode', encoder, (input_image, noise)).convert(quant_config=quant_config).export(
   #     f'{output_dir}/encoder.tflite'
   # )
 
@@ -122,12 +134,12 @@ def convert_stable_diffusion_to_tflite(
       'diffusion',
       diffusion_model,
       (torch.repeat_interleave(input_latents, 2, 0), context, time_embedding),
-  ).convert().export(f'{output_dir}/diffusion.tflite')
+  ).convert(quant_config=quant_config).export(f'{output_dir}/diffusion.tflite')
 
   # Image decoder
-  ai_edge_torch.signature('decode', decoder_model, (input_latents,)).convert().export(
-      f'{output_dir}/decoder.tflite'
-  )
+  ai_edge_torch.signature('decode', decoder_model, (input_latents,)).convert(
+      quant_config=quant_config
+  ).export(f'{output_dir}/decoder.tflite')
 
 
 if __name__ == '__main__':
@@ -139,4 +151,5 @@ if __name__ == '__main__':
       decoder_ckpt_path=args.decoder_ckpt,
       image_height=512,
       image_width=512,
+      quantize=True,
   )

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",
@@ -293,12 +296,15 @@ def get_model_config() -> unet_cfg.AutoEncoderConfig:
           qkv_fused_interleaved=False,
           rotary_percentage=0.0,
       ),
+      enable_hlfb=False,
   )
 
   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,
   )
@@ -307,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)
 ]
@@ -294,6 +313,7 @@ class Diffusion(nn.Module):
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=attention_config,
+                            enable_hlfb=False,
                         ),
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
@@ -301,6 +321,7 @@ class Diffusion(nn.Module):
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=attention_config,
+                            enable_hlfb=False,
                         ),
                         pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                         feed_forward_block_config=unet_cfg.FeedForwardBlock2DConfig(
@@ -354,6 +375,7 @@ class Diffusion(nn.Module):
                     attention_batch_size=config.transformer_batch_size,
                     normalization_config=config.transformer_norm_config,
                     attention_config=attention_config,
+                    enable_hlfb=False,
                 ),
                 cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                     query_dim=mid_block_channels,
@@ -361,6 +383,7 @@ class Diffusion(nn.Module):
                     attention_batch_size=config.transformer_batch_size,
                     normalization_config=config.transformer_norm_config,
                     attention_config=attention_config,
+                    enable_hlfb=False,
                 ),
                 pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                 feed_forward_block_config=unet_cfg.FeedForwardBlock2DConfig(
@@ -415,6 +438,7 @@ class Diffusion(nn.Module):
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=attention_config,
+                            enable_hlfb=False,
                         ),
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
@@ -422,6 +446,7 @@ class Diffusion(nn.Module):
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
                             attention_config=attention_config,
+                            enable_hlfb=False,
                         ),
                         pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                         feed_forward_block_config=unet_cfg.FeedForwardBlock2DConfig(
@@ -469,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
@@ -490,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,14 +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.'
+    ),
 )
 
 
@@ -148,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. '
@@ -167,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)
@@ -219,4 +248,5 @@ if __name__ == '__main__':
       output_path=args.output_path,
       sampler=args.sampler,
       n_inference_steps=args.n_inference_steps,
+      seed=args.seed,
   )

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