File size: 1,358 Bytes
b094968 |
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 |
from modules import devices
import torch
def get_delta(latent):
mean = torch.mean(latent)
return torch.sub(latent, mean)
def to_abs(latent):
return torch.abs(latent)
def zeros(latent):
return torch.zeros_like(latent)
def ones(latent):
return torch.ones_like(latent)
def gaussian_noise(latent):
return torch.rand_like(latent)
def normal_noise(latent):
return torch.randn_like(latent)
def multires_noise(latent, use_zero:bool, iterations=8, discount=0.4):
"""
Reference: https://wandb.ai/johnowhitaker/multires_noise/reports/Multi-Resolution-Noise-for-Diffusion-Model-Training--VmlldzozNjYyOTU2
Credit: Kohya_SS
"""
noise = zeros(latent) if use_zero else ones(latent)
batchSize = noise.size(0)
height = noise.size(2)
width = noise.size(3)
device = devices.get_optimal_device()
upsampler = torch.nn.Upsample(size=(height, width), mode="bilinear").to(device)
for b in range(batchSize):
for i in range(iterations):
r = torch.rand(1).item() * 2 + 2
wn = max(1, int(width / (r**i)))
hn = max(1, int(height / (r**i)))
for c in range(4):
noise[b, c] += upsampler(torch.randn(1, 1, hn, wn).to(device))[0, 0] * discount**i
if wn == 1 or hn == 1:
break
return noise / noise.std()
|