class MochiUpBlock3D(nn.Module):
    r"""
    An upsampling block used in the Mochi model.

    Args:
        in_channels (`int`):
            Number of input channels.
        out_channels (`int`, *optional*):
            Number of output channels. If None, defaults to `in_channels`.
        num_layers (`int`, defaults to `1`):
            Number of resnet blocks in the block.
        temporal_expansion (`int`, defaults to `2`):
            Temporal expansion factor.
        spatial_expansion (`int`, defaults to `2`):
            Spatial expansion factor.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_layers: int = 1,
        temporal_expansion: int = 2,
        spatial_expansion: int = 2,
    ):
        super().__init__()
        self.temporal_expansion = temporal_expansion
        self.spatial_expansion = spatial_expansion

        resnets = []
        for _ in range(num_layers):
            resnets.append(MochiResnetBlock3D(in_channels=in_channels))
        self.resnets = nn.ModuleList(resnets)

        self.proj = nn.Linear(in_channels, out_channels * temporal_expansion * spatial_expansion**2)

        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 `MochiUpBlock3D` class."""

        new_conv_cache = {}
        conv_cache = conv_cache or {}

        for i, resnet in enumerate(self.resnets):
            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)
                )

        hidden_states = hidden_states.permute(0, 2, 3, 4, 1)
        hidden_states = self.proj(hidden_states)
        hidden_states = hidden_states.permute(0, 4, 1, 2, 3)

        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        st = self.temporal_expansion
        sh = self.spatial_expansion
        sw = self.spatial_expansion

        # Reshape and unpatchify
        hidden_states = hidden_states.view(batch_size, -1, st, sh, sw, num_frames, height, width)
        hidden_states = hidden_states.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
        hidden_states = hidden_states.view(batch_size, -1, num_frames * st, height * sh, width * sw)

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