[Core] Introduce class variants for Transformer2DModel (#7647)

* init for patches

* finish patched model.

* continuous transformer

* vectorized transformer2d.

* style.

* inits.

* fix-copies.

* introduce DiTTransformer2DModel.

* fixes

* use REMAPPING as suggested by @DN6

* better logging.

* add pixart transformer model.

* inits.

* caption_channels.

* attention masking.

* fix use_additional_conditions.

* remove print.

* debug

* flatten

* fix: assertion for sigma

* handle remapping for modeling_utils

* add tests for dit transformer2d

* quality

* placeholder for pixart tests

* pixart tests

* add _no_split_modules

* add docs.

* check

* check

* check

* check

* fix tests

* fix tests

* move Transformer output to modeling_output

* move errors better and bring back use_additional_conditions attribute.

* add unnecessary things from DiT.

* clean up pixart

* fix remapping

* fix device_map things in pixart2d.

* replace Transformer2DModel with appropriate classes in dit, pixart tests

* empty

* legacy mixin classes./

* use a remapping dict for fetching class names.

* change to specifc model types in the pipeline implementations.

* move _fetch_remapped_cls_from_config to modeling_loading_utils.py

* fix dependency problems.

* add deprecation note.

Author: sayakpaul

docs/source/en/_toctree.yml

@@ -238,6 +238,10 @@
       title: ConsistencyDecoderVAE
     - local: api/models/transformer2d
       title: Transformer2D
+    - local: api/models/pixart_transformer2d
+      title: PixArtTransformer2D
+    - local: api/models/dit_transformer2d
+      title: DiTTransformer2D
     - local: api/models/transformer_temporal
       title: Transformer Temporal
     - local: api/models/prior_transformer

docs/source/en/api/models/dit_transformer2d.md

@@ -0,0 +1,19 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+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.
+-->
+
+# DiTTransformer2D
+
+A Transformer model for image-like data from [DiT](https://huggingface.co/papers/2212.09748).
+
+## DiTTransformer2DModel
+
+[[autodoc]] DiTTransformer2DModel

docs/source/en/api/models/pixart_transformer2d.md

@@ -0,0 +1,19 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+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.
+-->
+
+# PixArtTransformer2D
+
+A Transformer model for image-like data from [PixArt-Alpha](https://huggingface.co/papers/2310.00426) and [PixArt-Sigma](https://huggingface.co/papers/2403.04692). 
+
+## PixArtTransformer2DModel
+
+[[autodoc]] PixArtTransformer2DModel

src/diffusers/__init__.py

@@ -82,11 +82,13 @@
             "ConsistencyDecoderVAE",
             "ControlNetModel",
             "ControlNetXSAdapter",
+            "DiTTransformer2DModel",
             "I2VGenXLUNet",
             "Kandinsky3UNet",
             "ModelMixin",
             "MotionAdapter",
             "MultiAdapter",
+            "PixArtTransformer2DModel",
             "PriorTransformer",
             "StableCascadeUNet",
             "T2IAdapter",
@@ -484,11 +486,13 @@
             ConsistencyDecoderVAE,
             ControlNetModel,
             ControlNetXSAdapter,
+            DiTTransformer2DModel,
             I2VGenXLUNet,
             Kandinsky3UNet,
             ModelMixin,
             MotionAdapter,
             MultiAdapter,
+            PixArtTransformer2DModel,
             PriorTransformer,
             T2IAdapter,
             T5FilmDecoder,

src/diffusers/configuration_utils.py

@@ -706,3 +706,20 @@ def init(self, *args, **kwargs):
 
     cls.__init__ = init
     return cls
+
+
+class LegacyConfigMixin(ConfigMixin):
+    r"""
+    A subclass of `ConfigMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more
+    pipeline-specific classes (like `DiTTransformer2DModel`).
+    """
+
+    @classmethod
+    def from_config(cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs):
+        # To prevent depedency import problem.
+        from .models.model_loading_utils import _fetch_remapped_cls_from_config
+
+        # resolve remapping
+        remapped_class = _fetch_remapped_cls_from_config(config, cls)
+
+        return remapped_class.from_config(config, return_unused_kwargs, **kwargs)

src/diffusers/models/__init__.py

@@ -37,6 +37,8 @@
     _import_structure["dual_transformer_2d"] = ["DualTransformer2DModel"]
     _import_structure["embeddings"] = ["ImageProjection"]
     _import_structure["modeling_utils"] = ["ModelMixin"]
+    _import_structure["transformers.dit_transformer_2d"] = ["DiTTransformer2DModel"]
+    _import_structure["transformers.pixart_transformer_2d"] = ["PixArtTransformer2DModel"]
     _import_structure["transformers.prior_transformer"] = ["PriorTransformer"]
     _import_structure["transformers.t5_film_transformer"] = ["T5FilmDecoder"]
     _import_structure["transformers.transformer_2d"] = ["Transformer2DModel"]
@@ -74,7 +76,9 @@
         from .embeddings import ImageProjection
         from .modeling_utils import ModelMixin
         from .transformers import (
+            DiTTransformer2DModel,
             DualTransformer2DModel,
+            PixArtTransformer2DModel,
             PriorTransformer,
             T5FilmDecoder,
             Transformer2DModel,

src/diffusers/models/model_loading_utils.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import inspect
 import os
 from collections import OrderedDict
@@ -32,6 +33,13 @@
 
 logger = logging.get_logger(__name__)
 
+_CLASS_REMAPPING_DICT = {
+    "Transformer2DModel": {
+        "ada_norm_zero": "DiTTransformer2DModel",
+        "ada_norm_single": "PixArtTransformer2DModel",
+    }
+}
+
 
 if is_accelerate_available():
     from accelerate import infer_auto_device_map
@@ -61,6 +69,22 @@ def _determine_device_map(model: torch.nn.Module, device_map, max_memory, torch_
     return device_map
 
 
+def _fetch_remapped_cls_from_config(config, old_class):
+    previous_class_name = old_class.__name__
+    remapped_class_name = _CLASS_REMAPPING_DICT.get(previous_class_name).get(config["norm_type"])
+
+    # load diffusers library to import compatible and original scheduler
+    diffusers_library = importlib.import_module(__name__.split(".")[0])
+    remapped_class = getattr(diffusers_library, remapped_class_name)
+    logger.info(
+        f"Changing class object to be of `{remapped_class_name}` type from `{previous_class_name}` type."
+        "This is because `previous_class_name` is scheduled to be deprecated in a future version. Note that this"
+        " DOESN'T affect the final results."
+    )
+
+    return remapped_class
+
+
 def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[str] = None):
     """
     Reads a checkpoint file, returning properly formatted errors if they arise.

src/diffusers/models/modeling_outputs.py

@@ -15,3 +15,17 @@ class AutoencoderKLOutput(BaseOutput):
     """
 
     latent_dist: "DiagonalGaussianDistribution"  # noqa: F821
+
+
+@dataclass
+class Transformer2DModelOutput(BaseOutput):
+    """
+    The output of [`Transformer2DModel`].
+
+    Args:
+        sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
+            The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
+            distributions for the unnoised latent pixels.
+    """
+
+    sample: "torch.Tensor"  # noqa: F821

src/diffusers/models/modeling_utils.py

@@ -42,7 +42,11 @@
     is_torch_version,
     logging,
 )
-from ..utils.hub_utils import PushToHubMixin, load_or_create_model_card, populate_model_card
+from ..utils.hub_utils import (
+    PushToHubMixin,
+    load_or_create_model_card,
+    populate_model_card,
+)
 from .model_loading_utils import (
     _determine_device_map,
     _load_state_dict_into_model,
@@ -1039,3 +1043,55 @@ def recursive_find_attn_block(module) -> None:
             del module.key
             del module.value
             del module.proj_attn
+
+
+class LegacyModelMixin(ModelMixin):
+    r"""
+    A subclass of `ModelMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more
+    pipeline-specific classes (like `DiTTransformer2DModel`).
+    """
+
+    @classmethod
+    @validate_hf_hub_args
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
+        # To prevent depedency import problem.
+        from .model_loading_utils import _fetch_remapped_cls_from_config
+
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        resume_download = kwargs.pop("resume_download", None)
+        proxies = kwargs.pop("proxies", None)
+        local_files_only = kwargs.pop("local_files_only", None)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", None)
+
+        # Load config if we don't provide a configuration
+        config_path = pretrained_model_name_or_path
+
+        user_agent = {
+            "diffusers": __version__,
+            "file_type": "model",
+            "framework": "pytorch",
+        }
+
+        # load config
+        config, _, _ = cls.load_config(
+            config_path,
+            cache_dir=cache_dir,
+            return_unused_kwargs=True,
+            return_commit_hash=True,
+            force_download=force_download,
+            resume_download=resume_download,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            subfolder=subfolder,
+            user_agent=user_agent,
+            **kwargs,
+        )
+        # resolve remapping
+        remapped_class = _fetch_remapped_cls_from_config(config, cls)
+
+        return remapped_class.from_pretrained(pretrained_model_name_or_path, **kwargs)

src/diffusers/models/transformers/__init__.py

@@ -2,7 +2,9 @@
 
 
 if is_torch_available():
+    from .dit_transformer_2d import DiTTransformer2DModel
     from .dual_transformer_2d import DualTransformer2DModel
+    from .pixart_transformer_2d import PixArtTransformer2DModel
     from .prior_transformer import PriorTransformer
     from .t5_film_transformer import T5FilmDecoder
     from .transformer_2d import Transformer2DModel