ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/generative/examples/smollm/verify.py

@@ -0,0 +1,86 @@
+# 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.
+# ==============================================================================
+
+"""Verifies the reauthored SmolLM-135M model."""
+
+import logging
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.smollm import smollm
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+_MODEL_VERSION = flags.DEFINE_enum(
+    "model_version",
+    "v1",
+    ["v1", "v2"],
+    "The version of SmolLm to verify.",
+)
+_CHECKPOINT = {
+    "v1": "HuggingFaceTB/SmolLM-135M",
+    "v2": "HuggingFaceTB/SmolLM2-135M",
+}
+
+_BUILDER = {
+    "v1": smollm.build_model,
+    "v2": smollm.build_model_v2,
+}
+
+
+def main(_):
+  checkpoint = _CHECKPOINT[_MODEL_VERSION.value]
+  builder = _BUILDER[_MODEL_VERSION.value]
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
+
+  # Locate the cached dir.
+  cached_config_file = transformers.utils.cached_file(
+      checkpoint, transformers.utils.CONFIG_NAME
+  )
+  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = builder(reauthored_checkpoint)
+
+  logging.info("Loading the tokenizer from: %s", checkpoint)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(tokenizer),
+      generate_prompts=_PROMPTS.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

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,108 @@
+# 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,185 @@
+# 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 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
+import torch
+from torch import nn
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj=(
+        "cond_stage_model.transformer.text_model.encoder.layers.{}.mlp.fc1"
+    ),
+    ff_down_proj=(
+        "cond_stage_model.transformer.text_model.encoder.layers.{}.mlp.fc2"
+    ),
+    attn_query_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.q_proj",
+    attn_key_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.k_proj",
+    attn_value_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.v_proj",
+    attn_output_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.out_proj",
+    pre_attn_norm=(
+        "cond_stage_model.transformer.text_model.encoder.layers.{}.layer_norm1"
+    ),
+    post_attn_norm=(
+        "cond_stage_model.transformer.text_model.encoder.layers.{}.layer_norm2"
+    ),
+    embedding=(
+        "cond_stage_model.transformer.text_model.embeddings.token_embedding"
+    ),
+    embedding_position="cond_stage_model.transformer.text_model.embeddings.position_embedding.weight",
+    final_norm="cond_stage_model.transformer.text_model.final_layer_norm",
+    lm_head=None,
+)
+
+
+class CLIP(nn.Module):
+  """CLIP text encoder.
+
+  For details, see https://arxiv.org/abs/2103.00020
+  """
+
+  def __init__(self, config: cfg.ModelConfig):
+    super().__init__()
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.tok_embedding_position = nn.Parameter(
+        torch.zeros((config.max_seq_len, config.embedding_dim))
+    )
+
+    self.config = config
+    # CLIP has only one block config.
+    block_config = config.block_config(0)
+    self.transformer_blocks = nn.ModuleList(
+        TransformerBlock(block_config, config) for _ in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim, config.final_norm_config
+    )
+
+    self.mask_cache = attention_utils.build_causal_mask_cache(
+        size=config.max_seq_len, dtype=torch.float32
+    )
+
+  @torch.inference_mode
+  def forward(self, tokens: torch.IntTensor) -> torch.FloatTensor:
+    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:
+  """Get configs for the CLIP of Stable Diffusion v1.5."""
+  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,
+      head_dim=embedding_dim // num_heads,
+      num_query_groups=num_query_groups,
+      rotary_base=0,
+      rotary_percentage=0.0,
+      qkv_use_bias=True,
+      qkv_transpose_before_split=True,
+      qkv_fused_interleaved=False,
+      output_proj_use_bias=True,
+      enable_kv_cache=False,
+  )
+
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.SEQUENTIAL,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_QUICK),
+      intermediate_size=embedding_dim * 4,
+      use_bias=True,
+  )
+
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
+
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+
+  config = cfg.ModelConfig(
+      vocab_size=vocab_size,
+      num_layers=num_layers,
+      max_seq_len=max_seq_len,
+      embedding_dim=embedding_dim,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+
+  return config
+
+
+def get_fake_model_config() -> cfg.ModelConfig:
+  """Get fake configs for the CLIP of Stable Diffusion v1.5 for testing."""
+  max_seq_len = 6
+  vocab_size = 100
+  num_layers = 2
+  num_heads = 12
+  num_query_groups = 12
+  embedding_dim = 24
+
+  attn_config = cfg.AttentionConfig(
+      num_heads=num_heads,
+      head_dim=embedding_dim // num_heads,
+      num_query_groups=num_query_groups,
+      rotary_base=0,
+      rotary_percentage=0.0,
+      qkv_use_bias=True,
+      qkv_transpose_before_split=True,
+      qkv_fused_interleaved=False,
+      output_proj_use_bias=True,
+      enable_kv_cache=False,
+  )
+
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.SEQUENTIAL,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_QUICK),
+      intermediate_size=embedding_dim * 4,
+      use_bias=True,
+  )
+
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
+
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+
+  config = cfg.ModelConfig(
+      vocab_size=vocab_size,
+      num_layers=num_layers,
+      max_seq_len=max_seq_len,
+      embedding_dim=embedding_dim,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+
+  return config

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -0,0 +1,173 @@
+# 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 pathlib
+
+from absl import app
+from absl import flags
+import ai_edge_torch
+from ai_edge_torch.generative.examples.stable_diffusion import clip
+from ai_edge_torch.generative.examples.stable_diffusion import decoder
+from ai_edge_torch.generative.examples.stable_diffusion import diffusion
+from ai_edge_torch.generative.examples.stable_diffusion import util
+from ai_edge_torch.generative.quantize import quant_recipes
+from ai_edge_torch.generative.utilities import stable_diffusion_loader
+import torch
+
+_CLIP_CKPT = flags.DEFINE_string(
+    'clip_ckpt',
+    None,
+    help='Path to source CLIP model checkpoint',
+    required=True,
+)
+
+_DIFFUSION_CKPT = flags.DEFINE_string(
+    'diffusion_ckpt',
+    None,
+    help='Path to source diffusion model checkpoint',
+    required=True,
+)
+
+_DECODER_CKPT = flags.DEFINE_string(
+    'decoder_ckpt',
+    None,
+    help='Path to source image decoder model checkpoint',
+    required=True,
+)
+
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir',
+    None,
+    help='Path to the converted TF Lite directory.',
+    required=True,
+)
+
+_QUANTIZE = flags.DEFINE_bool(
+    'quantize',
+    help='Whether to quantize the model during conversion.',
+    default=True,
+)
+
+_DEVICE_TYPE = flags.DEFINE_string(
+    'device_type',
+    'cpu',
+    help='The device type of the model. Currently supported: cpu, gpu.',
+)
+
+
+@torch.inference_mode
+def convert_stable_diffusion_to_tflite(
+    output_dir: str,
+    clip_ckpt_path: str,
+    diffusion_ckpt_path: str,
+    decoder_ckpt_path: str,
+    image_height: int = 512,
+    image_width: int = 512,
+    quantize: bool = True,
+):
+
+  clip_model = clip.CLIP(clip.get_model_config())
+  loader = stable_diffusion_loader.ClipModelLoader(
+      clip_ckpt_path,
+      clip.TENSOR_NAMES,
+  )
+  loader.load(clip_model, strict=False)
+
+  diffusion_model = diffusion.Diffusion(
+      diffusion.get_model_config(batch_size=2, device_type=_DEVICE_TYPE.value)
+  )
+  diffusion_loader = stable_diffusion_loader.DiffusionModelLoader(
+      diffusion_ckpt_path, diffusion.TENSOR_NAMES
+  )
+  diffusion_loader.load(diffusion_model, strict=False)
+
+  decoder_model = decoder.Decoder(
+      decoder.get_model_config(device_type=_DEVICE_TYPE.value)
+  )
+  decoder_loader = stable_diffusion_loader.AutoEncoderModelLoader(
+      decoder_ckpt_path, decoder.TENSOR_NAMES
+  )
+  decoder_loader.load(decoder_model, strict=False)
+
+  # TODO(yichunk): enable image encoder conversion
+  # if encoder_ckpt_path is not None:
+  #   encoder = Encoder()
+  #   encoder.load_state_dict(torch.load(encoder_ckpt_path))
+
+  # Tensors used to trace the model graph during conversion.
+  n_tokens = 77
+  timestamp = 0
+  len_prompt = 1
+  prompt_tokens = torch.full((1, n_tokens), 0, dtype=torch.int)
+  input_image = torch.full(
+      (1, 3, image_height, image_width), 0, dtype=torch.float32
+  )
+  noise = torch.full(
+      (len_prompt, 4, image_height // 8, image_width // 8),
+      0,
+      dtype=torch.float32,
+  )
+
+  input_latents = torch.zeros_like(noise)
+  context_cond = clip_model(prompt_tokens)
+  context_uncond = torch.zeros_like(context_cond)
+  context = torch.cat([context_cond, context_uncond], axis=0)
+  time_embedding = util.get_time_embedding(timestamp)
+
+  if not os.path.exists(output_dir):
+    pathlib.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(
+      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(quant_config=quant_config).export(
+  #     f'{output_dir}/encoder.tflite'
+  # )
+
+  # Diffusion
+  ai_edge_torch.signature(
+      'diffusion',
+      diffusion_model,
+      (torch.repeat_interleave(input_latents, 2, 0), context, time_embedding),
+  ).convert(quant_config=quant_config).export(f'{output_dir}/diffusion.tflite')
+
+  # Image decoder
+  ai_edge_torch.signature('decode', decoder_model, (input_latents,)).convert(
+      quant_config=quant_config
+  ).export(f'{output_dir}/decoder.tflite')
+
+
+def main(_):
+  convert_stable_diffusion_to_tflite(
+      output_dir=_OUTPUT_DIR.value,
+      clip_ckpt_path=_CLIP_CKPT.value,
+      diffusion_ckpt_path=_DIFFUSION_CKPT.value,
+      decoder_ckpt_path=_DECODER_CKPT.value,
+      quantize=_QUANTIZE.value,
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)