Spaces:
Running on Zero
Running on Zero
| 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 | |
| def predict(self, *args, **kwargs): | |
| pass | |
| def predict_unconditional(self, *args, **kwargs): | |
| pass | |
| def predict_conditional(self, *args, **kwargs): | |
| pass | |
| 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 | |