Function noise_like

diffusion/diffusion_ddpm_pan.py:79–88

(shape, device, repeat=False)

def noise_like(shape, device, repeat=False):
    def repeat_noise():
        return torch.randn((1, *shape[1:]), device=device).repeat(
            shape[0], *((1,) * (len(shape) - 1))
        )

    def noise():
        return torch.randn(shape, device=device)

    return repeat_noise() if repeat else noise()


def normal_kl(mean1, logvar1, mean2, logvar2):