[Refactor] Move transformer blocks out of attention module

src/diffusers/hooks/_helpers.py

@@ -151,8 +151,8 @@ def _register_attention_processors_metadata():
 
 
 def _register_transformer_blocks_metadata():
-    from ..models.attention import BasicTransformerBlock
     from ..models.transformers.cogvideox_transformer_3d import CogVideoXBlock
+    from ..models.transformers.transformer_2d import BasicTransformerBlock
     from ..models.transformers.transformer_bria import BriaTransformerBlock
     from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock
     from ..models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock

src/diffusers/models/controlnets/controlnet_flux.py

@@ -24,8 +24,8 @@
 from ..attention_processor import AttentionProcessor
 from ..controlnets.controlnet import ControlNetConditioningEmbedding, zero_module
 from ..embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed
-from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
+from ..transformers.modeling_common import Transformer2DModelOutput
 from ..transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock
 
 

src/diffusers/models/controlnets/controlnet_qwenimage.py

@@ -24,8 +24,8 @@
 from ..attention_processor import AttentionProcessor
 from ..cache_utils import CacheMixin
 from ..controlnets.controlnet import zero_module
-from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
+from ..transformers.modeling_common import Transformer2DModelOutput
 from ..transformers.transformer_qwenimage import (
     QwenEmbedRope,
     QwenImageTransformerBlock,

src/diffusers/models/controlnets/controlnet_sana.py

@@ -23,9 +23,9 @@
 from ...utils import USE_PEFT_BACKEND, BaseOutput, logging, scale_lora_layers, unscale_lora_layers
 from ..attention_processor import AttentionProcessor
 from ..embeddings import PatchEmbed, PixArtAlphaTextProjection
-from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormSingle, RMSNorm
+from ..transformers.modeling_common import Transformer2DModelOutput
 from ..transformers.sana_transformer import SanaTransformerBlock
 from .controlnet import zero_module
 

src/diffusers/models/controlnets/controlnet_sd3.py

@@ -22,12 +22,11 @@
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
 from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
-from ..attention import JointTransformerBlock
 from ..attention_processor import Attention, AttentionProcessor, FusedJointAttnProcessor2_0
 from ..embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
-from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
-from ..transformers.transformer_sd3 import SD3SingleTransformerBlock
+from ..transformers.modeling_common import Transformer2DModelOutput
+from ..transformers.transformer_sd3 import SD3SingleTransformerBlock, SD3TransformerBlock
 from .controlnet import BaseOutput, zero_module
 
 
@@ -132,7 +131,7 @@ def __init__(
             # It needs to crafted when we get the actual checkpoints.
             self.transformer_blocks = nn.ModuleList(
                 [
-                    JointTransformerBlock(
+                    SD3TransformerBlock(
                         dim=self.inner_dim,
                         num_attention_heads=num_attention_heads,
                         attention_head_dim=attention_head_dim,

src/diffusers/models/embeddings.py

@@ -1530,7 +1530,7 @@ def forward(self, image_embeds: torch.Tensor):
 class IPAdapterFullImageProjection(nn.Module):
     def __init__(self, image_embed_dim=1024, cross_attention_dim=1024):
         super().__init__()
-        from .attention import FeedForward
+        from .transformers.modeling_common import FeedForward
 
         self.ff = FeedForward(image_embed_dim, cross_attention_dim, mult=1, activation_fn="gelu")
         self.norm = nn.LayerNorm(cross_attention_dim)
@@ -1542,7 +1542,7 @@ def forward(self, image_embeds: torch.Tensor):
 class IPAdapterFaceIDImageProjection(nn.Module):
     def __init__(self, image_embed_dim=1024, cross_attention_dim=1024, mult=1, num_tokens=1):
         super().__init__()
-        from .attention import FeedForward
+        from .transformers.modeling_common import FeedForward
 
         self.num_tokens = num_tokens
         self.cross_attention_dim = cross_attention_dim
@@ -2219,7 +2219,7 @@ def __init__(
         ffn_ratio: float = 4,
     ) -> None:
         super().__init__()
-        from .attention import FeedForward
+        from .transformers.modeling_common import FeedForward
 
         self.ln0 = nn.LayerNorm(embed_dims)
         self.ln1 = nn.LayerNorm(embed_dims)
@@ -2334,7 +2334,7 @@ def __init__(
         ffproj_ratio: int = 2,
     ) -> None:
         super().__init__()
-        from .attention import FeedForward
+        from .transformers.modeling_common import FeedForward
 
         self.num_tokens = num_tokens
         self.embed_dim = embed_dims
@@ -2404,7 +2404,7 @@ def __init__(
         ffn_ratio: int = 4,
     ) -> None:
         super().__init__()
-        from .attention import FeedForward
+        from .transformers.modeling_common import FeedForward
 
         self.ln0 = nn.LayerNorm(hidden_dim)
         self.ln1 = nn.LayerNorm(hidden_dim)

src/diffusers/models/modeling_outputs.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 
-from ..utils import BaseOutput
+from ..utils import BaseOutput, deprecate
 
 
 @dataclass
@@ -17,8 +17,7 @@ class AutoencoderKLOutput(BaseOutput):
     latent_dist: "DiagonalGaussianDistribution"  # noqa: F821
 
 
-@dataclass
-class Transformer2DModelOutput(BaseOutput):
+class Transformer2DModelOutput:
     """
     The output of [`Transformer2DModel`].
 
@@ -28,4 +27,13 @@ class Transformer2DModelOutput(BaseOutput):
             distributions for the unnoised latent pixels.
     """
 
-    sample: "torch.Tensor"  # noqa: F821
+    def __new__(cls, *args, **kwargs):
+        deprecate(
+            "Transformer2DModelOutput",
+            "1.0.0",
+            "Importing `Transformer2DModelOutput` from `diffusers.models.modeling_outputs` is deprecated. Please use `from diffusers.models.transformers.modeling_common import Transformer2DModelOutput` instead.",
+            standard_warn=False,
+        )
+        from .transformers.modeling_common import Transformer2DModelOutput
+
+        return Transformer2DModelOutput(*args, **kwargs)

src/diffusers/models/transformers/auraflow_transformer_2d.py

@@ -30,9 +30,9 @@
     FusedAuraFlowAttnProcessor2_0,
 )
 from ..embeddings import TimestepEmbedding, Timesteps
-from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormZero, FP32LayerNorm
+from .modeling_common import Transformer2DModelOutput
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -194,7 +194,8 @@ def forward(
 
 
 @maybe_allow_in_graph
-class AuraFlowJointTransformerBlock(nn.Module):
+# Copied from diffusers.models.transformers.transformer_sd3.SD3TransformerBlock with SD3->AuraFlow
+class AuraFlowTransformerBlock(nn.Module):
     r"""
     Transformer block for Aura Flow. Similar to SD3 MMDiT. Differences (non-exhaustive):
 
@@ -337,7 +338,7 @@ def __init__(
 
         self.joint_transformer_blocks = nn.ModuleList(
             [
-                AuraFlowJointTransformerBlock(
+                AuraFlowTransformerBlock(
                     dim=self.inner_dim,
                     num_attention_heads=self.config.num_attention_heads,
                     attention_head_dim=self.config.attention_head_dim,

src/diffusers/models/transformers/cogvideox_transformer_3d.py

@@ -22,13 +22,13 @@
 from ...loaders import PeftAdapterMixin
 from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
 from ...utils.torch_utils import maybe_allow_in_graph
-from ..attention import Attention, FeedForward
+from ..attention import Attention
 from ..attention_processor import AttentionProcessor, CogVideoXAttnProcessor2_0, FusedCogVideoXAttnProcessor2_0
 from ..cache_utils import CacheMixin
 from ..embeddings import CogVideoXPatchEmbed, TimestepEmbedding, Timesteps
-from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNorm, CogVideoXLayerNormZero
+from .modeling_common import FeedForward, Transformer2DModelOutput
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

src/diffusers/models/transformers/consisid_transformer_3d.py

@@ -22,12 +22,12 @@
 from ...loaders import PeftAdapterMixin
 from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
 from ...utils.torch_utils import maybe_allow_in_graph
-from ..attention import Attention, FeedForward
+from ..attention import Attention
 from ..attention_processor import AttentionProcessor, CogVideoXAttnProcessor2_0
 from ..embeddings import CogVideoXPatchEmbed, TimestepEmbedding, Timesteps
-from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNorm, CogVideoXLayerNormZero
+from .modeling_common import FeedForward, Transformer2DModelOutput
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name