class UTransformer2DModel(ModelMixin, ConfigMixin):
    """
    Transformer model based on the [U-ViT](https://github.com/baofff/U-ViT) architecture for image-like data. Compared
    to [`Transformer2DModel`], this model has skip connections between transformer blocks in a "U"-shaped fashion,
    similar to a U-Net. Supports only continuous (actual embeddings) inputs, which are embedded via a [`PatchEmbed`]
    layer and then reshaped to (b, t, d).

    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*):
            Pass if the input is continuous. The number of channels in the input.
        out_channels (`int`, *optional*):
            The number of output channels; if `None`, defaults to `in_channels`.
        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.
        norm_num_groups (`int`, *optional*, defaults to `32`):
            The number of groups to use when performing Group Normalization.
        cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use.
        attention_bias (`bool`, *optional*):
            Configure if the TransformerBlocks' attention should contain a bias parameter.
        sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images.
            Note that this is fixed at training time as it is used for learning a number of position embeddings. See
            `ImagePositionalEmbeddings`.
        num_vector_embeds (`int`, *optional*):
            Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels.
            Includes the class for the masked latent pixel.
        patch_size (`int`, *optional*, defaults to 2):
            The patch size to use in the patch embedding.
        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
        num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`.
            The number of diffusion steps used during training. Note that this is fixed at training time as 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 than steps than `num_embeds_ada_norm`.
        use_linear_projection (int, *optional*): TODO: Not used
        only_cross_attention (`bool`, *optional*):
            Whether to use only cross-attention layers. In this case two cross attention layers are used in each
            transformer block.
        upcast_attention (`bool`, *optional*):
            Whether to upcast the query and key to float() when performing the attention calculation.
        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
            The Layer Normalization implementation to use. Defaults to `torch.nn.LayerNorm`.
        block_type (`str`, *optional*, defaults to `"unidiffuser"`):
            The transformer block implementation to use. If `"unidiffuser"`, has the LayerNorms on the residual
            backbone of each transformer block; otherwise has them in the attention/feedforward branches (the standard
            behavior in `diffusers`.)
        pre_layer_norm (`bool`, *optional*):
            Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"),
            as opposed to after ("post-LayerNorm"). The original UniDiffuser implementation is post-LayerNorm
            (`pre_layer_norm = False`).
        norm_elementwise_affine (`bool`, *optional*):
            Whether to use learnable per-element affine parameters during layer normalization.
        use_patch_pos_embed (`bool`, *optional*):
            Whether to use position embeddings inside the patch embedding layer (`PatchEmbed`).
        final_dropout (`bool`, *optional*):
            Whether to use a final Dropout layer after the feedforward network.
    """

    @register_to_config
    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,
        num_vector_embeds: Optional[int] = None,
        patch_size: Optional[int] = 2,
        activation_fn: str = "geglu",
        num_embeds_ada_norm: Optional[int] = None,
        use_linear_projection: bool = False,
        only_cross_attention: bool = False,
        upcast_attention: bool = False,
        norm_type: str = "layer_norm",
        block_type: str = "unidiffuser",
        pre_layer_norm: bool = False,
        norm_elementwise_affine: bool = True,
        use_patch_pos_embed=False,
        ff_final_dropout: bool = False,
    ):
        super().__init__()
        self.use_linear_projection = use_linear_projection
        self.num_attention_heads = num_attention_heads
        self.attention_head_dim = attention_head_dim
        inner_dim = num_attention_heads * attention_head_dim

        # 1. Input
        # Only support patch input of shape (batch_size, num_channels, height, width) for now
        assert in_channels is not None and patch_size is not None, "Patch input requires in_channels and patch_size."

        assert sample_size is not None, "UTransformer2DModel over patched input must provide sample_size"

        # 2. Define input layers
        self.height = sample_size
        self.width = sample_size

        self.patch_size = patch_size
        self.pos_embed = PatchEmbed(
            height=sample_size,
            width=sample_size,
            patch_size=patch_size,
            in_channels=in_channels,
            embed_dim=inner_dim,
            use_pos_embed=use_patch_pos_embed,
        )

        # 3. Define transformers blocks
        # Modify this to have in_blocks ("downsample" blocks, even though we don't actually downsample), a mid_block,
        # and out_blocks ("upsample" blocks). Like a U-Net, there are skip connections from in_blocks to out_blocks in
        # a "U"-shaped fashion (e.g. first in_block to last out_block, etc.).
        # Quick hack to make the transformer block type configurable
        if block_type == "unidiffuser":
            block_cls = UniDiffuserBlock
        else:
            block_cls = UTransformerBlock
        self.transformer_in_blocks = nn.ModuleList(
            [
                block_cls(
                    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,
                    upcast_attention=upcast_attention,
                    norm_type=norm_type,
                    pre_layer_norm=pre_layer_norm,
                    norm_elementwise_affine=norm_elementwise_affine,
                    final_dropout=ff_final_dropout,
                )
                for d in range(num_layers // 2)
            ]
        )

        self.transformer_mid_block = block_cls(
            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,
            upcast_attention=upcast_attention,
            norm_type=norm_type,
            pre_layer_norm=pre_layer_norm,
            norm_elementwise_affine=norm_elementwise_affine,
            final_dropout=ff_final_dropout,
        )

        # For each skip connection, we use a SkipBlock (concatenation + Linear + LayerNorm) to process the inputs
        # before each transformer out_block.
        self.transformer_out_blocks = nn.ModuleList(
            [
                nn.ModuleDict(
                    {
                        "skip": SkipBlock(
                            inner_dim,
                        ),
                        "block": block_cls(
                            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,
                            upcast_attention=upcast_attention,
                            norm_type=norm_type,
                            pre_layer_norm=pre_layer_norm,
                            norm_elementwise_affine=norm_elementwise_affine,
                            final_dropout=ff_final_dropout,
                        ),
                    }
                )
                for d in range(num_layers // 2)
            ]
        )

        # 4. Define output layers
        self.out_channels = in_channels if out_channels is None else out_channels

        # Following the UniDiffuser U-ViT implementation, we process the transformer output with
        # a LayerNorm layer with per-element affine params
        self.norm_out = nn.LayerNorm(inner_dim)

    def forward(
        self,
        hidden_states,
        encoder_hidden_states=None,
        timestep=None,
        class_labels=None,
        cross_attention_kwargs=None,
        return_dict: bool = True,
        hidden_states_is_embedding: bool = False,
        unpatchify: bool = True,
    ):
        """
        Args:
            hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`.
                When continuous, `torch.Tensor` of shape `(batch size, channel, height, width)`): Input hidden_states
            encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*):
                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
                self-attention.
            timestep ( `torch.long`, *optional*):
                Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step.
            class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
                Optional class labels to be applied as an embedding in AdaLayerZeroNorm. Used to indicate class labels
                conditioning.
            cross_attention_kwargs (*optional*):
                Keyword arguments to supply to the cross attention layers, if used.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
                tuple.
            hidden_states_is_embedding (`bool`, *optional*, defaults to `False`):
                Whether or not hidden_states is an embedding directly usable by the transformer. In this case we will
                ignore input handling (e.g. continuous, vectorized, etc.) and directly feed hidden_states into the
                transformer blocks.
            unpatchify (`bool`, *optional*, defaults to `True`):
                Whether to unpatchify the transformer output.

        Returns:
            [`~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.
        """
        # 0. Check inputs

        if not unpatchify and return_dict:
            raise ValueError(
                f"Cannot both define `unpatchify`: {unpatchify} and `return_dict`: {return_dict} since when"
                f" `unpatchify` is {unpatchify} the returned output is of shape (batch_size, seq_len, hidden_dim)"
                " rather than (batch_size, num_channels, height, width)."
            )

        # 1. Input
        if not hidden_states_is_embedding:
            hidden_states = self.pos_embed(hidden_states)

        # 2. Blocks

        # In ("downsample") blocks
        skips = []
        for in_block in self.transformer_in_blocks:
            hidden_states = in_block(
                hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                timestep=timestep,
                cross_attention_kwargs=cross_attention_kwargs,
                class_labels=class_labels,
            )
            skips.append(hidden_states)

        # Mid block
        hidden_states = self.transformer_mid_block(hidden_states)

        # Out ("upsample") blocks
        for out_block in self.transformer_out_blocks:
            hidden_states = out_block["skip"](hidden_states, skips.pop())
            hidden_states = out_block["block"](
                hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                timestep=timestep,
                cross_attention_kwargs=cross_attention_kwargs,
                class_labels=class_labels,
            )

        # 3. Output
        # Don't support AdaLayerNorm for now, so no conditioning/scale/shift logic
        hidden_states = self.norm_out(hidden_states)
        # hidden_states = self.proj_out(hidden_states)

        if unpatchify:
            # unpatchify
            height = width = int(hidden_states.shape[1] ** 0.5)
            hidden_states = hidden_states.reshape(
                shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)
            )
            hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
            output = hidden_states.reshape(
                shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)
            )
        else:
            output = hidden_states

        if not return_dict:
            return (output,)

        return Transformer2DModelOutput(sample=output)
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/pipelines/unidiffuser/modeling_uvit.py