class MotionAdapter(ModelMixin, ConfigMixin, FromOriginalModelMixin):
    @register_to_config
    def __init__(
        self,
        block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
        motion_layers_per_block: Union[int, Tuple[int]] = 2,
        motion_transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple[int]]] = 1,
        motion_mid_block_layers_per_block: int = 1,
        motion_transformer_layers_per_mid_block: Union[int, Tuple[int]] = 1,
        motion_num_attention_heads: Union[int, Tuple[int]] = 8,
        motion_norm_num_groups: int = 32,
        motion_max_seq_length: int = 32,
        use_motion_mid_block: bool = True,
        conv_in_channels: Optional[int] = None,
    ):
        """Container to store AnimateDiff Motion Modules

        Args:
            block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
            The tuple of output channels for each UNet block.
            motion_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 2):
                The number of motion layers per UNet block.
            motion_transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple[int]]`, *optional*, defaults to 1):
                The number of transformer layers to use in each motion layer in each block.
            motion_mid_block_layers_per_block (`int`, *optional*, defaults to 1):
                The number of motion layers in the middle UNet block.
            motion_transformer_layers_per_mid_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
                The number of transformer layers to use in each motion layer in the middle block.
            motion_num_attention_heads (`int` or `Tuple[int]`, *optional*, defaults to 8):
                The number of heads to use in each attention layer of the motion module.
            motion_norm_num_groups (`int`, *optional*, defaults to 32):
                The number of groups to use in each group normalization layer of the motion module.
            motion_max_seq_length (`int`, *optional*, defaults to 32):
                The maximum sequence length to use in the motion module.
            use_motion_mid_block (`bool`, *optional*, defaults to True):
                Whether to use a motion module in the middle of the UNet.
        """

        super().__init__()
        down_blocks = []
        up_blocks = []

        if isinstance(motion_layers_per_block, int):
            motion_layers_per_block = (motion_layers_per_block,) * len(block_out_channels)
        elif len(motion_layers_per_block) != len(block_out_channels):
            raise ValueError(
                f"The number of motion layers per block must match the number of blocks, "
                f"got {len(block_out_channels)} and {len(motion_layers_per_block)}"
            )

        if isinstance(motion_transformer_layers_per_block, int):
            motion_transformer_layers_per_block = (motion_transformer_layers_per_block,) * len(block_out_channels)

        if isinstance(motion_transformer_layers_per_mid_block, int):
            motion_transformer_layers_per_mid_block = (
                motion_transformer_layers_per_mid_block,
            ) * motion_mid_block_layers_per_block
        elif len(motion_transformer_layers_per_mid_block) != motion_mid_block_layers_per_block:
            raise ValueError(
                f"The number of layers per mid block ({motion_mid_block_layers_per_block}) "
                f"must match the length of motion_transformer_layers_per_mid_block ({len(motion_transformer_layers_per_mid_block)})"
            )

        if isinstance(motion_num_attention_heads, int):
            motion_num_attention_heads = (motion_num_attention_heads,) * len(block_out_channels)
        elif len(motion_num_attention_heads) != len(block_out_channels):
            raise ValueError(
                f"The length of the attention head number tuple in the motion module must match the "
                f"number of block, got {len(motion_num_attention_heads)} and {len(block_out_channels)}"
            )

        if conv_in_channels:
            # input
            self.conv_in = nn.Conv2d(conv_in_channels, block_out_channels[0], kernel_size=3, padding=1)
        else:
            self.conv_in = None

        for i, channel in enumerate(block_out_channels):
            output_channel = block_out_channels[i]
            down_blocks.append(
                MotionModules(
                    in_channels=output_channel,
                    norm_num_groups=motion_norm_num_groups,
                    cross_attention_dim=None,
                    activation_fn="geglu",
                    attention_bias=False,
                    num_attention_heads=motion_num_attention_heads[i],
                    max_seq_length=motion_max_seq_length,
                    layers_per_block=motion_layers_per_block[i],
                    transformer_layers_per_block=motion_transformer_layers_per_block[i],
                )
            )

        if use_motion_mid_block:
            self.mid_block = MotionModules(
                in_channels=block_out_channels[-1],
                norm_num_groups=motion_norm_num_groups,
                cross_attention_dim=None,
                activation_fn="geglu",
                attention_bias=False,
                num_attention_heads=motion_num_attention_heads[-1],
                max_seq_length=motion_max_seq_length,
                layers_per_block=motion_mid_block_layers_per_block,
                transformer_layers_per_block=motion_transformer_layers_per_mid_block,
            )
        else:
            self.mid_block = None

        reversed_block_out_channels = list(reversed(block_out_channels))
        output_channel = reversed_block_out_channels[0]

        reversed_motion_layers_per_block = list(reversed(motion_layers_per_block))
        reversed_motion_transformer_layers_per_block = list(reversed(motion_transformer_layers_per_block))
        reversed_motion_num_attention_heads = list(reversed(motion_num_attention_heads))
        for i, channel in enumerate(reversed_block_out_channels):
            output_channel = reversed_block_out_channels[i]
            up_blocks.append(
                MotionModules(
                    in_channels=output_channel,
                    norm_num_groups=motion_norm_num_groups,
                    cross_attention_dim=None,
                    activation_fn="geglu",
                    attention_bias=False,
                    num_attention_heads=reversed_motion_num_attention_heads[i],
                    max_seq_length=motion_max_seq_length,
                    layers_per_block=reversed_motion_layers_per_block[i] + 1,
                    transformer_layers_per_block=reversed_motion_transformer_layers_per_block[i],
                )
            )

        self.down_blocks = nn.ModuleList(down_blocks)
        self.up_blocks = nn.ModuleList(up_blocks)

    def forward(self, sample):
        pass
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/unets/unet_motion_model.py