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

ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py

@@ -99,6 +99,36 @@ def _aten_hardswish(gm: GraphModule, node: Node):
   node.target = hardswish
 
 
+@_register_composite_builder(torch.ops.aten.gelu.default)
+def _aten_gelu(gm: GraphModule, node: Node):
+  op = node.target
+  args_mapper = TorchOpArgumentsMapper(op)
+
+  def gelu(*args, **kwargs):
+    nonlocal op, args_mapper
+
+    full_kwargs = args_mapper.get_full_kwargs(args, kwargs)
+
+    # TFLite supports exact and tanh approximate.
+    if full_kwargs["approximate"] != "none" and full_kwargs["approximate"] != "tanh":
+      return op(*args, **kwargs)
+
+    builder = StableHLOCompositeBuilder(
+        "aten.gelu.default",
+        attr=_tree_map_to_composite_attr_values(
+            {
+                "approximate": full_kwargs["approximate"],
+            }
+        ),
+    )
+    full_kwargs["self"] = builder.mark_inputs(full_kwargs["self"])
+    output = op(full_kwargs["self"])
+    output = builder.mark_outputs(output)
+    return output
+
+  node.target = gelu
+
+
 @_register_composite_builder(torch.ops.aten.avg_pool2d.default)
 def _aten_avg_pool2d(gm: GraphModule, node: Node):
   op = node.target

ai_edge_torch/convert/test/test_convert_composites.py

@@ -169,6 +169,24 @@ class TestConvertComposites(unittest.TestCase):
         model_coverage.compare_tflite_torch(edge_model, torch_module, tracing_args)
     )
 
+  def test_convert_gelu(self):
+    """Tests conversion of a GELU module."""
+
+    args = (torch.randn((5, 10)),)
+    torch_module = torch.nn.GELU().eval()
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+
+  def test_convert_gelu_approximate(self):
+    """Tests conversion of an Approximate GELU module."""
+
+    args = (torch.randn((5, 10)),)
+    torch_module = torch.nn.GELU('tanh').eval()
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(model_coverage.compare_tflite_torch(edge_model, torch_module, args))
+
 
 if __name__ == '__main__':
   unittest.main()

ai_edge_torch/generative/examples/stable_diffusion/clip.py

@@ -15,65 +15,99 @@
 
 import torch
 from torch import nn
-from torch._prims_common import mask_tensor
-from torch._prims_common.wrappers import out_wrapper
 
-from ai_edge_torch.generative.examples.stable_diffusion.attention import SelfAttention  # NOQA
+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="layers.{}.linear_1",
+    ff_down_proj="layers.{}.linear_2",
+    ff_gate_proj="layers.{}.linear_1",
+    attn_fused_qkv_proj="layers.{}.attention.in_proj",
+    attn_output_proj="layers.{}.attention.out_proj",
+    pre_attn_norm="layers.{}.layernorm_1",
+    pre_ff_norm="layers.{}.layernorm_2",
+    embedding="embedding.token_embedding",
+    embedding_position="embedding.position_value",
+    final_norm="layernorm",
+    lm_head=None,
+)
 
 
-class CLIPEmbedding(nn.Module):
-
-  def __init__(self, n_vocab: int, n_embd: int, n_token: int):
-    super().__init__()
-    self.token_embedding = nn.Embedding(n_vocab, n_embd)
-    self.position_value = nn.Parameter(torch.zeros((n_token, n_embd)))
-
-  def forward(self, tokens):
-    x = self.token_embedding(tokens)
-    x += self.position_value
-    return x
-
-
-class CLIPLayer(nn.Module):
+class CLIP(nn.Module):
+  """CLIP text encoder
+  For details, see https://arxiv.org/abs/2103.00020
+  """
 
-  def __init__(self, n_head: int, n_embd: int):
+  def __init__(self, config: cfg.ModelConfig):
     super().__init__()
-    self.layernorm_1 = nn.LayerNorm(n_embd)
-    self.attention = SelfAttention(n_head, n_embd)
-    self.layernorm_2 = nn.LayerNorm(n_embd)
-    self.linear_1 = nn.Linear(n_embd, 4 * n_embd)
-    self.linear_2 = nn.Linear(4 * n_embd, n_embd)
-
-  def forward(self, x):
-    residue = x
-    x = self.layernorm_1(x)
-    x = self.attention(x, causal_mask=True)
-    x += residue
+    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))
+    )
 
-    residue = x
-    x = self.layernorm_2(x)
-    x = self.linear_1(x)
-    x = x * torch.sigmoid(1.702 * x)  # QuickGELU activation function
-    x = self.linear_2(x)
-    x += residue
+    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)
 
-    return x
-
-
-class CLIP(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.embedding = CLIPEmbedding(49408, 768, 77)
-    self.layers = nn.ModuleList([CLIPLayer(12, 768) for i in range(12)])
-    self.layernorm = nn.LayerNorm(768)
+    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.embedding(tokens)
-    for layer in self.layers:
-      state = layer(state)
-    output = self.layernorm(state)
+    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,
+      output_proj_use_bias=True,
+      enable_kv_cache=False,
+  )
+
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.SEQUENTIAL,
+      activation=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

@@ -19,11 +19,12 @@ from pathlib import Path
 import torch
 
 import ai_edge_torch
-from ai_edge_torch.generative.examples.stable_diffusion.clip import CLIP
+import ai_edge_torch.generative.examples.stable_diffusion.clip as clip
 from ai_edge_torch.generative.examples.stable_diffusion.decoder import 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.loader as loading_utils
 
 
 @torch.inference_mode
@@ -36,8 +37,9 @@ def convert_stable_diffusion_to_tflite(
     image_width: int = 512,
 ):
 
-  clip = CLIP()
-  clip.load_state_dict(torch.load(clip_ckpt_path))
+  clip_model = clip.CLIP(clip.get_model_config())
+  loader = loading_utils.ModelLoader(clip_ckpt_path, clip.TENSOR_NAMES)
+  loader.load(clip_model, strict=False)
 
   encoder = Encoder()
   encoder.load_state_dict(torch.load(encoder_ckpt_path))
@@ -59,13 +61,13 @@ def convert_stable_diffusion_to_tflite(
   )
 
   input_latents = encoder(input_image, noise)
-  context_cond = clip(prompt_tokens)
+  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)
 
   # CLIP text encoder
-  ai_edge_torch.signature('encode', clip, (prompt_tokens,)).convert().export(
+  ai_edge_torch.signature('encode', clip_model, (prompt_tokens,)).convert().export(
       '/tmp/stable_diffusion/clip.tflite'
   )
 

ai_edge_torch/generative/examples/stable_diffusion/diffusion.py

@@ -202,11 +202,6 @@ class UNet(nn.Module):
 
     x = self.bottleneck(x, context, time)
 
-    # print('x shape:')
-    # print(list(x.shape))
-    # print('time shape:')
-    # print(list(time.shape))
-
     for layers in self.decoders:
       x = torch.cat((x, skip_connections.pop()), dim=1)
       x = layers(x, context, time)
@@ -214,199 +209,6 @@ class UNet(nn.Module):
     return x
 
 
-# The encoder component.
-class UNetEncoder(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.time_embedding = TimeEmbedding(320)
-    self.encoders = nn.ModuleList(
-        [
-            SwitchSequential(nn.Conv2d(4, 320, kernel_size=3, padding=1)),
-            SwitchSequential(ResidualBlock(320, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(320, 320), AttentionBlock(8, 40)),
-            SwitchSequential(nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)),
-            SwitchSequential(ResidualBlock(320, 640), AttentionBlock(8, 80)),
-            SwitchSequential(ResidualBlock(640, 640), AttentionBlock(8, 80)),
-            SwitchSequential(nn.Conv2d(640, 640, kernel_size=3, stride=2, padding=1)),
-            SwitchSequential(ResidualBlock(640, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(ResidualBlock(1280, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(nn.Conv2d(1280, 1280, kernel_size=3, stride=2, padding=1)),
-            SwitchSequential(ResidualBlock(1280, 1280)),
-            SwitchSequential(ResidualBlock(1280, 1280)),
-        ]
-    )
-
-  def forward(self, x, context, time):
-    time_embedding = self.time_embedding(time)
-    skip_connections = []
-    for layers in self.encoders:
-      x = layers(x, context, time_embedding)
-      skip_connections.append(x)
-
-    return x, skip_connections, time_embedding
-
-
-class UNetBottleNeck(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.bottleneck = SwitchSequential(
-        ResidualBlock(1280, 1280),
-        AttentionBlock(8, 160),
-        ResidualBlock(1280, 1280),
-    )
-
-  def forward(self, x, context, time):
-    x = self.bottleneck(x, context, time)
-    # print('shape')
-    # print(list(x.shape))
-    return x
-
-
-# Unet decoder.
-class UNetDecoder1(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.decoders = nn.ModuleList(
-        [
-            SwitchSequential(ResidualBlock(2560, 1280)),
-            SwitchSequential(ResidualBlock(2560, 1280)),
-            SwitchSequential(ResidualBlock(2560, 1280), Upsample(1280)),
-            SwitchSequential(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),
-        ]
-    )
-
-  def forward(self, x, context, time, s9, s10, s11, s12):
-    x = torch.cat((x, s12), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s11), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s10), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s9), dim=1)
-    x = self.decoders[3](x, context, time)
-
-    return x
-
-
-class UNetDecoder2(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.decoders = nn.ModuleList(
-        [
-            SwitchSequential(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(
-                ResidualBlock(1920, 1280), AttentionBlock(8, 160), Upsample(1280)
-            ),
-            SwitchSequential(ResidualBlock(1920, 640), AttentionBlock(8, 80)),
-            SwitchSequential(ResidualBlock(1280, 640), AttentionBlock(8, 80)),
-        ]
-    )
-
-  def forward(self, x, context, time, s5, s6, s7, s8):
-    x = torch.cat((x, s8), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s7), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s6), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s5), dim=1)
-    x = self.decoders[3](x, context, time)
-    return x
-
-
-class UNetDecoder3(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.decoders = nn.ModuleList(
-        [
-            SwitchSequential(
-                ResidualBlock(960, 640), AttentionBlock(8, 80), Upsample(640)
-            ),
-            SwitchSequential(ResidualBlock(960, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(640, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(640, 320), AttentionBlock(8, 40)),
-        ]
-    )
-    self.final = FinalLayer(320, 4)
-
-  def forward(self, x, context, time, s1, s2, s3, s4):
-    x = torch.cat((x, s4), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s3), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s2), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s1), dim=1)
-    x = self.decoders[3](x, context, time)
-
-    x = self.final(x)
-    return x
-
-
-class UNetDecoder(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.decoders = nn.ModuleList(
-        [
-            SwitchSequential(ResidualBlock(2560, 1280)),
-            SwitchSequential(ResidualBlock(2560, 1280)),
-            SwitchSequential(ResidualBlock(2560, 1280), Upsample(1280)),
-            SwitchSequential(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(ResidualBlock(2560, 1280), AttentionBlock(8, 160)),
-            SwitchSequential(
-                ResidualBlock(1920, 1280), AttentionBlock(8, 160), Upsample(1280)
-            ),
-            SwitchSequential(ResidualBlock(1920, 640), AttentionBlock(8, 80)),
-            SwitchSequential(ResidualBlock(1280, 640), AttentionBlock(8, 80)),
-            SwitchSequential(
-                ResidualBlock(960, 640), AttentionBlock(8, 80), Upsample(640)
-            ),
-            SwitchSequential(ResidualBlock(960, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(640, 320), AttentionBlock(8, 40)),
-            SwitchSequential(ResidualBlock(640, 320), AttentionBlock(8, 40)),
-        ]
-    )
-    self.final = FinalLayer(320, 4)
-
-  def forward(
-      self, x, context, time, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12
-  ):
-    x = torch.cat((x, s12), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s11), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s10), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s9), dim=1)
-    x = self.decoders[3](x, context, time)
-    x = torch.cat((x, s8), dim=1)
-    x = self.decoders[4](x, context, time)
-    x = torch.cat((x, s7), dim=1)
-    x = self.decoders[5](x, context, time)
-    x = torch.cat((x, s6), dim=1)
-    x = self.decoders[6](x, context, time)
-    x = torch.cat((x, s5), dim=1)
-    x = self.decoders[7](x, context, time)
-    x = torch.cat((x, s4), dim=1)
-    x = self.decoders[0](x, context, time)
-    x = torch.cat((x, s3), dim=1)
-    x = self.decoders[1](x, context, time)
-    x = torch.cat((x, s2), dim=1)
-    x = self.decoders[2](x, context, time)
-    x = torch.cat((x, s1), dim=1)
-    x = self.decoders[3](x, context, time)
-
-    x = self.final(x)
-
-    return x
-
-
 class FinalLayer(nn.Module):
 
   def __init__(self, in_channels, out_channels):
@@ -432,68 +234,6 @@ class Diffusion(nn.Module):
   @torch.inference_mode
   def forward(self, latent, context, time):
     time = self.time_embedding(time)
-    # print('time:')
-    # print(list(time.shape))
     output = self.unet(latent, context, time)
     output = self.final(output)
     return output
-
-
-# Calling code as if Diffusion is splitted into two parts.
-class DiffusionSplitted(nn.Module):
-
-  def __init__(self):
-    super().__init__()
-    self.unet_encoder = UNetEncoder()
-    self.bottleneck = UNetBottleNeck()
-    self.unet_decoder1 = UNetDecoder1()
-    self.unet_decoder2 = UNetDecoder2()
-    self.unet_decoder3 = UNetDecoder3()
-
-  def get_skip_connections(self, latent, context, time):
-    _, skip_connections, _ = self.unet_encoder(latent, context, time)
-    return skip_connections
-
-  def forward(self, latent, context, time):
-    output, skip_connections, time = self.unet_encoder(latent, context, time)
-    # print("output shape of unet encoder...")
-    # print(list(output.shape))
-    # print("output shape of time...")
-    # print(list(time.shape))
-    output = self.bottleneck(output, context, time)
-    # print("output shape of bn")
-    # print(list(output.shape))
-    output = self.unet_decoder1(
-        output,
-        context,
-        time,
-        skip_connections[8],
-        skip_connections[9],
-        skip_connections[10],
-        skip_connections[11],
-    )
-    # print("output shape of d1:")
-    # print(list(output.shape))
-
-    output = self.unet_decoder2(
-        output,
-        context,
-        time,
-        skip_connections[4],
-        skip_connections[5],
-        skip_connections[6],
-        skip_connections[7],
-    )
-
-    # print("output shape of d2:")
-    # print(list(output.shape))
-    output = self.unet_decoder3(
-        output,
-        context,
-        time,
-        skip_connections[0],
-        skip_connections[1],
-        skip_connections[2],
-        skip_connections[3],
-    )
-    return output

ai_edge_torch/generative/examples/t5/t5_attention.py

@@ -20,11 +20,11 @@ import torch
 from torch import nn
 import torch.nn.functional as F
 
-from ai_edge_torch.generative.layers.attention import scaled_dot_product_attention  # NOQA
-from ai_edge_torch.generative.layers.attention import scaled_dot_product_attention_with_hlfb  # NOQA
 import ai_edge_torch.generative.layers.builder as builder
 from ai_edge_torch.generative.layers.kv_cache import KVCache
 import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention  # NOQA
+from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention_with_hlfb  # NOQA
 
 
 class EncoderDecoderBlock(nn.Module):