Merge branch 'main' into transformer2d-as-blocks

src/diffusers/models/controlnet_xs.py

@@ -33,12 +33,12 @@ from .attention_processor import (
 from .controlnet import ControlNetConditioningEmbedding
 from .embeddings import TimestepEmbedding, Timesteps
 from .modeling_utils import ModelMixin
+from .transformers import ContinuousTransformer2DModelBlock
 from .unets.unet_2d_blocks import (
     CrossAttnDownBlock2D,
     CrossAttnUpBlock2D,
     Downsample2D,
     ResnetBlock2D,
-    Transformer2DModel,
     UNetMidBlock2DCrossAttn,
     Upsample2D,
 )
@@ -147,7 +147,7 @@ def get_down_block_adapter(
 
         if has_crossattn:
             attentions.append(
-                Transformer2DModel(
+                ContinuousTransformer2DModelBlock(
                     num_attention_heads,
                     ctrl_out_channels // num_attention_heads,
                     in_channels=ctrl_out_channels,
@@ -1281,7 +1281,7 @@ class ControlNetXSCrossAttnDownBlock2D(nn.Module):
 
             if has_crossattn:
                 base_attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         base_num_attention_heads,
                         base_out_channels // base_num_attention_heads,
                         in_channels=base_out_channels,
@@ -1293,7 +1293,7 @@ class ControlNetXSCrossAttnDownBlock2D(nn.Module):
                     )
                 )
                 ctrl_attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         ctrl_num_attention_heads,
                         ctrl_out_channels // ctrl_num_attention_heads,
                         in_channels=ctrl_out_channels,
@@ -1732,7 +1732,7 @@ class ControlNetXSCrossAttnUpBlock2D(nn.Module):
 
             if has_crossattn:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,

src/diffusers/models/transformers/__init__.py

@@ -12,5 +12,6 @@ if is_torch_available():
     from .prior_transformer import PriorTransformer
     from .t5_film_transformer import T5FilmDecoder
     from .transformer_2d import Transformer2DModel
+    from .transformer_2d_block import ContinuousTransformer2DModelBlock
     from .transformer_sd3 import SD3Transformer2DModel
     from .transformer_temporal import TransformerTemporalModel

src/diffusers/models/transformers/transformer_2d.py

@@ -115,6 +115,10 @@ class Transformer2DModel(LegacyModelMixin, LegacyConfigMixin):
         self.is_input_vectorized = num_vector_embeds is not None
         self.is_input_patches = in_channels is not None and patch_size is not None
 
+        if self.is_input_continuous:
+            deprecation_message = "Using `Transformer2DModel` when the input is continuous is deprecared and it will be removed in a future version. Please use `ContinuousTransformer2DModelBlock`, importing from `diffusers.models.transformers.transformer_2d_block`."
+            deprecate("Continuous transformer block.", "1.0.0", deprecation_message)
+
         if self.is_input_continuous and self.is_input_vectorized:
             raise ValueError(
                 f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"

src/diffusers/models/transformers/transformer_2d_block.py

@@ -0,0 +1,282 @@
+# 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.
+from typing import Any, Dict, Optional
+
+import torch
+from torch import nn
+
+from ...utils import is_torch_version, logging
+from ..attention import BasicTransformerBlock
+from ..modeling_outputs import Transformer2DModelOutput
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class ContinuousTransformer2DModelBlock(nn.Module):
+    """
+    A 2D Transformer block for continuous image-like data.
+
+    Parameters:
+        num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
+        in_channels (`int`, *optional*):
+            The number of channels in the input and output (specify if the input is **continuous**).
+        num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
+            This is fixed during training since it is used to learn a number of position embeddings.
+        num_vector_embeds (`int`, *optional*):
+            The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
+            Includes the class for the masked latent pixel.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
+        num_embeds_ada_norm ( `int`, *optional*):
+            The number of diffusion steps used during training. Pass if at least one of the norm_layers is
+            `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
+            added to the hidden states.
+
+            During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
+        attention_bias (`bool`, *optional*):
+            Configure if the `TransformerBlocks` attention should contain a bias parameter.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["BasicTransformerBlock"]
+
+    def __init__(
+        self,
+        num_attention_heads: int = 16,
+        attention_head_dim: int = 88,
+        in_channels: Optional[int] = None,
+        out_channels: Optional[int] = None,
+        num_layers: int = 1,
+        dropout: float = 0.0,
+        norm_num_groups: int = 32,
+        cross_attention_dim: Optional[int] = None,
+        attention_bias: bool = False,
+        sample_size: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        use_linear_projection: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_type: str = "layer_norm",  # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        attention_type: str = "default",
+        use_additional_conditions: Optional[bool] = None,
+    ):
+        super().__init__()
+
+        # Set some common variables used across the board.
+        self.use_linear_projection = use_linear_projection
+        out_channels = in_channels if out_channels is None else out_channels
+        self.gradient_checkpointing = False
+
+        if use_additional_conditions is None:
+            if norm_type == "ada_norm_single" and sample_size == 128:
+                use_additional_conditions = True
+            else:
+                use_additional_conditions = False
+        self.use_additional_conditions = use_additional_conditions
+
+        # Norm
+        inner_dim = num_attention_heads * attention_head_dim
+        self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+        # Input projection.
+        if self.use_linear_projection:
+            self.proj_in = torch.nn.Linear(in_channels, inner_dim)
+        else:
+            self.proj_in = torch.nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+
+        # Transformer blocks.
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    num_attention_heads,
+                    attention_head_dim,
+                    dropout=dropout,
+                    cross_attention_dim=cross_attention_dim,
+                    activation_fn=activation_fn,
+                    num_embeds_ada_norm=num_embeds_ada_norm,
+                    attention_bias=attention_bias,
+                    only_cross_attention=only_cross_attention,
+                    double_self_attention=double_self_attention,
+                    upcast_attention=upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=norm_elementwise_affine,
+                    norm_eps=norm_eps,
+                    attention_type=attention_type,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # Output projection.
+        out_channels = in_channels if out_channels is None else out_channels
+        if self.use_linear_projection:
+            self.proj_out = torch.nn.Linear(inner_dim, out_channels)
+        else:
+            self.proj_out = torch.nn.Conv2d(inner_dim, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ):
+        """
+        The [`Transformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.Tensor` of shape `(batch size, channel, height, width)` if continuous):
+                Input `hidden_states`.
+            encoder_hidden_states ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
+                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
+                self-attention.
+            timestep ( `torch.LongTensor`, *optional*):
+                Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
+            class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
+                Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
+                `AdaLayerZeroNorm`.
+            cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            attention_mask ( `torch.Tensor`, *optional*):
+                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                negative values to the attention scores corresponding to "discard" tokens.
+            encoder_attention_mask ( `torch.Tensor`, *optional*):
+                Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
+
+                    * Mask `(batch, sequence_length)` True = keep, False = discard.
+                    * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
+
+                If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
+                above. This bias will be added to the cross-attention scores.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformers.transformer_2d.Transformer2DModelOutput`] is returned,
+            otherwise a `tuple` where the first element is the sample tensor.
+        """
+        if cross_attention_kwargs is not None:
+            if cross_attention_kwargs.get("scale", None) is not None:
+                logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
+        #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
+        #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        if attention_mask is not None and attention_mask.ndim == 2:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 1. Input
+        batch_size, _, height, width = hidden_states.shape
+        residual = hidden_states
+        hidden_states = self.norm(hidden_states)
+
+        if not self.use_linear_projection:
+            hidden_states = self.proj_in(hidden_states)
+            inner_dim = hidden_states.shape[1]
+            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, inner_dim)
+        else:
+            inner_dim = hidden_states.shape[1]
+            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, inner_dim)
+            hidden_states = self.proj_in(hidden_states)
+
+        # 2. Blocks
+        for block in self.transformer_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    attention_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    timestep,
+                    cross_attention_kwargs,
+                    class_labels,
+                    **ckpt_kwargs,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    timestep=timestep,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    class_labels=class_labels,
+                )
+
+        # 3. Output
+        if not self.use_linear_projection:
+            hidden_states = (
+                hidden_states.reshape(batch_size, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
+            )
+            hidden_states = self.proj_out(hidden_states)
+        else:
+            hidden_states = self.proj_out(hidden_states)
+            hidden_states = (
+                hidden_states.reshape(batch_size, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
+            )
+        output = hidden_states + residual
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

src/diffusers/models/unets/unet_2d_blocks.py

@@ -34,7 +34,7 @@ from ..resnet import (
     Upsample2D,
 )
 from ..transformers.dual_transformer_2d import DualTransformer2DModel
-from ..transformers.transformer_2d import Transformer2DModel
+from ..transformers.transformer_2d_block import ContinuousTransformer2DModelBlock
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -801,7 +801,7 @@ class UNetMidBlock2DCrossAttn(nn.Module):
         for i in range(num_layers):
             if not dual_cross_attention:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,
@@ -1198,7 +1198,7 @@ class CrossAttnDownBlock2D(nn.Module):
             )
             if not dual_cross_attention:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,
@@ -2444,7 +2444,7 @@ class CrossAttnUpBlock2D(nn.Module):
             )
             if not dual_cross_attention:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,

src/diffusers/models/unets/unet_3d_blocks.py

@@ -28,7 +28,7 @@ from ..resnet import (
     Upsample2D,
 )
 from ..transformers.dual_transformer_2d import DualTransformer2DModel
-from ..transformers.transformer_2d import Transformer2DModel
+from ..transformers.transformer_2d_block import ContinuousTransformer2DModelBlock
 from ..transformers.transformer_temporal import (
     TransformerSpatioTemporalModel,
     TransformerTemporalModel,
@@ -375,7 +375,7 @@ class UNetMidBlock3DCrossAttn(nn.Module):
 
         for _ in range(num_layers):
             attentions.append(
-                Transformer2DModel(
+                ContinuousTransformer2DModelBlock(
                     in_channels // num_attention_heads,
                     num_attention_heads,
                     in_channels=in_channels,
@@ -513,7 +513,7 @@ class CrossAttnDownBlock3D(nn.Module):
                 )
             )
             attentions.append(
-                Transformer2DModel(
+                ContinuousTransformer2DModelBlock(
                     out_channels // num_attention_heads,
                     num_attention_heads,
                     in_channels=out_channels,
@@ -747,7 +747,7 @@ class CrossAttnUpBlock3D(nn.Module):
                 )
             )
             attentions.append(
-                Transformer2DModel(
+                ContinuousTransformer2DModelBlock(
                     out_channels // num_attention_heads,
                     num_attention_heads,
                     in_channels=out_channels,
@@ -1162,7 +1162,7 @@ class CrossAttnDownBlockMotion(nn.Module):
 
             if not dual_cross_attention:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,
@@ -1373,7 +1373,7 @@ class CrossAttnUpBlockMotion(nn.Module):
 
             if not dual_cross_attention:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,
@@ -1736,7 +1736,7 @@ class UNetMidBlockCrossAttnMotion(nn.Module):
         for i in range(num_layers):
             if not dual_cross_attention:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         in_channels // num_attention_heads,
                         in_channels=in_channels,

src/diffusers/pipelines/audioldm2/modeling_audioldm2.py

@@ -35,7 +35,7 @@ from ...models.embeddings import (
 )
 from ...models.modeling_utils import ModelMixin
 from ...models.resnet import Downsample2D, ResnetBlock2D, Upsample2D
-from ...models.transformers.transformer_2d import Transformer2DModel
+from ...models.transformers import ContinuousTransformer2DModelBlock
 from ...models.unets.unet_2d_blocks import DownBlock2D, UpBlock2D
 from ...models.unets.unet_2d_condition import UNet2DConditionOutput
 from ...utils import BaseOutput, is_torch_version, logging
@@ -1060,7 +1060,7 @@ class CrossAttnDownBlock2D(nn.Module):
             )
             for j in range(len(cross_attention_dim)):
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,
@@ -1236,7 +1236,7 @@ class UNetMidBlock2DCrossAttn(nn.Module):
         for i in range(num_layers):
             for j in range(len(cross_attention_dim)):
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         in_channels // num_attention_heads,
                         in_channels=in_channels,
@@ -1412,7 +1412,7 @@ class CrossAttnUpBlock2D(nn.Module):
             )
             for j in range(len(cross_attention_dim)):
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,

src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py

@@ -31,8 +31,8 @@ from ....models.embeddings import (
     Timesteps,
 )
 from ....models.resnet import ResnetBlockCondNorm2D
+from ....models.transformers import ContinuousTransformer2DModelBlock
 from ....models.transformers.dual_transformer_2d import DualTransformer2DModel
-from ....models.transformers.transformer_2d import Transformer2DModel
 from ....models.unets.unet_2d_condition import UNet2DConditionOutput
 from ....utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
 from ....utils.torch_utils import apply_freeu
@@ -1677,7 +1677,7 @@ class CrossAttnDownBlockFlat(nn.Module):
             )
             if not dual_cross_attention:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,
@@ -1957,7 +1957,7 @@ class CrossAttnUpBlockFlat(nn.Module):
             )
             if not dual_cross_attention:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,
@@ -2294,7 +2294,7 @@ class UNetMidBlockFlatCrossAttn(nn.Module):
         for i in range(num_layers):
             if not dual_cross_attention:
                 attentions.append(
-                    Transformer2DModel(
+                    ContinuousTransformer2DModelBlock(
                         num_attention_heads,
                         out_channels // num_attention_heads,
                         in_channels=out_channels,

tests/models/unets/test_models_unet_2d_condition.py

@@ -628,7 +628,7 @@ class UNet2DConditionModelTests(ModelTesterMixin, UNetTesterMixin, unittest.Test
             "CrossAttnDownBlock2D",
             "UNetMidBlock2DCrossAttn",
             "UpBlock2D",
-            "Transformer2DModel",
+            "ContinuousTransformer2DModelBlock",
             "DownBlock2D",
         }
 

tests/models/unets/test_models_unet_controlnetxs.py

@@ -291,7 +291,7 @@ class UNetControlNetXSModelTests(ModelTesterMixin, UNetTesterMixin, unittest.Tes
         model.enable_gradient_checkpointing()
 
         EXPECTED_SET = {
-            "Transformer2DModel",
+            "ContinuousTransformer2DModelBlock",
             "UNetMidBlock2DCrossAttn",
             "ControlNetXSCrossAttnDownBlock2D",
             "ControlNetXSCrossAttnMidBlock2D",

tests/models/unets/test_models_unet_motion.py

@@ -186,7 +186,7 @@ class UNetMotionModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase)
             "CrossAttnDownBlockMotion",
             "UNetMidBlockCrossAttnMotion",
             "UpBlockMotion",
-            "Transformer2DModel",
+            "ContinuousTransformer2DModelBlock",
             "DownBlockMotion",
         }