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

ai_edge_torch/generative/examples/stable_diffusion/tokenizer.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.
+# ==============================================================================
+
+import functools
+import json
+import os
+from typing import List, Tuple
+import unicodedata
+
+import regex as re
+
+
+def create_bytes_table() -> dict:
+  table = {}
+  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
+      table[byte] = chr(byte)
+    else:  # ith character IS control char or space
+      table[byte] = chr(special_count + 256)
+      special_count += 1
+  return table
+
+
+def pairwise(seq):
+  a = iter(seq)
+  b = iter(seq)
+  next(b)
+  return zip(a, b)
+
+
+class Tokenizer:
+
+  def __init__(self, vocab_dir: str):
+    with open(os.path.join(vocab_dir, 'vocab.json'), encoding='utf-8') as f:
+      self.vocab = json.load(f)
+
+    with open(os.path.join(vocab_dir, 'merges.txt'), encoding='utf-8') as f:
+      lines = f.read().split('\n')
+      lines = lines[1:-1]
+      self.merges = {tuple(bigram.split()): i for i, bigram in enumerate(lines)}
+
+    self.bos_token = self.vocab['<|startoftext|>']
+    self.eos_token = self.vocab['<|endoftext|>']
+    self.pad_token = self.vocab['<|endoftext|>']
+    self.max_length = 77
+    self.bytes_table = create_bytes_table()
+    self.chunk_pattern = re.compile(
+        r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
+        re.IGNORECASE,
+    )
+
+  def encode(self, text: str) -> List[int]:
+    text = unicodedata.normalize('NFC', text)
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    text = text.lower()
+
+    tokens = [self.bos_token]
+    for chunk in re.findall(self.chunk_pattern, text):
+      chunk = ''.join(self.bytes_table[byte] for byte in chunk.encode('utf-8'))
+      tokens.extend(self.vocab[word] for word in self.bpe(chunk))
+    tokens.append(self.eos_token)
+
+    tokens = tokens[: self.max_length]
+    token_length = len(tokens)
+    pad_length = self.max_length - token_length
+    tokens += [self.pad_token] * pad_length
+    return tokens
+
+  def encode_batch(self, texts: List[str]) -> List[List[int]]:
+    return [self.encode(text) for text in texts]
+
+  @functools.lru_cache(maxsize=10000)
+  def bpe(self, chunk: str) -> Tuple[str]:
+    words = list(chunk)
+    words[-1] += '</w>'
+
+    while len(words) > 1:
+      valid_pairs = [pair for pair in pairwise(words) if pair in self.merges]
+      if not valid_pairs:
+        break
+
+      bigram = min(valid_pairs, key=lambda pair: self.merges[pair])
+      first, second = bigram
+
+      new_words = []
+      for word in words:
+        if word == second and new_words and new_words[-1] == first:
+          new_words[-1] = first + second
+        else:
+          new_words.append(word)
+      words = new_words
+
+    return tuple(words)

ai_edge_torch/generative/examples/stable_diffusion/util.py

@@ -0,0 +1,77 @@
+# 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":
+    if isinstance(image, torch.Tensor):
+      return image.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+    if isinstance(image, np.ndarray):
+      return image.transpose(0, 3, 1, 2)
+  elif to == "last":
+    if isinstance(image, torch.Tensor):
+      return image.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+    if isinstance(image, np.ndarray):
+      return image.transpose(0, 2, 3, 1)
+  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:
+    if isinstance(x, torch.Tensor):
+      x = x.clamp(new_min, new_max)
+    elif isinstance(x, np.ndarray):
+      x = x.clip(new_min, new_max)
+  return x

ai_edge_torch/generative/examples/t5/__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/t5/convert_to_tflite.py

@@ -0,0 +1,138 @@
+# 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
+from pathlib import Path
+
+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.
+def convert_t5_to_tflite_singlesig(checkpoint_path: str):
+  pytorch_model = t5.build_t5_model(checkpoint_path)
+
+  # encoder
+  seq_len = 512
+  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
+  prompt_e_token = [1, 2, 3, 4, 5, 6]
+  prefill_e_tokens[0, : len(prompt_e_token)] = torch.tensor(
+      prompt_e_token, dtype=torch.int
+  )
+  prefill_e_input_pos = torch.arange(0, seq_len, dtype=torch.int)
+  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
+  prompt_d_token = [1, 2, 3, 4, 5, 6]
+  prefill_d_tokens[0, : len(prompt_d_token)] = torch.tensor(
+      prompt_d_token, dtype=torch.int
+  )
+  prefill_d_input_pos = torch.arange(0, seq_len, dtype=torch.int)
+
+  # decoder
+  decode_token = torch.tensor([[1]], dtype=torch.int)
+  decode_input_pos = torch.tensor([0], dtype=torch.int)
+  decode_d_token = torch.tensor([[1]], dtype=torch.int)
+  decode_d_input_pos = torch.tensor([0], dtype=torch.int)
+
+  # Pad mask for self attention only on "real" tokens.
+  # Pad with `-inf` for any tokens indices that aren't desired.
+  pad_mask = torch.zeros([seq_len], dtype=torch.float32)
+
+  edge_model = ai_edge_torch.signature(
+      'decode',
+      pytorch_model,
+      (
+          prefill_e_tokens,
+          prefill_e_input_pos,
+          decode_d_token,
+          decode_d_input_pos,
+          pad_mask,
+      ),
+  ).convert()
+
+  edge_model.export('/tmp/t5_encode_decode.tflite')
+
+
+def convert_t5_to_tflite_multisig(checkpoint_path: str):
+  config = t5.get_model_config_t5()
+  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
+  )
+
+  # encoder
+  seq_len = 512
+  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
+  prompt_e_token = [1, 2, 3, 4, 5, 6]
+  prefill_e_tokens[0, : len(prompt_e_token)] = torch.tensor(
+      prompt_e_token, dtype=torch.int
+  )
+  prefill_e_input_pos = torch.arange(0, seq_len, dtype=torch.int)
+  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.int)
+  prompt_d_token = [1, 2, 3, 4, 5, 6]
+  prefill_d_tokens[0, : len(prompt_d_token)] = torch.tensor(
+      prompt_d_token, dtype=torch.int
+  )
+  prefill_d_input_pos = torch.arange(0, seq_len, dtype=torch.int)
+
+  # decoder
+  decode_token = torch.tensor([[1]], dtype=torch.int)
+  decode_input_pos = torch.tensor([0], dtype=torch.int)
+  decode_d_token = torch.tensor([[1]], dtype=torch.int)
+  decode_d_input_pos = torch.tensor([0], dtype=torch.int)
+
+  # Pad mask for self attention only on "real" tokens.
+  # Pad with `-inf` for any tokens indices that aren't desired.
+  pad_mask = torch.zeros([seq_len], dtype=torch.float32)
+  hidden_states = torch.zeros((1, 512, 768), dtype=torch.float32)
+  quant_config = quant_recipes.full_int8_dynamic_recipe()
+
+  edge_model = (
+      ai_edge_torch.signature(
+          'encode',
+          t5_encoder_model,
+          (
+              prefill_e_tokens,
+              prefill_e_input_pos,
+              pad_mask,
+          ),
+      )
+      .signature(
+          'decode',
+          t5_decoder_model,
+          (
+              hidden_states,
+              decode_d_token,
+              decode_d_input_pos,
+              pad_mask,
+          ),
+      )
+      .convert(quant_config=quant_config)
+  )
+
+  edge_model.export('/tmp/t5_encode_decode_2_sigs.tflite')
+
+
+if __name__ == '__main__':
+  checkpoint_path = os.path.join(Path.home(), 'Downloads/llm_data/t5')
+  # convert_t5_to_tflite_singlesig(checkpoint_path)
+  convert_t5_to_tflite_multisig(checkpoint_path)