Naive group norm implementation that somehow beats the native kernel? For channels last. But for some reason also faster on channels first.

channels_last_group_norm.py

def cl_weight_hook(state_dict, *args, **kwargs):
    for key in state_dict.keys():
        state_dict[key] = state_dict[key].reshape(1, -1, 1, 1).to(memory_format=torch.channels_last)

class CLGroupNorm(torch.nn.GroupNorm):
    def __init__(self, num_groups: int, num_channels: int, eps: float = 0.00001, affine: bool = True, device=None, dtype=None) -> None:
        super().__init__(num_groups, num_channels, eps, affine, device, dtype)
        if self.weight.ndim == 1:
            self.weight.data = self.weight.data.reshape(1, -1, 1, 1).to(memory_format=torch.channels_last)
        if self.bias.ndim == 1:
            self.bias.data = self.bias.data.reshape(1, -1, 1, 1).to(memory_format=torch.channels_last)
        self._register_load_state_dict_pre_hook(cl_weight_hook)

    def forward(self, x: torch.Tensor):
        N,C,H,W = x.size()
        G = self.num_groups
        assert C % G == 0

        x = x.view(N, G, C//G, -1)
        mean = x.mean((2, 3), keepdim=True)
        var = x.var((2, 3), keepdim=True)

        x = (x-mean) / (var+self.eps).sqrt()
        x = x.view(N,C,H,W)
        return x * self.weight + self.bias

def NormalizeCL(in_channels):
    return CLGroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)

def force_model_fp16():
    """
    ldm and sgm has modules.diffusionmodules.util.GroupNorm32.forward, which
    force conversion of input to float32. If force_fp16 is enabled, we need to
    prevent this casting.
    """
    assert force_fp16
    import sgm.modules.diffusionmodules.util as sgm_util
    import ldm.modules.diffusionmodules.util as ldm_util
    import ldm.modules.attention as ldm_attn
    sgm_util.GroupNorm32 = CLGroupNorm
    ldm_util.GroupNorm32 = CLGroupNorm
    ldm_attn.Normalize = NormalizeCL

    print("ldm/sgm GroupNorm32 replaced with naive groupnorm impl due to `--precision half`.")