class MochiDecoder3D(nn.Module):
    r"""
    The `MochiDecoder3D` layer of a variational autoencoder that decodes its latent representation into an output
    sample.

    Args:
        in_channels (`int`, *optional*):
            The number of input channels.
        out_channels (`int`, *optional*):
            The number of output channels.
        block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(128, 256, 512, 768)`):
            The number of output channels for each block.
        layers_per_block (`Tuple[int, ...]`, *optional*, defaults to `(3, 3, 4, 6, 3)`):
            The number of resnet blocks for each block.
        temporal_expansions (`Tuple[int, ...]`, *optional*, defaults to `(1, 2, 3)`):
            The temporal expansion factor for each of the up blocks.
        spatial_expansions (`Tuple[int, ...]`, *optional*, defaults to `(2, 2, 2)`):
            The spatial expansion factor for each of the up blocks.
        non_linearity (`str`, *optional*, defaults to `"swish"`):
            The non-linearity to use in the decoder.
    """

    def __init__(
        self,
        in_channels: int,  # 12
        out_channels: int,  # 3
        block_out_channels: Tuple[int, ...] = (128, 256, 512, 768),
        layers_per_block: Tuple[int, ...] = (3, 3, 4, 6, 3),
        temporal_expansions: Tuple[int, ...] = (1, 2, 3),
        spatial_expansions: Tuple[int, ...] = (2, 2, 2),
        act_fn: str = "swish",
    ):
        super().__init__()

        self.nonlinearity = get_activation(act_fn)

        self.conv_in = nn.Conv3d(in_channels, block_out_channels[-1], kernel_size=(1, 1, 1))
        self.block_in = MochiMidBlock3D(
            in_channels=block_out_channels[-1],
            num_layers=layers_per_block[-1],
            add_attention=False,
        )

        up_blocks = []
        for i in range(len(block_out_channels) - 1):
            up_block = MochiUpBlock3D(
                in_channels=block_out_channels[-i - 1],
                out_channels=block_out_channels[-i - 2],
                num_layers=layers_per_block[-i - 2],
                temporal_expansion=temporal_expansions[-i - 1],
                spatial_expansion=spatial_expansions[-i - 1],
            )
            up_blocks.append(up_block)
        self.up_blocks = nn.ModuleList(up_blocks)

        self.block_out = MochiMidBlock3D(
            in_channels=block_out_channels[0],
            num_layers=layers_per_block[0],
            add_attention=False,
        )
        self.proj_out = nn.Linear(block_out_channels[0], out_channels)

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

        new_conv_cache = {}
        conv_cache = conv_cache or {}

        hidden_states = self.conv_in(hidden_states)

        # 1. Mid
        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["block_in"] = torch.utils.checkpoint.checkpoint(
                create_custom_forward(self.block_in), hidden_states, conv_cache=conv_cache.get("block_in")
            )

            for i, up_block in enumerate(self.up_blocks):
                conv_cache_key = f"up_block_{i}"
                hidden_states, new_conv_cache[conv_cache_key] = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(up_block), hidden_states, conv_cache=conv_cache.get(conv_cache_key)
                )
        else:
            hidden_states, new_conv_cache["block_in"] = self.block_in(
                hidden_states, conv_cache=conv_cache.get("block_in")
            )

            for i, up_block in enumerate(self.up_blocks):
                conv_cache_key = f"up_block_{i}"
                hidden_states, new_conv_cache[conv_cache_key] = up_block(
                    hidden_states, conv_cache=conv_cache.get(conv_cache_key)
                )

        hidden_states, new_conv_cache["block_out"] = self.block_out(
            hidden_states, conv_cache=conv_cache.get("block_out")
        )

        hidden_states = self.nonlinearity(hidden_states)

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

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