Refactors serialization

Author: xadupre

_doc/api/torch_export_patches/onnx_export_serialization_impl.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.torch_export_patches.serialization.diffusers_impl
+=================================================================
+
+.. automodule:: onnx_diagnostic.torch_export_patches.serialization.diffusers_impl
+    :members:
+    :no-undoc-members:

_doc/api/torch_export_patches/serialization/diffusers_impl.rst

@@ -0,0 +1,13 @@
+onnx_diagnostic.torch_export_patches.serialization
+==================================================
+
+.. toctree::
+    :maxdepth: 1
+    :caption: submodules
+
+    diffusers_impl
+    transformers_impl
+
+.. automodule:: onnx_diagnostic.torch_export_patches.serialization
+    :members:
+    :no-undoc-members:

_doc/api/torch_export_patches/serialization/index.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.torch_export_patches.serialization.transformers_impl
+====================================================================
+
+.. automodule:: onnx_diagnostic.torch_export_patches.serialization.transformers_impl
+    :members:
+    :no-undoc-members:

_doc/api/torch_export_patches/serialization/transformers_impl.rst

@@ -134,11 +134,17 @@ def unpatch_module_or_classes(mod, info: Dict[type, Dict[type, Callable]], verbo
 
 @contextlib.contextmanager
 def register_additional_serialization_functions(
-    patch_transformers: bool = False, verbose: int = 0
+    patch_transformers: bool = False, patch_diffusers: bool = False, verbose: int = 0
 ) -> Callable:
     """The necessary modifications to run the fx Graph."""
-    fct_callable = replacement_before_exporting if patch_transformers else (lambda x: x)
-    done = register_cache_serialization(verbose=verbose)
+    fct_callable = (
+        replacement_before_exporting
+        if patch_transformers or patch_diffusers
+        else (lambda x: x)
+    )
+    done = register_cache_serialization(
+        patch_transformers=patch_transformers, patch_diffusers=patch_diffusers, verbose=verbose
+    )
     try:
         yield fct_callable
     finally:
@@ -150,6 +156,7 @@ def torch_export_patches(
     patch_sympy: bool = True,
     patch_torch: bool = True,
     patch_transformers: bool = False,
+    patch_diffusers: bool = False,
     catch_constraints: bool = True,
     stop_if_static: int = 0,
     verbose: int = 0,
@@ -165,6 +172,7 @@ def torch_export_patches(
     :param patch_sympy: fix missing method ``name`` for IntegerConstant
     :param patch_torch: patches :epkg:`torch` with supported implementation
     :param patch_transformers: patches :epkg:`transformers` with supported implementation
+    :param patch_diffusers: patches :epkg:`diffusers` with supported implementation
     :param catch_constraints: catch constraints related to dynamic shapes,
         as a result, some dynamic dimension may turn into static ones,
         the environment variable ``SKIP_SOLVE_CONSTRAINTS=0``
@@ -249,6 +257,7 @@ def torch_export_patches(
             patch_sympy=patch_sympy,
             patch_torch=patch_torch,
             patch_transformers=patch_transformers,
+            patch_diffusers=patch_diffusers,
             catch_constraints=catch_constraints,
             stop_if_static=stop_if_static,
             verbose=verbose,
@@ -281,7 +290,11 @@ def torch_export_patches(
         # caches
         ########
 
-        cache_done = register_cache_serialization(verbose=verbose)
+        cache_done = register_cache_serialization(
+            patch_transformers=patch_transformers,
+            patch_diffusers=patch_diffusers,
+            verbose=verbose,
+        )
 
         #############
         # patch sympy

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -13,7 +13,7 @@
 )
 
 from ..helpers import string_type
-
+from .serialization import _lower_name_with_
 
 PATCH_OF_PATCHES: Set[Any] = set()
 
@@ -73,14 +73,25 @@ def register_class_serialization(
     return True
 
 
-def register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
+def register_cache_serialization(
+    patch_transformers: bool = False, patch_diffusers: bool = True, verbose: int = 0
+) -> Dict[str, bool]:
     """
     Registers many classes with :func:`register_class_serialization`.
     Returns information needed to undo the registration.
+
+    :param patch_transformers: add serialization function for
+        :epkg:`transformers` package
+    :param patch_diffusers: add serialization function for
+        :epkg:`diffusers` package
+    :param verbosity: verbosity level
+    :return: information to unpatch
     """
     from .onnx_export_serialization_impl import WRONG_REGISTRATIONS
 
-    registration_functions = serialization_functions(verbose=verbose)
+    registration_functions = serialization_functions(
+        patch_transformers=patch_transformers, patch_diffusers=patch_diffusers, verbose=verbose
+    )
 
     # DynamicCache serialization is different in transformers and does not
     # play way with torch.export.export.
@@ -124,68 +135,86 @@ def register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
     return done
 
 
-def serialization_functions(verbose: int = 0) -> Dict[type, Callable[[int], bool]]:
+def serialization_functions(
+    patch_transformers: bool = False, patch_diffusers: bool = False, verbose: int = 0
+) -> Dict[type, Callable[[int], bool]]:
     """Returns the list of serialization functions."""
-    from .onnx_export_serialization_impl import (
-        SUPPORTED_DATACLASSES,
-        _lower_name_with_,
-        __dict__ as all_functions,
-        flatten_dynamic_cache,
-        unflatten_dynamic_cache,
-        flatten_with_keys_dynamic_cache,
-        flatten_mamba_cache,
-        unflatten_mamba_cache,
-        flatten_with_keys_mamba_cache,
-        flatten_encoder_decoder_cache,
-        unflatten_encoder_decoder_cache,
-        flatten_with_keys_encoder_decoder_cache,
-        flatten_sliding_window_cache,
-        unflatten_sliding_window_cache,
-        flatten_with_keys_sliding_window_cache,
-        flatten_static_cache,
-        unflatten_static_cache,
-        flatten_with_keys_static_cache,
-    )
 
-    transformers_classes = {
-        DynamicCache: lambda verbose=verbose: register_class_serialization(
-            DynamicCache,
+    supported_classes = set()
+    classes = {}
+    all_functions = {}
+
+    if patch_transformers:
+        from .serialization.transformers_impl import (
+            __dict__ as dtr,
+            SUPPORTED_DATACLASSES,
             flatten_dynamic_cache,
             unflatten_dynamic_cache,
             flatten_with_keys_dynamic_cache,
-            # f_check=make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
-            verbose=verbose,
-        ),
-        MambaCache: lambda verbose=verbose: register_class_serialization(
-            MambaCache,
             flatten_mamba_cache,
             unflatten_mamba_cache,
             flatten_with_keys_mamba_cache,
-            verbose=verbose,
-        ),
-        EncoderDecoderCache: lambda verbose=verbose: register_class_serialization(
-            EncoderDecoderCache,
             flatten_encoder_decoder_cache,
             unflatten_encoder_decoder_cache,
             flatten_with_keys_encoder_decoder_cache,
-            verbose=verbose,
-        ),
-        SlidingWindowCache: lambda verbose=verbose: register_class_serialization(
-            SlidingWindowCache,
             flatten_sliding_window_cache,
             unflatten_sliding_window_cache,
             flatten_with_keys_sliding_window_cache,
-            verbose=verbose,
-        ),
-        StaticCache: lambda verbose=verbose: register_class_serialization(
-            StaticCache,
             flatten_static_cache,
             unflatten_static_cache,
             flatten_with_keys_static_cache,
-            verbose=verbose,
-        ),
-    }
-    for cls in SUPPORTED_DATACLASSES:
+        )
+
+        all_functions.update(dtr)
+        supported_classes |= SUPPORTED_DATACLASSES
+
+        transformers_classes = {
+            DynamicCache: lambda verbose=verbose: register_class_serialization(
+                DynamicCache,
+                flatten_dynamic_cache,
+                unflatten_dynamic_cache,
+                flatten_with_keys_dynamic_cache,
+                # f_check=make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
+                verbose=verbose,
+            ),
+            MambaCache: lambda verbose=verbose: register_class_serialization(
+                MambaCache,
+                flatten_mamba_cache,
+                unflatten_mamba_cache,
+                flatten_with_keys_mamba_cache,
+                verbose=verbose,
+            ),
+            EncoderDecoderCache: lambda verbose=verbose: register_class_serialization(
+                EncoderDecoderCache,
+                flatten_encoder_decoder_cache,
+                unflatten_encoder_decoder_cache,
+                flatten_with_keys_encoder_decoder_cache,
+                verbose=verbose,
+            ),
+            SlidingWindowCache: lambda verbose=verbose: register_class_serialization(
+                SlidingWindowCache,
+                flatten_sliding_window_cache,
+                unflatten_sliding_window_cache,
+                flatten_with_keys_sliding_window_cache,
+                verbose=verbose,
+            ),
+            StaticCache: lambda verbose=verbose: register_class_serialization(
+                StaticCache,
+                flatten_static_cache,
+                unflatten_static_cache,
+                flatten_with_keys_static_cache,
+                verbose=verbose,
+            ),
+        }
+        classes.update(patch_transformers)
+
+    if patch_diffusers:
+        from .serialization.diffusers_impl import SUPPORTED_DATACLASSES, __dict__ as dfu
+
+        all_functions.update(dfu)
+        supported_classes |= SUPPORTED_DATACLASSES
+
+    for cls in supported_classes:
         lname = _lower_name_with_(cls.__name__)
         assert (
             f"flatten_{lname}" in all_functions

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -0,0 +1,46 @@
+import re
+from typing import Any, Callable, List, Set, Tuple
+import torch
+
+
+def _lower_name_with_(name):
+    s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
+    return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()
+
+
+def make_serialization_function_for_dataclass(
+    cls: type, supported_classes: Set[type]
+) -> Tuple[Callable, Callable, Callable]:
+    """
+    Automatically creates serialization function for a class decorated with
+    ``dataclasses.dataclass``.
+    """
+
+    def flatten_cls(obj: cls) -> Tuple[List[Any], torch.utils._pytree.Context]:
+        """Serializes a ``%s`` with python objects."""
+        return list(obj.values()), list(obj.keys())
+
+    def flatten_with_keys_cls(
+        obj: cls,
+    ) -> Tuple[List[Tuple[torch.utils._pytree.KeyEntry, Any]], torch.utils._pytree.Context]:
+        """Serializes a ``%s`` with python objects with keys."""
+        values, context = list(obj.values()), list(obj.keys())
+        return [
+            (torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)
+        ], context
+
+    def unflatten_cls(
+        values: List[Any], context: torch.utils._pytree.Context, output_type=None
+    ) -> cls:
+        """Restores an instance of ``%s`` from python objects."""
+        return cls(**dict(zip(context, values)))
+
+    name = _lower_name_with_(cls.__name__)
+    flatten_cls.__name__ = f"flatten_{name}"
+    flatten_with_keys_cls.__name__ = f"flatten_with_keys_{name}"
+    unflatten_cls.__name__ = f"unflatten_{name}"
+    flatten_cls.__doc__ = flatten_cls.__doc__ % cls.__name__
+    flatten_with_keys_cls.__doc__ = flatten_with_keys_cls.__doc__ % cls.__name__
+    unflatten_cls.__doc__ = unflatten_cls.__doc__ % cls.__name__
+    supported_classes.add(cls)
+    return flatten_cls, flatten_with_keys_cls, unflatten_cls

onnx_diagnostic/torch_export_patches/serialization/__init__.py

@@ -0,0 +1,34 @@
+from typing import Dict, Optional
+
+try:
+    from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
+except ImportError as e:
+    try:
+        import diffusers
+    except ImportError:
+        diffusers = None
+        UNet2DConditionOutput = None
+    if diffusers:
+        raise e
+
+from . import make_serialization_function_for_dataclass
+
+
+def _make_wrong_registrations() -> Dict[str, Optional[str]]:
+    res = {}
+    for c in [UNet2DConditionOutput]:
+        if c is not None:
+            res[c] = None
+    return res
+
+
+SUPPORTED_DATACLASSES = set()
+WRONG_REGISTRATIONS = _make_wrong_registrations()
+
+
+if UNet2DConditionOutput is not None:
+    (
+        flatten_u_net2_d_condition_output,
+        flatten_with_keys_u_net2_d_condition_output,
+        unflatten_u_net2_d_condition_output,
+    ) = make_serialization_function_for_dataclass(UNet2DConditionOutput, SUPPORTED_DATACLASSES)