[Attention] Finish refactor attention file (huggingface#1879)

* [Attention] Finish refactor attention file

* correct more

* fix

* more fixes

* correct

* up

Author: patrickvonplaten

docs/source/api/models.mdx

@@ -56,7 +56,7 @@ The models are built on the base class ['ModelMixin'] that is a `torch.nn.module
 [[autodoc]] Transformer2DModel
 
 ## Transformer2DModelOutput
-[[autodoc]] models.attention.Transformer2DModelOutput
+[[autodoc]] models.transformer_2d.Transformer2DModelOutput
 
 ## PriorTransformer
 [[autodoc]] models.prior_transformer.PriorTransformer

src/diffusers/models/attention.py

@@ -119,9 +119,9 @@ def forward(
                 Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
 
         Returns:
-            [`~models.attention.Transformer2DModelOutput`] or `tuple`: [`~models.attention.Transformer2DModelOutput`]
-            if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is the sample
-            tensor.
+            [`~models.transformer_2d.Transformer2DModelOutput`] or `tuple`:
+            [`~models.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is the sample tensor.
         """
         input_states = hidden_states
 

src/diffusers/models/dual_transformer_2d.py

@@ -189,9 +189,9 @@ def forward(
                 Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
 
         Returns:
-            [`~models.attention.Transformer2DModelOutput`] or `tuple`: [`~models.attention.Transformer2DModelOutput`]
-            if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is the sample
-            tensor.
+            [`~models.transformer_2d.Transformer2DModelOutput`] or `tuple`:
+            [`~models.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is the sample tensor.
         """
         # 1. Input
         if self.is_input_continuous:

src/diffusers/models/transformer_2d.py

@@ -15,9 +15,11 @@
 import torch
 from torch import nn
 
-from .attention import AttentionBlock, DualTransformer2DModel, Transformer2DModel
+from .attention import AttentionBlock
 from .cross_attention import CrossAttention, CrossAttnAddedKVProcessor
+from .dual_transformer_2d import DualTransformer2DModel
 from .resnet import Downsample2D, FirDownsample2D, FirUpsample2D, ResnetBlock2D, Upsample2D
+from .transformer_2d import Transformer2DModel
 
 
 def get_down_block(

src/diffusers/models/unet_2d_blocks.py

@@ -6,9 +6,11 @@
 
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...models import ModelMixin
-from ...models.attention import CrossAttention, DualTransformer2DModel, Transformer2DModel
+from ...models.attention import CrossAttention
 from ...models.cross_attention import AttnProcessor, CrossAttnAddedKVProcessor
+from ...models.dual_transformer_2d import DualTransformer2DModel
 from ...models.embeddings import TimestepEmbedding, Timesteps
+from ...models.transformer_2d import Transformer2DModel
 from ...models.unet_2d_condition import UNet2DConditionOutput
 from ...utils import logging
 

src/diffusers/pipelines/versatile_diffusion/modeling_text_unet.py

@@ -20,9 +20,10 @@
 import torch
 from torch import nn
 
-from diffusers.models.attention import GEGLU, AdaLayerNorm, ApproximateGELU, AttentionBlock, Transformer2DModel
+from diffusers.models.attention import GEGLU, AdaLayerNorm, ApproximateGELU, AttentionBlock
 from diffusers.models.embeddings import get_timestep_embedding
 from diffusers.models.resnet import Downsample2D, Upsample2D
+from diffusers.models.transformer_2d import Transformer2DModel
 from diffusers.utils import torch_device