ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__py3-none-any.whl

ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py

@@ -28,6 +28,7 @@ class SamplerInterface(abc.ABC):
   @abc.abstractmethod
   def set_strength(self, strength: float = 1) -> None:
     """Set the strength of initial step.
+
     Conceptually, indicates how much to transform the reference `input_images`.
     """
     return NotImplemented

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.
@@ -30,24 +29,24 @@ def convert_t5_to_tflite_singlesig(checkpoint_path: str):
 
   # encoder
   seq_len = 512
-  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
+  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.long
+      prompt_e_token, dtype=torch.int
   )
-  prefill_e_input_pos = torch.arange(0, seq_len)
-  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
+  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.long
+      prompt_d_token, dtype=torch.int
   )
-  prefill_d_input_pos = torch.arange(0, seq_len)
+  prefill_d_input_pos = torch.arange(0, seq_len, dtype=torch.int)
 
   # decoder
-  decode_token = torch.tensor([[1]], dtype=torch.long)
-  decode_input_pos = torch.tensor([0], dtype=torch.int64)
-  decode_d_token = torch.tensor([[1]], dtype=torch.long)
-  decode_d_input_pos = torch.tensor([0], dtype=torch.int64)
+  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.
@@ -73,29 +72,33 @@ 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
-  prefill_e_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
+  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.long
+      prompt_e_token, dtype=torch.int
   )
-  prefill_e_input_pos = torch.arange(0, seq_len)
-  prefill_d_tokens = torch.full((1, seq_len), 0, dtype=torch.long)
+  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.long
+      prompt_d_token, dtype=torch.int
   )
-  prefill_d_input_pos = torch.arange(0, seq_len)
+  prefill_d_input_pos = torch.arange(0, seq_len, dtype=torch.int)
 
   # decoder
-  decode_token = torch.tensor([[1]], dtype=torch.long)
-  decode_input_pos = torch.tensor([0], dtype=torch.int64)
-  decode_d_token = torch.tensor([[1]], dtype=torch.long)
-  decode_d_input_pos = torch.tensor([0], dtype=torch.int64)
+  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.