class MochiChunkedGroupNorm3D(nn.Module):
    r"""
    Applies per-frame group normalization for 5D video inputs. It also supports memory-efficient chunked group
    normalization.

    Args:
        num_channels (int): Number of channels expected in input
        num_groups (int, optional): Number of groups to separate the channels into. Default: 32
        affine (bool, optional): If True, this module has learnable affine parameters. Default: True
        chunk_size (int, optional): Size of each chunk for processing. Default: 8

    """

    def __init__(
        self,
        num_channels: int,
        num_groups: int = 32,
        affine: bool = True,
        chunk_size: int = 8,
    ):
        super().__init__()
        self.norm_layer = nn.GroupNorm(num_channels=num_channels, num_groups=num_groups, affine=affine)
        self.chunk_size = chunk_size

    def forward(self, x: torch.Tensor = None) -> torch.Tensor:
        batch_size = x.size(0)

        x = x.permute(0, 2, 1, 3, 4).flatten(0, 1)
        output = torch.cat([self.norm_layer(chunk) for chunk in x.split(self.chunk_size, dim=0)], dim=0)
        output = output.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)

        return output
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py