Spaces:
Running on Zero
Running on Zero
File size: 2,927 Bytes
a95f6c0 | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | import torch
from diff_params.shared import SDE
class EDM(SDE):
"""
Definition of the diffusion parameterization, following ( Karras et al., "Elucidating...", 2022).
This includes only the utilities needed for training, not for sampling.
"""
def __init__(self,
type,
sde_hp):
super().__init__(type, sde_hp)
self.sigma_data = self.sde_hp.sigma_data #depends on the training data!! precalculated variance of the dataset
self.sigma_min = self.sde_hp.sigma_min
self.sigma_max = self.sde_hp.sigma_max
self.rho = self.sde_hp.rho
def sample_time_training(self, N):
"""
For training, getting t according to a similar criteria as sampling. Simpler and safer to what Karras et al. did
Args:
N (int): batch size
"""
a = torch.rand(N)
t = (self.sigma_max**(1/self.rho) +a *(self.sigma_min**(1/self.rho) - self.sigma_max**(1/self.rho)))**self.rho
return t
def sample_prior(self, shape):
"""
Just sample some gaussian noise, nothing more
Args:
shape (tuple): shape of the noise to sample, something like (B,T)
"""
n = torch.randn(shape)
return n
def cskip(self, sigma):
"""
Just one of the preconditioning parameters
Args:
sigma (float): noise level (equal to timestep is sigma=t, which is our default)
"""
return self.sigma_data**2 *(sigma**2+self.sigma_data**2)**-1
def cout(self, sigma):
"""
Just one of the preconditioning parameters
Args:
sigma (float): noise level (equal to timestep is sigma=t, which is our default)
"""
return sigma*self.sigma_data* (self.sigma_data**2+sigma**2)**(-0.5)
def cin(self, sigma):
"""
Just one of the preconditioning parameters
Args:
sigma (float): noise level (equal to timestep is sigma=t, which is our default)
"""
return (self.sigma_data**2+sigma**2)**(-0.5)
def cnoise(self, sigma):
"""
preconditioning of the noise embedding
Args:
sigma (float): noise level (equal to timestep is sigma=t, which is our default)
"""
return (1/4)*torch.log(sigma)
def lambda_w(self, sigma):
"""
Score matching loss weighting
"""
return (sigma*self.sigma_data)**(-2) * (self.sigma_data**2+sigma**2)
def Tweedie2score(self, tweedie, xt, t, *args, **kwargs):
return (tweedie - self._mean(xt, t)) / self._std(t)**2
def score2Tweedie(self, score, xt, t, *args, **kwargs):
return self._std(t)**2 * score + self._mean(xt, t)
def _mean(self, x, t):
return x
def _std(self, t):
return t
def _ode_integrand(self, x, t, score):
return -t * score |