Spaces:
Sleeping
Sleeping
| import torch | |
| from inference.sampler import Sampler | |
| class SamplerEulerHeun(Sampler): | |
| def __init__(self, model, diff_params, args): | |
| super().__init__(model, diff_params, args) | |
| # stochasticity parameters | |
| self.Schurn = self.args.tester.sampling_params.Schurn | |
| self.Snoise = self.args.tester.sampling_params.Snoise | |
| self.Stmin = self.args.tester.sampling_params.Stmin | |
| self.Stmax = self.args.tester.sampling_params.Stmax | |
| # order of the sampler | |
| self.order = self.args.tester.sampling_params.order | |
| self.cond=None | |
| self.cfg_scale = 1.0 | |
| def predict_DPS( | |
| self, | |
| shape, # observations (lowpssed signal) Tensor with shape ?? | |
| cond=None, | |
| cfg_scale=1.0, | |
| device=None, # device | |
| apply_inverse_transform=True, # whether to apply inverse transform | |
| taxonomy=None, # taxonomy for the conditional input | |
| masks=None, # masks for the conditional input | |
| fwd_operator=None, | |
| zeta=None, | |
| dtype=torch.float32, # data type | |
| ): | |
| self.cond = cond | |
| assert self.cond is not None, "Conditional input is None" | |
| self.taxonomy = taxonomy | |
| self.masks = masks | |
| self.cfg_scale = cfg_scale | |
| # get the noise schedule | |
| t = self.create_schedule().to(device).to(torch.float32) | |
| # sample prior | |
| x = self.diff_params.sample_prior(t=t[0], shape=shape, dtype=dtype) | |
| # parameter for langevin stochasticity, if Schurn is 0, gamma will be 0 to, so the sampler will be deterministic | |
| gamma = self.get_gamma(t).to(device) | |
| for i in range(0, self.T-1, 1): | |
| self.step_counter = i | |
| x, x_den = self.step_DPS(x, t[i], t[i + 1], gamma[i], fwd_operator, zeta) | |
| if apply_inverse_transform: | |
| with torch.no_grad(): | |
| x_den_wave=self.diff_params.transform_inverse(x_den.detach()) | |
| return x_den_wave.detach(), None | |
| else: | |
| return x_den.detach(), None | |
| def predict_conditional( | |
| self, | |
| shape, # observations (lowpssed signal) Tensor with shape ?? | |
| cond=None, | |
| cfg_scale=1.0, | |
| device=None, # device | |
| apply_inverse_transform=True, # whether to apply inverse transform | |
| taxonomy=None, # taxonomy for the conditional input | |
| masks=None, # masks for the conditional input | |
| dtype=torch.float32, # data type | |
| ): | |
| self.cond = cond | |
| assert self.cond is not None, "Conditional input is None" | |
| self.taxonomy = taxonomy | |
| self.masks = masks | |
| self.cfg_scale = cfg_scale | |
| # get the noise schedule | |
| t = self.create_schedule().to(device).to(torch.float32) | |
| # sample prior | |
| x = self.diff_params.sample_prior(t=t[0], shape=shape, dtype=dtype) | |
| # parameter for langevin stochasticity, if Schurn is 0, gamma will be 0 to, so the sampler will be deterministic | |
| gamma = self.get_gamma(t).to(device) | |
| for i in range(0, self.T-1, 1): | |
| self.step_counter = i | |
| x, x_den = self.step(x, t[i], t[i + 1], gamma[i]) | |
| if apply_inverse_transform: | |
| with torch.no_grad(): | |
| x_den_wave=self.diff_params.transform_inverse(x_den.detach()) | |
| return x_den_wave.detach(), None | |
| else: | |
| return x_den.detach(), None | |
| def predict_unconditional( | |
| self, | |
| shape, # observations (lowpssed signal) Tensor with shape ?? | |
| device | |
| ): | |
| self.y = None | |
| self.degradation = None | |
| return self.predict(shape, device) | |
| def get_gamma(self, t): | |
| """ | |
| Get the parameter gamma that defines the stochasticity of the sampler | |
| Args | |
| t (Tensor): shape: (N_steps, ) Tensor of timesteps, from which we will compute gamma | |
| """ | |
| N = t.shape[0] | |
| gamma = torch.zeros(t.shape).to(t.device) | |
| # If desired, only apply stochasticity between a certain range of noises Stmin is 0 by default and Stmax is a huge number by default. (Unless these parameters are specified, this does nothing) | |
| indexes = torch.logical_and(t > self.Stmin, t < self.Stmax) | |
| # We use Schurn=5 as the default in our experiments | |
| gamma[indexes] = gamma[indexes] + torch.min(torch.Tensor([self.Schurn / N, 2 ** (1 / 2) - 1])) | |
| return gamma | |
| def get_Tweedie_estimate(self, x, t_i): | |
| if x.ndim==2: | |
| x_=x.unsqueeze(1) | |
| elif x.ndim==3: | |
| pass | |
| if self.cond is not None: | |
| x_hat = self.diff_params.denoiser(x, self.model, t_i, cond=self.cond, cfg_scale=self.cfg_scale, taxonomy=self.taxonomy, masks=self.masks) | |
| else: | |
| x_hat = self.diff_params.denoiser(x, self.model, t_i, taxonomy=self.taxonomy, masks=self.masks) | |
| return x_hat | |
| def Tweedie2score(self, tweedie, xt, t): | |
| return self.diff_params.Tweedie2score(tweedie, xt, t) | |
| def score2Tweedie(self, score, xt, t): | |
| return self.diff_params.score2Tweedie(score, xt, t) | |
| def stochastic_timestep(self, x, t, gamma, Snoise=1): | |
| t_hat = t + gamma * t # if gamma_sig[i]==0 this is a deterministic step, make sure it doed not crash | |
| t_hat=torch.clamp(t_hat, 0, self.diff_params.max_t) | |
| epsilon = torch.randn(x.shape).to(x.device) * Snoise # sample Gaussiannoise, Snoise is 1 by default | |
| if t_hat <= t: | |
| x_hat = x | |
| #print(f"t_hat<=t, gamma {gamma}") | |
| else: | |
| #print(t_hat, t) | |
| x_hat = x + ((t_hat ** 2 - t ** 2) ** (1 / 2)) * epsilon # Perturb data | |
| return x_hat, t_hat | |
| def step_DPS(self, x_i, t_i, t_iplus1, gamma_i , fwd_operator=None, zeta=None): | |
| #with torch.no_grad(): | |
| x_hat, t_hat=self.stochastic_timestep(x_i, t_i, gamma_i) | |
| x_hat.requires_grad=True | |
| x_den = self.get_Tweedie_estimate(x_hat, t_hat) | |
| #optionally L2 normalize here... | |
| #compute likelihood score | |
| loss=fwd_operator(x_den) | |
| loss.backward(retain_graph=False) | |
| grads= x_hat.grad | |
| #lets normalize the grads | |
| norm_factor= torch.sqrt(torch.tensor(x_hat.view(-1).shape[0])).to(x_hat.device) | |
| normguide=torch.norm(grads)/ norm_factor | |
| zeta= zeta/(normguide+1e-8) | |
| lh_score=-zeta*grads/t_hat | |
| x_hat.detach_() # detach x_hat to avoid accumulating gradients | |
| x_den.detach_() # detach x_den to avoid accumulating gradients | |
| #compute normal score | |
| score = self.Tweedie2score(x_den, x_hat, t_hat) | |
| ode_integrand = self.diff_params._ode_integrand(x_hat, t_hat, score+ lh_score) | |
| dt = t_iplus1 - t_hat | |
| x_iplus1 = x_hat + dt * ode_integrand | |
| return x_iplus1, x_den | |
| def step(self, x_i, t_i, t_iplus1, gamma_i ): | |
| with torch.no_grad(): | |
| x_hat, t_hat = self.stochastic_timestep(x_i, t_i, gamma_i) | |
| x_den = self.get_Tweedie_estimate(x_hat, t_hat) | |
| score = self.Tweedie2score(x_den, x_hat, t_hat) | |
| ode_integrand = self.diff_params._ode_integrand(x_hat, t_hat, score) | |
| dt = t_iplus1 - t_hat | |
| if t_iplus1 != 0 and self.order == 2: # second order correction | |
| t_prime = t_iplus1 | |
| x_prime = x_hat + dt * ode_integrand | |
| x_den = self.get_Tweedie_estimate(x_prime, t_prime) | |
| score = self.Tweedie2score(x_den, x_prime, t_prime) | |
| ode_integrand_next = self.diff_params._ode_integrand(x_prime, t_prime, score) | |
| ode_integrand_midpoint = .5 * (ode_integrand + ode_integrand_next) | |
| x_iplus1 = x_hat + dt * ode_integrand_midpoint | |
| else: | |
| x_iplus1 = x_hat + dt * ode_integrand | |
| return x_iplus1, x_den | |
| def get_domain_shape(self, shape, device): | |
| x=torch.zeros(shape, dtype=torch.float32).to(device) | |
| X=self.diff_params.transform_forward(x) | |
| return X.shape, X.dtype | |
| def predict( | |
| self, | |
| shape, # observations (lowpssed signal) Tensor with shape ?? | |
| device, # lambda function | |
| dtype=torch.float32, # data type | |
| apply_inverse_transform=True # whether to apply inverse transform | |
| ): | |
| # get the noise schedule | |
| t = self.create_schedule().to(device).to(torch.float32) | |
| # sample prior | |
| x = self.diff_params.sample_prior(t=t[0], shape=shape, dtype=dtype) | |
| # parameter for langevin stochasticity, if Schurn is 0, gamma will be 0 to, so the sampler will be deterministic | |
| gamma = self.get_gamma(t).to(device) | |
| for i in range(0, self.T-1, 1): | |
| self.step_counter = i | |
| x, x_den = self.step(x, t[i], t[i + 1], gamma[i]) | |
| if apply_inverse_transform: | |
| with torch.no_grad(): | |
| x_den_wave=self.diff_params.transform_inverse(x_den.detach()) | |
| return x_den_wave.detach(), None | |
| else: | |
| return x_den.detach(), None | |
| def create_schedule(self, sigma_min=None, sigma_max=None, rho=None, T=None): | |
| """ | |
| EDM schedule by default | |
| """ | |
| if T is None: | |
| T=self.T | |
| if self.args.tester.sampling_params.schedule == "edm": | |
| 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 | |
| a = torch.arange(0, T) | |
| 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.schedule == "song": | |
| 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 | |
| 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 | |
| return t | |
| elif self.args.tester.sampling_params.schedule == "FM": | |
| t = torch.linspace(1, 0, T+1) | |
| return t | |
| else: | |
| raise NotImplementedError(f"schedule {self.args.tester.sampling_params.schedule} not implemented") | |