class MochiMidBlock3D(nn.Module):
    r"""
    A middle block used in the Mochi model.

    Args:
        in_channels (`int`):
            Number of input channels.
        num_layers (`int`, defaults to `3`):
            Number of resnet blocks in the block.
    """

    def __init__(
        self,
        in_channels: int,  # 768
        num_layers: int = 3,
        add_attention: bool = True,
    ):
        super().__init__()

        resnets = []
        norms = []
        attentions = []

        for _ in range(num_layers):
            resnets.append(MochiResnetBlock3D(in_channels=in_channels))

            if add_attention:
                norms.append(MochiChunkedGroupNorm3D(num_channels=in_channels))
                attentions.append(
                    Attention(
                        query_dim=in_channels,
                        heads=in_channels // 32,
                        dim_head=32,
                        qk_norm="l2",
                        is_causal=True,
                        processor=MochiVaeAttnProcessor2_0(),
                    )
                )
            else:
                norms.append(None)
                attentions.append(None)

        self.resnets = nn.ModuleList(resnets)
        self.norms = nn.ModuleList(norms)
        self.attentions = nn.ModuleList(attentions)

        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        conv_cache: Optional[Dict[str, torch.Tensor]] = None,
    ) -> torch.Tensor:
        r"""Forward method of the `MochiMidBlock3D` class."""

        new_conv_cache = {}
        conv_cache = conv_cache or {}

        for i, (resnet, norm, attn) in enumerate(zip(self.resnets, self.norms, self.attentions)):
            conv_cache_key = f"resnet_{i}"

            if torch.is_grad_enabled() and self.gradient_checkpointing:

                def create_custom_forward(module):
                    def create_forward(*inputs):
                        return module(*inputs)

                    return create_forward

                hidden_states, new_conv_cache[conv_cache_key] = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(resnet), hidden_states, conv_cache=conv_cache.get(conv_cache_key)
                )
            else:
                hidden_states, new_conv_cache[conv_cache_key] = resnet(
                    hidden_states, conv_cache=conv_cache.get(conv_cache_key)
                )

            if attn is not None:
                residual = hidden_states
                hidden_states = norm(hidden_states)

                batch_size, num_channels, num_frames, height, width = hidden_states.shape
                hidden_states = hidden_states.permute(0, 3, 4, 2, 1).flatten(0, 2).contiguous()
                hidden_states = attn(hidden_states)
                hidden_states = hidden_states.unflatten(0, (batch_size, height, width)).permute(0, 4, 3, 1, 2)

                hidden_states = residual + hidden_states

        return hidden_states, new_conv_cache
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py