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()