File size: 2,717 Bytes
864ec44 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
# Author: Bingxin Ke
# Last modified: 2024-04-18
import torch
import math
# adapted from: https://wandb.ai/johnowhitaker/multires_noise/reports/Multi-Resolution-Noise-for-Diffusion-Model-Training--VmlldzozNjYyOTU2?s=31
def multi_res_noise_like(
x, strength=0.9, downscale_strategy="original", generator=None, device=None
):
if torch.is_tensor(strength):
strength = strength.reshape((-1, 1, 1, 1))
b, c, w, h = x.shape
if device is None:
device = x.device
up_sampler = torch.nn.Upsample(size=(w, h), mode="bilinear")
noise = torch.randn(x.shape, device=x.device, generator=generator)
if "original" == downscale_strategy:
for i in range(10):
r = (
torch.rand(1, generator=generator, device=device) * 2 + 2
) # Rather than always going 2x,
w, h = max(1, int(w / (r**i))), max(1, int(h / (r**i)))
noise += (
up_sampler(
torch.randn(b, c, w, h, generator=generator, device=device).to(x)
)
* strength**i
)
if w == 1 or h == 1:
break # Lowest resolution is 1x1
elif "every_layer" == downscale_strategy:
for i in range(int(math.log2(min(w, h)))):
w, h = max(1, int(w / 2)), max(1, int(h / 2))
noise += (
up_sampler(
torch.randn(b, c, w, h, generator=generator, device=device).to(x)
)
* strength**i
)
elif "power_of_two" == downscale_strategy:
for i in range(10):
r = 2
w, h = max(1, int(w / (r**i))), max(1, int(h / (r**i)))
noise += (
up_sampler(
torch.randn(b, c, w, h, generator=generator, device=device).to(x)
)
* strength**i
)
if w == 1 or h == 1:
break # Lowest resolution is 1x1
elif "random_step" == downscale_strategy:
for i in range(10):
r = (
torch.rand(1, generator=generator, device=device) * 2 + 2
) # Rather than always going 2x,
w, h = max(1, int(w / (r))), max(1, int(h / (r)))
noise += (
up_sampler(
torch.randn(b, c, w, h, generator=generator, device=device).to(x)
)
* strength**i
)
if w == 1 or h == 1:
break # Lowest resolution is 1x1
else:
raise ValueError(f"unknown downscale strategy: {downscale_strategy}")
noise = noise / noise.std() # Scaled back to roughly unit variance
return noise
|