ai-edge-torch-nightly 0.2.0.dev20240710__py3-none-any.whl → 0.2.0.dev20240713__py3-none-any.whl

ai_edge_torch/convert/conversion.py

@@ -88,16 +88,14 @@ def convert_signatures(
   _warn_training_modules(signatures)
 
   exported_programs: torch.export.ExportedProgram = [
-      torch.export.export(
-          sig.module, sig.sample_args, dynamic_shapes=sig.dynamic_shapes
-      )
+      torch.export.export(sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes)
       for sig in signatures
   ]
 
   # Apply default fx passes
   exported_programs = list(map(_run_convert_passes, exported_programs))
   shlo_bundles: list[stablehlo.StableHLOModelBundle] = [
-      cutils.exported_program_to_stablehlo_bundle(exported, sig.sample_args)
+      cutils.exported_program_to_stablehlo_bundle(exported, sig.flat_args)
       for exported, sig in zip(exported_programs, signatures)
   ]
 

ai_edge_torch/convert/conversion_utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import collections
 import copy
 from dataclasses import dataclass
 import gc
@@ -22,6 +23,7 @@ import tempfile
 from typing import Any, Dict, Optional, Tuple, Union
 
 import torch
+import torch.utils._pytree as pytree
 from torch_xla import stablehlo
 
 from ai_edge_torch.generative.quantize.ai_edge_quantizer_glue import translate_recipe  # NOQA
@@ -47,8 +49,59 @@ class Signature:
   name: str
   module: torch.nn.Module
   sample_args: tuple[torch.Tensor]
+  sample_kwargs: dict[str, torch.Tensor]
   dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None
 
+  @property
+  def _normalized_sample_args_kwargs(self):
+    args, kwargs = self.sample_args, self.sample_kwargs
+    if args is not None:
+      if not isinstance(args, tuple):
+        # TODO(b/352584188): Check value types
+        raise ValueError("sample_args must be a tuple of torch tensors.")
+    if kwargs is not None:
+      if not isinstance(kwargs, dict) or not all(
+          isinstance(key, str) for key in kwargs.keys()
+      ):
+        # TODO(b/352584188): Check value types
+        raise ValueError("sample_kwargs must be a dict of string to tensor.")
+
+    args = args if args is not None else tuple()
+    kwargs = kwargs if kwargs is not None else {}
+    return args, kwargs
+
+  @property
+  def flat_arg_names(self) -> list[str]:
+    spec = pytree.tree_flatten(self._normalized_sample_args_kwargs)[1]
+    args_spec, kwargs_spec = spec.children_specs
+
+    names = []
+    for i in range(args_spec.num_leaves):
+      names.append(f"args_{i}")
+
+    dict_context = (
+        kwargs_spec.context
+        if kwargs_spec.type is not collections.defaultdict
+        # ignore mismatch of `default_factory` for defaultdict
+        else kwargs_spec.context[1]
+    )
+
+    for name, value_spec in zip(dict_context, kwargs_spec.children_specs):
+      if value_spec.num_leaves == 1:
+        names.append(name)
+      else:
+        # value_spec.num_leaves may be greater than 1 when the value is a (nested)
+        # tuple of tensors. We haven't decided how we should support flattenable
+        # tensor containers as inputs.
+        # TODO(b/352584188): Decide the behavior of tensor container as input (flatten or reject)
+        for i in range(value_spec.num_leaves):
+          names.append(f"{name}_{i}")
+    return names
+
+  @property
+  def flat_args(self) -> tuple[torch.Tensor]:
+    return tuple(pytree.tree_flatten(self._normalized_sample_args_kwargs)[0])
+
 
 def exported_program_to_stablehlo_bundle(
     exported_program: torch.export.ExportedProgram, sample_args: tuple[torch.Tensor]
@@ -189,7 +242,9 @@ def _make_tf_function(
 
 def _make_tf_signature(
     meta: stablehlo.StableHLOFunctionMeta,
+    signature: Signature,
 ) -> list[tf.TensorSpec]:
+  input_names = signature.flat_arg_names
   input_pos_to_spec = {
       loc.position: spec
       for loc, spec in itertools.chain(
@@ -197,9 +252,11 @@ def _make_tf_signature(
       )
       if loc.type_ == stablehlo.VariableType.INPUT_ARG
   }
+  assert len(input_pos_to_spec) == len(input_names)
+
   primitive_type_to_tf_type = {"int": "int32", "float": "float32"}
   ret: list[tf.TensorSpec] = []
-  for i in range(len(input_pos_to_spec)):
+  for i, name in enumerate(input_names):
     spec = input_pos_to_spec[i]
     shape = _get_shape_with_dynamic(spec)
     ret.append(
@@ -208,7 +265,7 @@ def _make_tf_signature(
             dtype=primitive_type_to_tf_type[spec.dtype]
             if spec.dtype in primitive_type_to_tf_type
             else spec.dtype,
-            name=f"args_{i}",
+            name=name,
         )
     )
   return ret
@@ -276,7 +333,8 @@ def convert_stablehlo_to_tflite(
       tf.Variable(v, trainable=False) for v in bundle.additional_constants
   ]
   tf_signatures: list[list[tf.TensorSpec]] = list(
-      _make_tf_signature(func.meta) for func in bundle.stablehlo_funcs
+      _make_tf_signature(func.meta, sig)
+      for func, sig in zip(bundle.stablehlo_funcs, signatures)
   )
 
   tf_functions = _make_tf_function(shlo_graph_module, bundle)

ai_edge_torch/convert/converter.py

@@ -34,17 +34,23 @@ class Converter:
       self,
       name: str,
       module: torch.nn.Module,
-      sample_args: tuple[cutils.TracingArg],
+      sample_args=None,
+      sample_kwargs=None,
+      *,
       dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
   ) -> Converter:
     """Alias to `add_signature`"""
-    return self.add_signature(name, module, sample_args, dynamic_shapes)
+    return self.add_signature(
+        name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+    )
 
   def add_signature(
       self,
       name: str,
       module: torch.nn.Module,
-      sample_args: tuple[cutils.TracingArg],
+      sample_args=None,
+      sample_kwargs=None,
+      *,
       dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
   ) -> Converter:
     """Allows adding a new named torch model along with sample args to the conversion.
@@ -52,7 +58,8 @@ class Converter:
     Args:
       name: The name of the signature included in the converted edge model.
       module: The torch module to be converted.
-      sample_args: Tuple of args by which the torch module will be traced prior to conversion.
+      sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+      sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
       dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
         See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
 
@@ -63,13 +70,21 @@ class Converter:
     if name in [sig.name for sig in self._signatures]:
       raise ValueError(f"A signature with the provided name ({name}) is already added.")
 
-    self._signatures.append(cutils.Signature(name, module, sample_args, dynamic_shapes))
+    if sample_args is None and sample_kwargs is None:
+      raise ValueError("sample_args or sample_kwargs must be provided.")
+
+    self._signatures.append(
+        cutils.Signature(
+            name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+        )
+    )
     return self
 
   def convert(
       self,
       module: torch.nn.Module = None,
-      sample_args: tuple[cutils.TracingArg] = None,
+      sample_args=None,
+      sample_kwargs=None,
       *,
       quant_config: Optional[qcfg.QuantConfig] = None,
       dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
@@ -88,7 +103,8 @@ class Converter:
     Args:
       name: The name of the signature included in the converted edge model.
       module: The torch module to be converted.
-      sample_args: Tuple of args by which the torch module will be traced prior to conversion.
+      sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+      sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
       quant_config: User-defined quantization method and scheme of the model.
       dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
         See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
@@ -100,12 +116,20 @@ class Converter:
       ValueError: If the arguments are not provided as expected. See the example in this functions's comment.
     """
     if module is not None:
-      if sample_args is not None:  # both module and args provided
+      if (
+          sample_args is not None or sample_kwargs is not None
+      ):  # both module and args provided
         self.add_signature(
-            cutils.DEFAULT_SIGNATURE_NAME, module, sample_args, dynamic_shapes
+            cutils.DEFAULT_SIGNATURE_NAME,
+            module,
+            sample_args,
+            sample_kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+      else:  # module is provided but not args
+        raise ValueError(
+            "sample_args or sample_kwargs must be provided if a module is specified."
         )
-      else:  # module is provided but not sample_args
-        raise ValueError("sample_args needs to be provided if a module is specified.")
 
     return conversion.convert_signatures(
         self._signatures,
@@ -117,7 +141,8 @@ class Converter:
 def signature(
     name: str,
     module: torch.nn.Module,
-    sample_args: tuple[cutils.TracingArg],
+    sample_args=None,
+    sample_kwargs=None,
     dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
 ) -> Converter:
   """Initiates a Converter object with the provided signature.
@@ -125,7 +150,8 @@ def signature(
   Args:
     name: The name of the signature included in the converted edge model.
     module: The torch module to be converted.
-    sample_args: Tuple of args by which the torch module will be traced prior to conversion.
+    sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+    sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
     dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
       See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
 
@@ -134,12 +160,15 @@ def signature(
     edge_model = converter.convert()
 
   """
-  return Converter().signature(name, module, sample_args, dynamic_shapes)
+  return Converter().signature(
+      name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+  )
 
 
 def convert(
     module: torch.nn.Module = None,
-    sample_args: tuple[cutils.TracingArg] = None,
+    sample_args=None,
+    sample_kwargs=None,
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
@@ -149,7 +178,8 @@ def convert(
 
   Args:
     module: The torch module to be converted.
-    sample_args: Tuple of args by which the torch module will be traced prior to conversion.
+    sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+    sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
     quant_config: User-defined quantization method and scheme of the model.
     dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
       See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
@@ -165,6 +195,7 @@ def convert(
   return Converter().convert(
       module,
       sample_args,
+      sample_kwargs,
       quant_config=quant_config,
       dynamic_shapes=dynamic_shapes,
       _ai_edge_converter_flags=_ai_edge_converter_flags,

ai_edge_torch/convert/test/test_convert.py

@@ -267,6 +267,45 @@ class TestConvert(unittest.TestCase):
           model_coverage.compare_tflite_torch(edge_model, model, validate_input)
       )
 
+  def test_convert_model_with_kwargs(self):
+    """
+    Test converting a simple model with sample_kwargs.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x, y):
+        return x + y
+
+    kwargs_gen = lambda: dict(x=torch.randn(10, 10), y=torch.randn(10, 10))
+
+    model = SampleModel().eval()
+    edge_model = ai_edge_torch.convert(model, sample_kwargs=kwargs_gen())
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, model, kwargs=kwargs_gen)
+    )
+
+  def test_convert_model_with_args_kwargs(self):
+    """
+    Test converting a simple model with both sample_args and sample_kwargs.
+    """
+
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x, y):
+        return x + y
+
+    args_gen = lambda: (torch.randn(10, 10),)
+    kwargs_gen = lambda: dict(y=torch.randn(10, 10))
+
+    model = SampleModel().eval()
+    edge_model = ai_edge_torch.convert(model, args_gen(), kwargs_gen())
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, model, args_gen, kwargs_gen)
+    )
+
 
 if __name__ == "__main__":
   unittest.main()

ai_edge_torch/generative/examples/stable_diffusion/clip.py

@@ -23,16 +23,17 @@ import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
-    ff_up_proj="layers.{}.linear_1",
-    ff_down_proj="layers.{}.linear_2",
-    ff_gate_proj="layers.{}.linear_1",
-    attn_fused_qkv_proj="layers.{}.attention.in_proj",
-    attn_output_proj="layers.{}.attention.out_proj",
-    pre_attn_norm="layers.{}.layernorm_1",
-    pre_ff_norm="layers.{}.layernorm_2",
-    embedding="embedding.token_embedding",
-    embedding_position="embedding.position_value",
-    final_norm="layernorm",
+    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",
+    pre_ff_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,
 )
 
@@ -84,6 +85,7 @@ def get_model_config() -> cfg.ModelConfig:
       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,
   )

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 
+import argparse
 import os
 from pathlib import Path
+from typing import Optional
 
 import torch
 
@@ -24,14 +26,36 @@ import ai_edge_torch.generative.examples.stable_diffusion.decoder as decoder
 import ai_edge_torch.generative.examples.stable_diffusion.diffusion as diffusion
 from ai_edge_torch.generative.examples.stable_diffusion.encoder import Encoder
 import ai_edge_torch.generative.examples.stable_diffusion.util as util
-import ai_edge_torch.generative.utilities.loader as loading_utils
 import ai_edge_torch.generative.utilities.stable_diffusion_loader as stable_diffusion_loader
 
+arg_parser = argparse.ArgumentParser()
+arg_parser.add_argument(
+    '--clip_ckpt', type=str, help='Path to source CLIP model checkpoint', required=True
+)
+arg_parser.add_argument(
+    '--diffusion_ckpt',
+    type=str,
+    help='Path to source diffusion model checkpoint',
+    required=True,
+)
+arg_parser.add_argument(
+    '--decoder_ckpt',
+    type=str,
+    help='Path to source image decoder model checkpoint',
+    required=True,
+)
+arg_parser.add_argument(
+    '--output_dir',
+    type=str,
+    help='Path to the converted TF Lite directory.',
+    required=True,
+)
+
 
 @torch.inference_mode
 def convert_stable_diffusion_to_tflite(
+    output_dir: str,
     clip_ckpt_path: str,
-    encoder_ckpt_path: str,
     diffusion_ckpt_path: str,
     decoder_ckpt_path: str,
     image_height: int = 512,
@@ -39,23 +63,28 @@ def convert_stable_diffusion_to_tflite(
 ):
 
   clip_model = clip.CLIP(clip.get_model_config())
-  loader = loading_utils.ModelLoader(clip_ckpt_path, clip.TENSOR_NAMES)
+  loader = stable_diffusion_loader.ClipModelLoader(
+      clip_ckpt_path,
+      clip.TENSOR_NAMES,
+  )
   loader.load(clip_model, strict=False)
 
-  encoder = Encoder()
-  encoder.load_state_dict(torch.load(encoder_ckpt_path))
-
   diffusion_model = diffusion.Diffusion(diffusion.get_model_config(2))
   diffusion_loader = stable_diffusion_loader.DiffusionModelLoader(
-      diffusion_ckpt_path, diffusion.TENSORS_NAMES
+      diffusion_ckpt_path, diffusion.TENSOR_NAMES
   )
-  diffusion_loader.load(diffusion_model)
+  diffusion_loader.load(diffusion_model, strict=False)
 
   decoder_model = decoder.Decoder(decoder.get_model_config())
   decoder_loader = stable_diffusion_loader.AutoEncoderModelLoader(
-      decoder_ckpt_path, decoder.TENSORS_NAMES
+      decoder_ckpt_path, decoder.TENSOR_NAMES
   )
-  decoder_loader.load(decoder_model)
+  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
@@ -67,50 +96,47 @@ def convert_stable_diffusion_to_tflite(
       (len_prompt, 4, image_height // 8, image_width // 8), 0, dtype=torch.float32
   )
 
-  input_latents = encoder(input_image, noise)
+  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):
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+  # TODO(yichunk): convert to multi signature tflite model.
   # CLIP text encoder
   ai_edge_torch.signature('encode', clip_model, (prompt_tokens,)).convert().export(
-      '/tmp/stable_diffusion/clip.tflite'
+      f'{output_dir}/clip.tflite'
   )
 
-  # TODO(yichunk): convert to multi signature tflite model.
+  # TODO(yichunk): enable image encoder conversion
   # Image encoder
-  ai_edge_torch.signature('encode', encoder, (input_image, noise)).convert().export(
-      '/tmp/stable_diffusion/encoder.tflite'
-  )
+  # ai_edge_torch.signature('encode', encoder, (input_image, noise)).convert().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().export('/tmp/stable_diffusion/diffusion.tflite')
+  ).convert().export(f'{output_dir}/diffusion.tflite')
 
   # Image decoder
   ai_edge_torch.signature('decode', decoder_model, (input_latents,)).convert().export(
-      '/tmp/stable_diffusion/decoder.tflite'
+      f'{output_dir}/decoder.tflite'
   )
 
 
 if __name__ == '__main__':
+  args = arg_parser.parse_args()
   convert_stable_diffusion_to_tflite(
-      clip_ckpt_path=os.path.join(
-          Path.home(), 'Downloads/stable_diffusion_data/ckpt/clip.pt'
-      ),
-      encoder_ckpt_path=os.path.join(
-          Path.home(), 'Downloads/stable_diffusion_data/ckpt/encoder.pt'
-      ),
-      diffusion_ckpt_path=os.path.join(
-          Path.home(), 'Downloads/stable_diffusion_data/ckpt/diffusion.pt'
-      ),
-      decoder_ckpt_path=os.path.join(
-          Path.home(), 'Downloads/stable_diffusion_data/ckpt/decoder.pt'
-      ),
+      output_dir=args.output_dir,
+      clip_ckpt_path=args.clip_ckpt,
+      diffusion_ckpt_path=args.diffusion_ckpt,
+      decoder_ckpt_path=args.decoder_ckpt,
       image_height=512,
       image_width=512,
   )