class SDCascadeLayerNorm(nn.LayerNorm):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, x):
        x = x.permute(0, 2, 3, 1)
        x = super().forward(x)
        return x.permute(0, 3, 1, 2)
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/unets/unet_stable_cascade.py