ai-edge-torch-nightly 0.1.dev202405131930__py3-none-any.whl → 0.2.0.dev20240527__py3-none-any.whl

ai_edge_torch/generative/examples/stable_diffusion/__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.
+# ==============================================================================

ai_edge_torch/generative/examples/stable_diffusion/attention.py

@@ -0,0 +1,106 @@
+# 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 math
+
+import torch
+from torch import _decomp
+from torch import nn
+from torch._prims_common import mask_tensor
+from torch._prims_common.wrappers import out_wrapper
+from torch.nn import functional as F
+
+
+def triu(a):
+  h, w = a.shape[-2:]
+  mask = (
+      torch.arange(w, device=a.device).unsqueeze(-2)
+      - torch.arange(h, device=a.device).unsqueeze(-1)
+  ) >= 1
+  mask = torch.broadcast_to(mask, a.shape)
+  return torch.ops.aten.logical_and(a, mask).contiguous()
+
+
+# _decomp.decomposition_table[torch.ops.aten.triu.default] = triu
+
+
+class SelfAttention(nn.Module):
+
+  def __init__(self, n_heads, d_embed, in_proj_bias=True, out_proj_bias=True):
+    super().__init__()
+    self.in_proj = nn.Linear(d_embed, 3 * d_embed, bias=in_proj_bias)
+    self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)
+    self.n_heads = n_heads
+    self.d_head = d_embed // n_heads
+
+  def forward(self, x, causal_mask=False):
+    input_shape = x.shape
+    batch_size, sequence_length, d_embed = input_shape
+    interim_shape = (batch_size, sequence_length, self.n_heads, self.d_head)
+
+    q, k, v = self.in_proj(x).chunk(3, dim=-1)
+
+    q = q.view(interim_shape).transpose(1, 2)
+    k = k.view(interim_shape).transpose(1, 2)
+    v = v.view(interim_shape).transpose(1, 2)
+
+    weight = q @ k.transpose(-1, -2)
+    if causal_mask:
+      # mask = torch.ones_like(weight, dtype=torch.bool).triu(1)
+      mask = triu(torch.ones_like(weight, dtype=torch.bool))
+      weight.masked_fill_(mask, -torch.inf)
+    weight /= math.sqrt(self.d_head)
+    weight = F.softmax(weight, dim=-1)
+
+    output = weight @ v
+    output = output.transpose(1, 2)
+    output = output.reshape(input_shape)
+    output = self.out_proj(output)
+    return output
+
+
+class CrossAttention(nn.Module):
+
+  def __init__(self, n_heads, d_embed, d_cross, in_proj_bias=True, out_proj_bias=True):
+    super().__init__()
+    self.q_proj = nn.Linear(d_embed, d_embed, bias=in_proj_bias)
+    self.k_proj = nn.Linear(d_cross, d_embed, bias=in_proj_bias)
+    self.v_proj = nn.Linear(d_cross, d_embed, bias=in_proj_bias)
+    self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)
+    self.n_heads = n_heads
+    self.d_head = d_embed // n_heads
+
+  def forward(self, x, y):
+    input_shape = x.shape
+    batch_size, sequence_length, d_embed = input_shape
+    interim_shape = (batch_size, -1, self.n_heads, self.d_head)
+
+    q = self.q_proj(x)
+    k = self.k_proj(y)
+    v = self.v_proj(y)
+
+    q = q.view(interim_shape).transpose(1, 2)
+    k = k.view(interim_shape).transpose(1, 2)
+    v = v.view(interim_shape).transpose(1, 2)
+
+    weight = q @ k.transpose(-1, -2)
+    weight /= math.sqrt(self.d_head)
+    weight = F.softmax(weight, dim=-1)
+
+    output = weight @ v
+    output = output.transpose(1, 2).contiguous()
+    output = output.view(input_shape)
+    output = self.out_proj(output)
+    return output

ai_edge_torch/generative/examples/stable_diffusion/clip.py

@@ -0,0 +1,78 @@
+# 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 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
+
+
+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):
+
+  def __init__(self, n_head: int, n_embd: int):
+    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
+
+    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
+
+    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)
+
+  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)
+    return output

ai_edge_torch/generative/examples/stable_diffusion/decoder.py

@@ -0,0 +1,111 @@
+# 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.
+# ==============================================================================
+
+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):
+
+  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 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),
+    )
+
+  def forward(self, x):
+    x = x / 0.18215
+    for module in self:
+      x = module(x)
+    return x

ai_edge_torch/generative/examples/stable_diffusion/diffusion.py

@@ -0,0 +1,498 @@
+# 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 torch.nn import functional as F
+
+from ai_edge_torch.generative.examples.stable_diffusion.attention import CrossAttention  # NOQA
+from ai_edge_torch.generative.examples.stable_diffusion.attention import SelfAttention  # NOQA
+
+
+class TimeEmbedding(nn.Module):
+
+  def __init__(self, n_embd):
+    super().__init__()
+    self.linear_1 = nn.Linear(n_embd, 4 * n_embd)
+    self.linear_2 = nn.Linear(4 * n_embd, 4 * n_embd)
+
+  def forward(self, x):
+    x = self.linear_1(x)
+    x = F.silu(x)
+    x = self.linear_2(x)
+    return x
+
+
+class ResidualBlock(nn.Module):
+
+  def __init__(self, in_channels, out_channels, n_time=1280):
+    super().__init__()
+    self.groupnorm_feature = nn.GroupNorm(32, in_channels)
+    self.conv_feature = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+    self.linear_time = nn.Linear(n_time, out_channels)
+
+    self.groupnorm_merged = nn.GroupNorm(32, out_channels)
+    self.conv_merged = 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, feature, time):
+    residue = feature
+
+    feature = self.groupnorm_feature(feature)
+    feature = F.silu(feature)
+    feature = self.conv_feature(feature)
+
+    time = F.silu(time)
+    time = self.linear_time(time)
+
+    merged = feature + time.unsqueeze(-1).unsqueeze(-1)
+    merged = self.groupnorm_merged(merged)
+    merged = F.silu(merged)
+    merged = self.conv_merged(merged)
+
+    return merged + self.residual_layer(residue)
+
+
+class AttentionBlock(nn.Module):
+
+  def __init__(self, n_head: int, n_embd: int, d_context=768):
+    super().__init__()
+    channels = n_head * n_embd
+
+    self.groupnorm = nn.GroupNorm(32, channels, eps=1e-6)
+    self.conv_input = nn.Conv2d(channels, channels, kernel_size=1, padding=0)
+
+    self.layernorm_1 = nn.LayerNorm(channels)
+    self.attention_1 = SelfAttention(n_head, channels, in_proj_bias=False)
+    self.layernorm_2 = nn.LayerNorm(channels)
+    self.attention_2 = CrossAttention(n_head, channels, d_context, in_proj_bias=False)
+    self.layernorm_3 = nn.LayerNorm(channels)
+    self.linear_geglu_1 = nn.Linear(channels, 4 * channels * 2)
+    self.linear_geglu_2 = nn.Linear(4 * channels, channels)
+
+    self.conv_output = nn.Conv2d(channels, channels, kernel_size=1, padding=0)
+
+  def forward(self, x, context):
+    residue_long = x
+
+    x = self.groupnorm(x)
+    x = self.conv_input(x)
+
+    n, c, h, w = x.shape
+    x = x.view((n, c, h * w))  # (n, c, hw)
+    x = x.transpose(-1, -2)  # (n, hw, c)
+
+    residue_short = x
+    x = self.layernorm_1(x)
+    x = self.attention_1(x)
+    x += residue_short
+
+    residue_short = x
+    x = self.layernorm_2(x)
+    x = self.attention_2(x, context)
+    x += residue_short
+
+    residue_short = x
+    x = self.layernorm_3(x)
+    x, gate = self.linear_geglu_1(x).chunk(2, dim=-1)
+    x = x * F.gelu(gate)
+    x = self.linear_geglu_2(x)
+    x += residue_short
+
+    x = x.transpose(-1, -2)  # (n, c, hw)
+    x = x.view((n, c, h, w))  # (n, c, h, w)
+
+    return self.conv_output(x) + residue_long
+
+
+class Upsample(nn.Module):
+
+  def __init__(self, channels):
+    super().__init__()
+    self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+
+  def forward(self, x):
+    x = F.interpolate(x, scale_factor=2, mode='nearest')
+    return self.conv(x)
+
+
+class SwitchSequential(nn.Sequential):
+
+  def forward(self, x, context, time):
+    for layer in self:
+      if isinstance(layer, AttentionBlock):
+        x = layer(x, context)
+      elif isinstance(layer, ResidualBlock):
+        x = layer(x, time)
+      else:
+        x = layer(x)
+    return x
+
+
+class UNet(nn.Module):
+
+  def __init__(self):
+    super().__init__()
+    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)),
+        ]
+    )
+    self.bottleneck = SwitchSequential(
+        ResidualBlock(1280, 1280),
+        AttentionBlock(8, 160),
+        ResidualBlock(1280, 1280),
+    )
+
+    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)),
+        ]
+    )
+
+  def forward(self, x, context, time):
+    skip_connections = []
+    for layers in self.encoders:
+      x = layers(x, context, time)
+      skip_connections.append(x)
+
+    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)
+
+    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):
+    super().__init__()
+    self.groupnorm = nn.GroupNorm(32, in_channels)
+    self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+
+  def forward(self, x):
+    x = self.groupnorm(x)
+    x = F.silu(x)
+    x = self.conv(x)
+    return x
+
+
+class Diffusion(nn.Module):
+
+  def __init__(self):
+    super().__init__()
+    self.time_embedding = TimeEmbedding(320)
+    self.unet = UNet()
+    self.final = FinalLayer(320, 4)
+
+  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/stable_diffusion/encoder.py

@@ -0,0 +1,65 @@
+# 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 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
+
+
+class Encoder(nn.Sequential):
+
+  def __init__(self):
+    super().__init__(
+        nn.Conv2d(3, 128, kernel_size=3, padding=1),
+        ResidualBlock(128, 128),
+        ResidualBlock(128, 128),
+        nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=0),
+        ResidualBlock(128, 256),
+        ResidualBlock(256, 256),
+        nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=0),
+        ResidualBlock(256, 512),
+        ResidualBlock(512, 512),
+        nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=0),
+        ResidualBlock(512, 512),
+        ResidualBlock(512, 512),
+        ResidualBlock(512, 512),
+        AttentionBlock(512),
+        ResidualBlock(512, 512),
+        nn.GroupNorm(32, 512),
+        nn.SiLU(),
+        nn.Conv2d(512, 8, kernel_size=3, padding=1),
+        nn.Conv2d(8, 8, kernel_size=1, padding=0),
+    )
+
+  def forward(self, x, noise):
+    for module in self:
+      if getattr(module, 'stride', None) == (
+          2,
+          2,
+      ):  # Padding at downsampling should be asymmetric (see #8)
+        x = F.pad(x, (0, 1, 0, 1))
+      x = module(x)
+
+    mean, log_variance = torch.chunk(x, 2, dim=1)
+    log_variance = torch.clamp(log_variance, -30, 20)
+    variance = log_variance.exp()
+    stdev = variance.sqrt()
+    x = mean + stdev * noise
+
+    x *= 0.18215
+    return x

ai_edge_torch/generative/examples/stable_diffusion/util.py

@@ -0,0 +1,62 @@
+# 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 os
+
+import numpy as np
+import torch
+
+
+def get_time_embedding(timestep):
+  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):
+  betas = (
+      np.linspace(beta_start**0.5, beta_end**0.5, n_training_steps, dtype=np.float32)
+      ** 2
+  )
+  alphas = 1.0 - betas
+  alphas_cumprod = np.cumprod(alphas, axis=0)
+  return alphas_cumprod
+
+
+def get_file_path(filename, url=None):
+  module_location = os.path.dirname(os.path.abspath(__file__))
+  parent_location = os.path.dirname(module_location)
+  file_location = os.path.join(parent_location, "data", filename)
+  return file_location
+
+
+def move_channel(image, to):
+  if to == "first":
+    return image.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+  elif to == "last":
+    return image.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+  else:
+    raise ValueError("to must be one of the following: first, last")
+
+
+def rescale(x, old_range, new_range, clamp=False):
+  old_min, old_max = old_range
+  new_min, new_max = new_range
+  x -= old_min
+  x *= (new_max - new_min) / (old_max - old_min)
+  x += new_min
+  if clamp:
+    x = x.clamp(new_min, new_max)
+  return x

ai_edge_torch/generative/test/loader_test.py

@@ -0,0 +1,80 @@
+# 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.
+# ==============================================================================
+# Testing weight loader utilities.
+
+import os
+import tempfile
+import unittest
+
+import safetensors.torch
+import torch
+
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+from ai_edge_torch.generative.utilities import loader as loading_utils
+
+
+class TestLoader(unittest.TestCase):
+  """Unit tests that check weight loader."""
+
+  def test_load_safetensors(self):
+    with tempfile.TemporaryDirectory() as temp_dir:
+      file_path = os.path.join(temp_dir, "test.safetensors")
+      test_data = {"weight": torch.randn(20, 10), "bias": torch.randn(20)}
+      safetensors.torch.save_file(test_data, file_path)
+
+      loaded_tensors = loading_utils.load_safetensors(file_path)
+      self.assertIn("weight", loaded_tensors)
+      self.assertIn("bias", loaded_tensors)
+
+  def test_load_statedict(self):
+    with tempfile.TemporaryDirectory() as temp_dir:
+      file_path = os.path.join(temp_dir, "test.pt")
+      model = torch.nn.Linear(10, 5)
+      state_dict = model.state_dict()
+      torch.save(state_dict, file_path)
+
+      loaded_tensors = loading_utils.load_pytorch_statedict(file_path)
+      self.assertIn("weight", loaded_tensors)
+      self.assertIn("bias", loaded_tensors)
+
+  def test_model_loader(self):
+    with tempfile.TemporaryDirectory() as temp_dir:
+      file_path = os.path.join(temp_dir, "test.safetensors")
+      test_weights = {
+          "lm_head.weight": torch.randn((32000, 2048)),
+          "model.embed_tokens.weight": torch.randn((32000, 2048)),
+          "model.layers.0.input_layernorm.weight": torch.randn((2048,)),
+          "model.layers.0.mlp.down_proj.weight": torch.randn((2048, 5632)),
+          "model.layers.0.mlp.gate_proj.weight": torch.randn((5632, 2048)),
+          "model.layers.0.mlp.up_proj.weight": torch.randn((5632, 2048)),
+          "model.layers.0.post_attention_layernorm.weight": torch.randn((2048,)),
+          "model.layers.0.self_attn.k_proj.weight": torch.randn((256, 2048)),
+          "model.layers.0.self_attn.o_proj.weight": torch.randn((2048, 2048)),
+          "model.layers.0.self_attn.q_proj.weight": torch.randn((2048, 2048)),
+          "model.layers.0.self_attn.v_proj.weight": torch.randn((256, 2048)),
+          "model.norm.weight": torch.randn((2048,)),
+      }
+      safetensors.torch.save_file(test_weights, file_path)
+      cfg = tiny_llama.get_model_config()
+      cfg.num_layers = 1
+      model = tiny_llama.TinyLLamma(cfg)
+
+      loader = loading_utils.ModelLoader(file_path, tiny_llama.TENSOR_NAMES)
+      # if returns successfully, it means all the tensors were initiallized.
+      loader.load(model, strict=True)
+
+
+if __name__ == "__main__":
+  unittest.main()

ai_edge_torch/generative/utilities/loader.py

@@ -228,14 +228,14 @@ class ModelLoader:
     q_name = self._names.attn_query_proj.format(idx)
     k_name = self._names.attn_key_proj.format(idx)
     v_name = self._names.attn_value_proj.format(idx)
-    converted_state[f"{prefix}.atten_func.attn.weight"] = self._fuse_qkv(
+    converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = self._fuse_qkv(
         config,
         state.pop(f"{q_name}.weight"),
         state.pop(f"{k_name}.weight"),
         state.pop(f"{v_name}.weight"),
     )
     if config.attn_config.qkv_use_bias:
-      converted_state[f"{prefix}.atten_func.attn.bias"] = self._fuse_qkv(
+      converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = self._fuse_qkv(
           config,
           state.pop(f"{q_name}.bias"),
           state.pop(f"{k_name}.bias"),
@@ -243,9 +243,13 @@ class ModelLoader:
       )
 
     o_name = self._names.attn_output_proj.format(idx)
-    converted_state[f"{prefix}.atten_func.proj.weight"] = state.pop(f"{o_name}.weight")
+    converted_state[f"{prefix}.atten_func.output_projection.weight"] = state.pop(
+        f"{o_name}.weight"
+    )
     if config.attn_config.output_proj_use_bias:
-      converted_state[f"{prefix}.atten_func.proj.bias"] = state.pop(f"{o_name}.bias")
+      converted_state[f"{prefix}.atten_func.output_projection.bias"] = state.pop(
+          f"{o_name}.bias"
+      )
 
   def _map_norm(
       self,

{ai_edge_torch_nightly-0.1.dev202405131930.dist-info → ai_edge_torch_nightly-0.2.0.dev20240527.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.1.dev202405131930
+Version: 0.2.0.dev20240527
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI
@@ -24,8 +24,8 @@ Requires-Python: >=3.9, <3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy
-Requires-Dist: safetensors
 Requires-Dist: scipy
+Requires-Dist: safetensors
 Requires-Dist: tabulate
 Requires-Dist: torch ==2.4.*
 

{ai_edge_torch_nightly-0.1.dev202405131930.dist-info → ai_edge_torch_nightly-0.2.0.dev20240527.dist-info}/RECORD

@@ -38,6 +38,13 @@ ai_edge_torch/generative/examples/gemma/gemma.py,sha256=YF4Ua-1lnL3qhQnh1sY5-HlY
 ai_edge_torch/generative/examples/phi2/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/phi2/convert_to_tflite.py,sha256=6nOuwx9q3AUlYcjXRRXSr_3M2JKqdJ-vUf-uE3VFYHE,2512
 ai_edge_torch/generative/examples/phi2/phi2.py,sha256=VvigzPQ_LJHeADTsMliwFwPe2BcnOhFgKDqr_WZ2JQ8,5540
+ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
+ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=Lo4Dq7a3Kg-lyH56iqGtqCo5UaClQHRCTDdNagXGTo8,3535
+ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=8W4X9PKdnMsWGxXbBfm5OX6mX4XhvaMZ2gZw8yCTScY,2410
+ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=beLCtogA32oYT2nlATpyT-1xzkyPF8zi4v3kfHpw6Mc,3239
+ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=nnsfgjSeL16U3TVdjTkRycaoWA2ChFeitx2RjGLpwyA,16200
+ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=X6ekByU19KNHNh5OaztZEROv-QwcCwVm1xiJjm2SCoo,2251
+ai_edge_torch/generative/examples/stable_diffusion/util.py,sha256=pG_dsV4xIaB7B8MgoRgSXBvLCVqDlF6bNunPN3GIm-s,2046
 ai_edge_torch/generative/examples/t5/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/t5/convert_to_tflite.py,sha256=bWtwtUacvJOEDUpuYvLTgkP7oTkXKJA-Tf4FPxlD1Cw,4536
 ai_edge_torch/generative/examples/t5/t5.py,sha256=q2gG5RRo7RgNzvHXYC0Juh6Tgt5d_RTMSWFaYvOKiZU,21065
@@ -65,10 +72,11 @@ ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=9ktL7fT8C5j1dnY_7
 ai_edge_torch/generative/quantize/quant_recipes.py,sha256=CRA2ENevS-3usHqidWDe2wrf_epILE_7Hx-XfZQ9buk,1798
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=OQ4ghQXknA1PPjuY-xBgAmOpaIBgYFM8F2YAIot06hE,1345
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
+ai_edge_torch/generative/test/loader_test.py,sha256=N88CbrLW7Q2x1EyurwdXQ6YjsA-ySQcPxpZH3QOGp-M,3317
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=1NfZxKo9Gx6CmVfd86K1FkmsNQnjzIV1ojBS85UGvT0,6500
 ai_edge_torch/generative/test/test_quantize.py,sha256=f70sH1ZFzdCwYj0MG-eg54WOC4LasR0D8CTUYpjxZYM,3728
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/loader.py,sha256=QrGZ3JlEN_tn8j6EdZOxVt_0u3yB5vBrR3KJtNaAwV8,10029
+ai_edge_torch/generative/utilities/loader.py,sha256=c-ZOIDBVnat_5l2W5sWU7HQm7CL-wducS8poSu5PlUg,10107
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=guDTv-12UUvJGl4eDvvZX3t4rRKewfXO8SpcYXM6gbc,16156
 ai_edge_torch/hlfb/__init__.py,sha256=rrje8a2iuKboBoV96bVq7nlS9HsnuEMbHE5JiWmCxFA,752
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=2VXnHcGf23VOuP-1GriGIpuL98leBB8twp_qaScMnmc,4799
@@ -84,8 +92,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=ExThdTXqnWmGC3-F6sdXbXr8nYzkEe_qCz
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=EIyKz-HY70DguWuSrJal8LpYXQ5ZSEUf3ZrVl7jikFM,4286
-ai_edge_torch_nightly-0.1.dev202405131930.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.1.dev202405131930.dist-info/METADATA,sha256=lQcAb0esNisYUqkzDRHamW4S9luvrJ4QU75042IAqWc,1750
-ai_edge_torch_nightly-0.1.dev202405131930.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.1.dev202405131930.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.1.dev202405131930.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240527.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240527.dist-info/METADATA,sha256=nbZoIm0s6CWdrMkaffTrpz-XooKzTR1q0SQ17rs-AKU,1748
+ai_edge_torch_nightly-0.2.0.dev20240527.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240527.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240527.dist-info/RECORD,,