Function pyramid_noise_like

k_diffusion/utils.py:462–470

(noise, iterations=6, discount=0.3)

def pyramid_noise_like(noise, iterations=6, discount=0.3):
    b, c, w, h = noise.shape
    u = torch.nn.Upsample(size=(w, h), mode='bilinear')
    for i in range(iterations):
        r = random.random()*2+2 # Rather than always going 2x, 
        w, h = max(1, int(w/(r**i))), max(1, int(h/(r**i)))
        noise += u(torch.randn((b, c, w, h), device=noise.device)) * discount**i
        if w==1 or h==1: break # Lowest resolution is 1x1
    return noise/noise.std() # Scaled back to roughly unit variance