51 lines
1.3 KiB
Python
51 lines
1.3 KiB
Python
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()
|