Spaces:
Sleeping
Sleeping
File size: 2,222 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 |
import torch
import abc
class Sampler():
def __init__(self, model, diff_params, args):
self.model = model.eval() #is it ok to do this here?
self.diff_params = diff_params #same as training, useful if we need to apply a wrapper or something
self.args=args
if self.args.tester.sampling_params.same_as_training:
self.sde_hp = diff_params.sde_hp
else:
self.sde_hp = self.args.tester.sampling_params.sde_hp
self.T = self.args.tester.sampling_params.T
self.step_counter = 0
@abc.abstractmethod
def predict(self, *args, **kwargs):
pass
@abc.abstractmethod
def predict_unconditional(self, *args, **kwargs):
pass
@abc.abstractmethod
def predict_conditional(self, *args, **kwargs):
pass
@abc.abstractmethod
def step(self, *args, **kwargs):
pass
def create_schedule(self, sigma_min=None, sigma_max=None, rho=None, T=None):
"""
EDM schedule by default
"""
if sigma_min is None:
sigma_min = self.sde_hp.sigma_min
if sigma_max is None:
sigma_max = self.sde_hp.sigma_max
if rho is None:
rho = self.sde_hp.rho
if T is None:
T=self.T
if self.args.tester.sampling_params.schedule == "edm":
a = torch.arange(0, T+1)
t = (sigma_max**(1/rho) + a/(T-1) *(sigma_min**(1/rho) - sigma_max**(1/rho)))**rho
t[-1] = 0
return t
elif self.args.tester.sampling_params == "song":
eps = 0. if not "t_eps" in self.args.tester.diff_params.keys() else self.args.tester.diff_params.t_eps
a = torch.arange(eps, T+1)
t = sigma_min**2 * (sigma_max / sigma_min)**(2*a)
t[-1] = 0
else:
raise NotImplementedError(f"schedule {self.args.tester.posterior_sampling.RED.schedule} not implemented")
def Tweedie2score(self, tweedie, xt, t):
return self.diff_params.Tweedie2score(tweedie, xt, t)
def get_Tweedie_estimate(self, x, t_i):
x_hat = self.diff_params.denoiser(x.unsqueeze(1), self.model, t_i).squeeze(1)
return x_hat
|