ai-edge-torch-nightly 0.2.0.dev20240801__py3-none-any.whl → 0.2.0.dev20240803__py3-none-any.whl

ai_edge_torch/generative/test/test_quantize.py

@@ -15,9 +15,6 @@
 
 import unittest
 
-from parameterized import parameterized
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.generative.examples.test_models import toy_model  # NOQA
 from ai_edge_torch.generative.quantize import quant_recipe
@@ -29,20 +26,20 @@ from ai_edge_torch.generative.quantize.quant_attrs import Granularity
 from ai_edge_torch.generative.quantize.quant_attrs import Mode
 from ai_edge_torch.quantize import quant_config
 from ai_edge_torch.testing import model_coverage
+from parameterized import parameterized
+import torch
 
 
 class TestVerifyRecipes(unittest.TestCase):
   """Unit tests that check for model quantization recipes."""
 
-  @parameterized.expand(
-      [
-          (Dtype.FP32, Dtype.FP32),
-          (Dtype.INT8, Dtype.INT8),
-          (Dtype.INT8, Dtype.FP16),
-          (Dtype.FP16, Dtype.INT8),
-          (Dtype.FP16, Dtype.FP16),
-      ]
-  )
+  @parameterized.expand([
+      (Dtype.FP32, Dtype.FP32),
+      (Dtype.INT8, Dtype.INT8),
+      (Dtype.INT8, Dtype.FP16),
+      (Dtype.FP16, Dtype.INT8),
+      (Dtype.FP16, Dtype.FP16),
+  ])
   def test_verify_invalid_recipes(
       self,
       activation,
@@ -54,31 +51,29 @@ class TestVerifyRecipes(unittest.TestCase):
           with self.assertRaises(ValueError):
             quant_recipe.LayerQuantRecipe(activation, weight, m, a, g).verify()
 
-  @parameterized.expand(
-      [
-          (
-              Dtype.FP32,
-              Dtype.INT8,
-              Mode.DYNAMIC_RANGE,
-              Algorithm.MIN_MAX,
-              Granularity.CHANNELWISE,
-          ),
-          (
-              Dtype.FP32,
-              Dtype.INT8,
-              Mode.WEIGHT_ONLY,
-              Algorithm.MIN_MAX,
-              Granularity.CHANNELWISE,
-          ),
-          (
-              Dtype.FP32,
-              Dtype.FP16,
-              Mode.WEIGHT_ONLY,
-              Algorithm.FLOAT_CAST,
-              Granularity.NONE,
-          ),
-      ]
-  )
+  @parameterized.expand([
+      (
+          Dtype.FP32,
+          Dtype.INT8,
+          Mode.DYNAMIC_RANGE,
+          Algorithm.MIN_MAX,
+          Granularity.CHANNELWISE,
+      ),
+      (
+          Dtype.FP32,
+          Dtype.INT8,
+          Mode.WEIGHT_ONLY,
+          Algorithm.MIN_MAX,
+          Granularity.CHANNELWISE,
+      ),
+      (
+          Dtype.FP32,
+          Dtype.FP16,
+          Mode.WEIGHT_ONLY,
+          Algorithm.FLOAT_CAST,
+          Granularity.NONE,
+      ),
+  ])
   def test_verify_valid_recipes(
       self,
       activation,
@@ -87,7 +82,9 @@ class TestVerifyRecipes(unittest.TestCase):
       algo,
       granularity,
   ):
-    quant_recipe.LayerQuantRecipe(activation, weight, mode, algo, granularity).verify()
+    quant_recipe.LayerQuantRecipe(
+        activation, weight, mode, algo, granularity
+    ).verify()
 
 
 class TestQuantizeConvert(unittest.TestCase):
@@ -107,15 +104,13 @@ class TestQuantizeConvert(unittest.TestCase):
         )
     )
 
-  @parameterized.expand(
-      [
-          (quant_recipes.full_fp16_recipe()),
-          (quant_recipes.full_int8_dynamic_recipe()),
-          (quant_recipes.full_int8_weight_only_recipe()),
-          (_attention_int8_dynamic_recipe()),
-          (_feedforward_int8_dynamic_recipe()),
-      ]
-  )
+  @parameterized.expand([
+      (quant_recipes.full_fp16_recipe()),
+      (quant_recipes.full_int8_dynamic_recipe()),
+      (quant_recipes.full_int8_weight_only_recipe()),
+      (_attention_int8_dynamic_recipe()),
+      (_feedforward_int8_dynamic_recipe()),
+  ])
   def test_quantize_convert_toy_sizes(self, quant_config):
     config = toy_model.get_model_config()
     pytorch_model = toy_model.ToySingleLayerModel(config)
@@ -146,7 +141,9 @@ class TestQuantizeConvert(unittest.TestCase):
     )
     float_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
 
-    self.assertLess(len(quantized_model._tflite_model), len(float_model._tflite_model))
+    self.assertLess(
+        len(quantized_model._tflite_model), len(float_model._tflite_model)
+    )
     self.assertTrue(
         model_coverage.compare_tflite_torch(
             quantized_model,

ai_edge_torch/generative/utilities/loader.py

@@ -18,11 +18,10 @@ import glob
 import os
 from typing import Callable, Dict, List, Tuple
 
+from ai_edge_torch.generative.layers import model_config
 from safetensors import safe_open
 import torch
 
-from ai_edge_torch.generative.layers import model_config
-
 
 def load_safetensors(full_path: str):
   """Loads safetensors into a single state dictionary.
@@ -158,14 +157,22 @@ class ModelLoader:
             f"{self._names.embedding_position}"
         )
     if self._names.lm_head is not None:
-      converted_state["lm_head.weight"] = state.pop(f"{self._names.lm_head}.weight")
+      converted_state["lm_head.weight"] = state.pop(
+          f"{self._names.lm_head}.weight"
+      )
       if model.config.lm_head_use_bias:
-        converted_state["lm_head.bias"] = state.pop(f"{self._names.lm_head}.bias")
+        converted_state["lm_head.bias"] = state.pop(
+            f"{self._names.lm_head}.bias"
+        )
     if self._names.final_norm is not None:
       final_norm_name = self._names.final_norm
-      converted_state["final_norm.weight"] = state.pop(f"{final_norm_name}.weight")
+      converted_state["final_norm.weight"] = state.pop(
+          f"{final_norm_name}.weight"
+      )
       if f"{final_norm_name}.bias" in state:
-        converted_state["final_norm.bias"] = state.pop(f"{final_norm_name}.bias")
+        converted_state["final_norm.bias"] = state.pop(
+            f"{final_norm_name}.bias"
+        )
 
     for i in range(model.config.num_layers):
       self._map_norm(i, model.config, state, converted_state)
@@ -214,18 +221,26 @@ class ModelLoader:
     if config.ff_config.type == model_config.FeedForwardType.SEQUENTIAL:
       ff_up_proj_name = self._names.ff_up_proj.format(idx)
       ff_down_proj_name = self._names.ff_down_proj.format(idx)
-      converted_state[f"{prefix}.ff.w1.weight"] = state.pop(f"{ff_up_proj_name}.weight")
+      converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
+          f"{ff_up_proj_name}.weight"
+      )
       converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
           f"{ff_down_proj_name}.weight"
       )
       if config.ff_config.use_bias:
-        converted_state[f"{prefix}.ff.w1.bias"] = state.pop(f"{ff_up_proj_name}.bias")
-        converted_state[f"{prefix}.ff.w2.bias"] = state.pop(f"{ff_down_proj_name}.bias")
+        converted_state[f"{prefix}.ff.w1.bias"] = state.pop(
+            f"{ff_up_proj_name}.bias"
+        )
+        converted_state[f"{prefix}.ff.w2.bias"] = state.pop(
+            f"{ff_down_proj_name}.bias"
+        )
     else:
       ff_up_proj_name = self._names.ff_up_proj.format(idx)
       ff_down_proj_name = self._names.ff_down_proj.format(idx)
       ff_gate_proj_name = self._names.ff_gate_proj.format(idx)
-      converted_state[f"{prefix}.ff.w3.weight"] = state.pop(f"{ff_up_proj_name}.weight")
+      converted_state[f"{prefix}.ff.w3.weight"] = state.pop(
+          f"{ff_up_proj_name}.weight"
+      )
       converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
           f"{ff_down_proj_name}.weight"
       )
@@ -233,9 +248,15 @@ class ModelLoader:
           f"{ff_gate_proj_name}.weight"
       )
       if config.ff_config.use_bias:
-        converted_state[f"{prefix}.ff.w3.bias"] = state.pop(f"{ff_up_proj_name}.bias")
-        converted_state[f"{prefix}.ff.w2.bias"] = state.pop(f"{ff_down_proj_name}.bias")
-        converted_state[f"{prefix}.ff.w1.bias"] = state.pop(f"{ff_gate_proj_name}.bias")
+        converted_state[f"{prefix}.ff.w3.bias"] = state.pop(
+            f"{ff_up_proj_name}.bias"
+        )
+        converted_state[f"{prefix}.ff.w2.bias"] = state.pop(
+            f"{ff_down_proj_name}.bias"
+        )
+        converted_state[f"{prefix}.ff.w1.bias"] = state.pop(
+            f"{ff_gate_proj_name}.bias"
+        )
 
   def _map_attention(
       self,
@@ -254,11 +275,13 @@ 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.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"),
+      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:
       if self._names.attn_fused_qkv_proj:
@@ -266,20 +289,22 @@ class ModelLoader:
             f"{fused_qkv_name}.bias"
         )
       else:
-        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"),
-            state.pop(f"{v_name}.bias"),
+        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"),
+                state.pop(f"{v_name}.bias"),
+            )
         )
 
     o_name = self._names.attn_output_proj.format(idx)
-    converted_state[f"{prefix}.atten_func.output_projection.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.output_projection.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(
@@ -318,7 +343,9 @@ class ModelLoader:
       v: torch.Tensor,
   ) -> torch.Tensor:
     if config.attn_config.qkv_fused_interleaved:
-      q_per_kv = config.attn_config.num_heads // config.attn_config.num_query_groups
+      q_per_kv = (
+          config.attn_config.num_heads // config.attn_config.num_query_groups
+      )
       qs = torch.split(q, config.head_dim * q_per_kv)
       ks = torch.split(k, config.head_dim)
       vs = torch.split(v, config.head_dim)

ai_edge_torch/generative/utilities/stable_diffusion_loader.py

@@ -16,11 +16,10 @@
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Tuple
 
-import torch
-
 import ai_edge_torch.generative.layers.model_config as layers_config
 import ai_edge_torch.generative.layers.unet.model_config as unet_config
 import ai_edge_torch.generative.utilities.loader as loader
+import torch
 
 
 @dataclass
@@ -80,27 +79,35 @@ class TransformerBlockTensorNames:
 class MidBlockTensorNames:
   residual_block_tensor_names: List[ResidualBlockTensorNames]
   attention_block_tensor_names: Optional[List[AttentionBlockTensorNames]] = None
-  transformer_block_tensor_names: Optional[List[TransformerBlockTensorNames]] = None
+  transformer_block_tensor_names: Optional[
+      List[TransformerBlockTensorNames]
+  ] = None
 
 
 @dataclass
 class DownEncoderBlockTensorNames:
   residual_block_tensor_names: List[ResidualBlockTensorNames]
-  transformer_block_tensor_names: Optional[List[TransformerBlockTensorNames]] = None
+  transformer_block_tensor_names: Optional[
+      List[TransformerBlockTensorNames]
+  ] = None
   downsample_conv: str = None
 
 
 @dataclass
 class UpDecoderBlockTensorNames:
   residual_block_tensor_names: List[ResidualBlockTensorNames]
-  transformer_block_tensor_names: Optional[List[TransformerBlockTensorNames]] = None
+  transformer_block_tensor_names: Optional[
+      List[TransformerBlockTensorNames]
+  ] = None
   upsample_conv: str = None
 
 
 @dataclass
 class SkipUpDecoderBlockTensorNames:
   residual_block_tensor_names: List[ResidualBlockTensorNames]
-  transformer_block_tensor_names: Optional[List[TransformerBlockTensorNames]] = None
+  transformer_block_tensor_names: Optional[
+      List[TransformerBlockTensorNames]
+  ] = None
   upsample_conv: str = None
 
 
@@ -119,7 +126,9 @@ def _map_to_converted_state(
         converted_state[f"{converted_state_param}.weight"]
     )
   if f"{state_param}.bias" in state:
-    converted_state[f"{converted_state_param}.bias"] = state.pop(f"{state_param}.bias")
+    converted_state[f"{converted_state_param}.bias"] = state.pop(
+        f"{state_param}.bias"
+    )
     if squeeze_dims:
       converted_state[f"{converted_state_param}.bias"] = torch.squeeze(
           converted_state[f"{converted_state_param}.bias"]
@@ -220,25 +229,41 @@ class BaseLoader(loader.ModelLoader):
           f"{attention_layer_prefix}.v_projection",
           squeeze_dims=True,
       )
-      converted_state[f"{attention_layer_prefix}.qkv_projection.weight"] = torch.concat(
-          [
-              converted_state[f"{attention_layer_prefix}.q_projection.weight"],
-              converted_state[f"{attention_layer_prefix}.k_projection.weight"],
-              converted_state[f"{attention_layer_prefix}.v_projection.weight"],
-          ],
-          axis=0,
+      converted_state[f"{attention_layer_prefix}.qkv_projection.weight"] = (
+          torch.concat(
+              [
+                  converted_state[
+                      f"{attention_layer_prefix}.q_projection.weight"
+                  ],
+                  converted_state[
+                      f"{attention_layer_prefix}.k_projection.weight"
+                  ],
+                  converted_state[
+                      f"{attention_layer_prefix}.v_projection.weight"
+                  ],
+              ],
+              axis=0,
+          )
       )
       del converted_state[f"{attention_layer_prefix}.q_projection.weight"]
       del converted_state[f"{attention_layer_prefix}.k_projection.weight"]
       del converted_state[f"{attention_layer_prefix}.v_projection.weight"]
       if config.attention_config.qkv_use_bias:
-        converted_state[f"{attention_layer_prefix}.qkv_projection.bias"] = torch.concat(
-            [
-                converted_state[f"{attention_layer_prefix}.q_projection.bias"],
-                converted_state[f"{attention_layer_prefix}.k_projection.bias"],
-                converted_state[f"{attention_layer_prefix}.v_projection.bias"],
-            ],
-            axis=0,
+        converted_state[f"{attention_layer_prefix}.qkv_projection.bias"] = (
+            torch.concat(
+                [
+                    converted_state[
+                        f"{attention_layer_prefix}.q_projection.bias"
+                    ],
+                    converted_state[
+                        f"{attention_layer_prefix}.k_projection.bias"
+                    ],
+                    converted_state[
+                        f"{attention_layer_prefix}.v_projection.bias"
+                    ],
+                ],
+                axis=0,
+            )
         )
         del converted_state[f"{attention_layer_prefix}.q_projection.bias"]
         del converted_state[f"{attention_layer_prefix}.k_projection.bias"]
@@ -316,11 +341,17 @@ class BaseLoader(loader.ModelLoader):
       )
     else:
       _map_to_converted_state(
-          state, tensor_names.w1, converted_state, f"{converted_state_param_prefix}.w1"
+          state,
+          tensor_names.w1,
+          converted_state,
+          f"{converted_state_param_prefix}.w1",
       )
 
     _map_to_converted_state(
-        state, tensor_names.w2, converted_state, f"{converted_state_param_prefix}.w2"
+        state,
+        tensor_names.w2,
+        converted_state,
+        f"{converted_state_param_prefix}.w2",
     )
 
   def _map_transformer_block(
@@ -509,9 +540,13 @@ class BaseLoader(loader.ModelLoader):
   ):
     for i in range(config.num_layers):
       res_skip_channels = (
-          config.in_channels if (i == config.num_layers - 1) else config.out_channels
+          config.in_channels
+          if (i == config.num_layers - 1)
+          else config.out_channels
+      )
+      resnet_in_channels = (
+          config.prev_out_channels if i == 0 else config.out_channels
       )
-      resnet_in_channels = config.prev_out_channels if i == 0 else config.out_channels
       self._map_residual_block(
           state,
           converted_state,
@@ -599,9 +634,13 @@ class AutoEncoderModelLoader(BaseLoader):
           state, self._names.post_quant_conv, converted_state, "post_quant_conv"
       )
     if self._names.conv_in is not None:
-      _map_to_converted_state(state, self._names.conv_in, converted_state, "conv_in")
+      _map_to_converted_state(
+          state, self._names.conv_in, converted_state, "conv_in"
+      )
     if self._names.conv_out is not None:
-      _map_to_converted_state(state, self._names.conv_out, converted_state, "conv_out")
+      _map_to_converted_state(
+          state, self._names.conv_out, converted_state, "conv_out"
+      )
     if self._names.final_norm is not None:
       _map_to_converted_state(
           state, self._names.final_norm, converted_state, "final_norm"
@@ -614,7 +653,9 @@ class AutoEncoderModelLoader(BaseLoader):
         model.config.mid_block_config,
     )
 
-    reversed_block_out_channels = list(reversed(model.config.block_out_channels))
+    reversed_block_out_channels = list(
+        reversed(model.config.block_out_channels)
+    )
     block_out_channels = reversed_block_out_channels[0]
     for i, out_channels in enumerate(reversed_block_out_channels):
       prev_output_channel = block_out_channels
@@ -690,8 +731,12 @@ class DiffusionModelLoader(BaseLoader):
     self._map_time_embedding(
         state, converted_state, "time_embedding", self._names.time_embedding
     )
-    _map_to_converted_state(state, self._names.conv_in, converted_state, "conv_in")
-    _map_to_converted_state(state, self._names.conv_out, converted_state, "conv_out")
+    _map_to_converted_state(
+        state, self._names.conv_in, converted_state, "conv_in"
+    )
+    _map_to_converted_state(
+        state, self._names.conv_out, converted_state, "conv_out"
+    )
     _map_to_converted_state(
         state, self._names.final_norm, converted_state, "final_norm"
     )
@@ -825,7 +870,9 @@ class DiffusionModelLoader(BaseLoader):
     )
 
     # Map up_decoders.
-    reversed_block_out_channels = list(reversed(model.config.block_out_channels))
+    reversed_block_out_channels = list(
+        reversed(model.config.block_out_channels)
+    )
     up_decoder_layers_per_block = config.layers_per_block + 1
     output_channel = reversed_block_out_channels[0]
     for i, block_out_channel in enumerate(reversed_block_out_channels):
@@ -917,8 +964,14 @@ class DiffusionModelLoader(BaseLoader):
       tensor_names: TimeEmbeddingTensorNames,
   ):
     _map_to_converted_state(
-        state, tensor_names.w1, converted_state, f"{converted_state_param_prefix}.w1"
+        state,
+        tensor_names.w1,
+        converted_state,
+        f"{converted_state_param_prefix}.w1",
     )
     _map_to_converted_state(
-        state, tensor_names.w2, converted_state, f"{converted_state_param_prefix}.w2"
+        state,
+        tensor_names.w2,
+        converted_state,
+        f"{converted_state_param_prefix}.w2",
     )