import math import torch from tqdm.auto import trange from tqdm import trange, tqdm import torch.distributions as dist try: import torchsde except ImportError as e: print(f"Failed to import torchsde. You won't be able to use the GPT model. {e}") torchsde = None ##### # code here is taken from k_diffusion, exposed for future modification ##### def get_alphas_sigmas(t): """Returns the scaling factors for the clean image (alpha) and for the noise (sigma), given a timestep.""" return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2) @torch.no_grad() def sample_ddim(model, x, steps, eta=0.3, disable=False, **extra_args): """DDIM sampling from a model given starting noise. v-diffusion""" ts = x.new_ones([x.shape[0]]) # Create the noise schedule t = torch.linspace(1, 0, steps + 1)[:-1] alphas, sigmas = get_alphas_sigmas(t) # The sampling loop for i in trange(steps, disable=disable): # Get the model output (v, the predicted velocity) v = model.forward_inference(x, ts * t[i], **extra_args) # Predict the noise and the denoised image pred = x * alphas[i] - v * sigmas[i] eps = x * sigmas[i] + v * alphas[i] # If we are not on the last timestep, compute the noisy image for the # next timestep. if i < steps - 1: # If eta > 0, adjust the scaling factor for the predicted noise # downward according to the amount of additional noise to add ddim_sigma = ( eta * (sigmas[i + 1] ** 2 / sigmas[i] ** 2).sqrt() * (1 - alphas[i] ** 2 / alphas[i + 1] ** 2).sqrt() ) adjusted_sigma = (sigmas[i + 1] ** 2 - ddim_sigma**2).sqrt() # Recombine the predicted noise and predicted denoised image in the # correct proportions for the next step x = pred * alphas[i + 1] + eps * adjusted_sigma # Add the correct amount of fresh noise if eta: x += torch.randn_like(x) * ddim_sigma # If we are on the last timestep, output the denoised image return pred def append_dims(x, target_dims): """Appends dimensions to the end of a tensor until it has target_dims dimensions.""" dims_to_append = target_dims - x.ndim if dims_to_append < 0: raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less") return x[(...,) + (None,) * dims_to_append] class VDenoiser(torch.nn.Module): """A v-diffusion-pytorch model wrapper for k-diffusion.""" def __init__(self, inner_model): super().__init__() self.inner_model = inner_model self.sigma_data = 1.0 def get_scalings(self, sigma): c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2) c_out = -sigma * self.sigma_data / (sigma**2 + self.sigma_data**2) ** 0.5 c_in = 1 / (sigma**2 + self.sigma_data**2) ** 0.5 return c_skip, c_out, c_in def sigma_to_t(self, sigma): return sigma.atan() / math.pi * 2 def t_to_sigma(self, t): return (t * math.pi / 2).tan() def forward(self, input, sigma, **kwargs): c_skip, c_out, c_in = [append_dims(x, input.ndim) for x in self.get_scalings(sigma)] if "text_codes" in kwargs and "semantic_codes" in kwargs: # used in training for metrics out = self.inner_model.forward(input * c_in, self.sigma_to_t(sigma), **kwargs) else: kwargs_1 = { k: v for k, v in kwargs.items() if k not in [ "cross_attn_cond_2", "no_text_cross_attn_cond_2", "no_sem_cross_attn_cond_2", "no_ctx_cross_attn_cond_2", ] } out = self.inner_model.forward_inference(input * c_in, self.sigma_to_t(sigma), **kwargs_1) return out * c_out + input * c_skip class BatchedBrownianTree: """A wrapper around torchsde.BrownianTree that enables batches of entropy.""" def __init__(self, x, t0, t1, seed=None, **kwargs): t0, t1, self.sign = self.sort(t0, t1) w0 = kwargs.get("w0", torch.zeros_like(x)) if seed is None: seed = torch.randint(0, 2**63 - 1, []).item() self.batched = True try: assert len(seed) == x.shape[0] w0 = w0[0] except TypeError: seed = [seed] self.batched = False self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed] @staticmethod def sort(a, b): return (a, b, 1) if a < b else (b, a, -1) def __call__(self, t0, t1): t0, t1, sign = self.sort(t0, t1) w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign) return w if self.batched else w[0] class BrownianTreeNoiseSampler: """A noise sampler backed by a torchsde.BrownianTree. Args: x (Tensor): The tensor whose shape, device and dtype to use to generate random samples. sigma_min (float): The low end of the valid interval. sigma_max (float): The high end of the valid interval. seed (int or List[int]): The random seed. If a list of seeds is supplied instead of a single integer, then the noise sampler will use one BrownianTree per batch item, each with its own seed. transform (callable): A function that maps sigma to the sampler's internal timestep. """ def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x): self.transform = transform t0, t1 = ( self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max)), ) self.tree = BatchedBrownianTree(x, t0, t1, seed) def __call__(self, sigma, sigma_next): t0, t1 = ( self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next)), ) return self.tree(t0, t1) / (t1 - t0).abs().sqrt() def append_zero(x): return torch.cat([x, x.new_zeros([1])]) def get_sigmas_polyexponential(n, sigma_min, sigma_max, rho=1.0, device="cpu"): """Constructs an polynomial in log sigma noise schedule.""" ramp = torch.linspace(1, 0, n, device=device) ** rho sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min)) return append_zero(sigmas) @torch.no_grad() def sample_dpmpp_3m_sde( model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1.0, noise_sampler=None, ): """DPM-Solver++(3M) SDE.""" sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() noise_sampler = ( BrownianTreeNoiseSampler(x, sigma_min, sigma_max) if noise_sampler is None else noise_sampler ) extra_args = {} if extra_args is None else extra_args s_in = x.new_ones([x.shape[0]]) denoised_1, denoised_2 = None, None h_1, h_2 = None, None for i in trange(len(sigmas) - 1, disable=disable): denoised = model(x, sigmas[i] * s_in, **extra_args) if callback is not None: callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised}) if sigmas[i + 1] == 0: # Denoising step x = denoised else: t, s = -sigmas[i].log(), -sigmas[i + 1].log() h = s - t h_eta = h * (eta + 1) x = torch.exp(-h_eta) * x + (-h_eta).expm1().neg() * denoised if h_2 is not None: r0 = h_1 / h r1 = h_2 / h d1_0 = (denoised - denoised_1) / r0 d1_1 = (denoised_1 - denoised_2) / r1 d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1) d2 = (d1_0 - d1_1) / (r0 + r1) phi_2 = h_eta.neg().expm1() / h_eta + 1 phi_3 = phi_2 / h_eta - 0.5 x = x + phi_2 * d1 - phi_3 * d2 elif h_1 is not None: r = h_1 / h d = (denoised - denoised_1) / r phi_2 = h_eta.neg().expm1() / h_eta + 1 x = x + phi_2 * d if eta: x = ( x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise ) denoised_1, denoised_2 = denoised, denoised_1 h_1, h_2 = h, h_1 return x @torch.no_grad() def sample_discrete_euler( model, x, steps=None, sigma_max=1, sigmas=None, callback=None, dist_shift=None, disable_tqdm=False, **extra_args, ): """Draws samples from a model given starting noise. Euler method""" assert steps is not None or sigmas is not None, "Either steps or sigmas must be provided" # Make tensor of ones to broadcast the single t values ts = x.new_ones([x.shape[0]]) if sigmas is None: # Create the noise schedule t = torch.linspace(sigma_max, 0, steps + 1) if dist_shift is not None: t = dist_shift.time_shift(t, x.shape[-1]) else: t = sigmas # alphas, sigmas = 1-t, t for i, (t_curr, t_prev) in enumerate(tqdm(zip(t[:-1], t[1:]), disable=disable_tqdm)): # Broadcast the current timestep to the correct shape t_curr_tensor = t_curr * torch.ones((x.shape[0],), dtype=x.dtype, device=x.device) dt = t_prev - t_curr # we solve backwards in our formulation v = model(x, t_curr_tensor, **extra_args) x = x + dt * v if callback is not None: denoised = x - t_prev * v callback({"x": x, "t": t_curr, "sigma": t_curr, "i": i + 1, "denoised": denoised}) # If we are on the last timestep, output the denoised data return x @torch.no_grad() def sample_rk4( model, x, steps=None, sigma_max=1, sigmas=None, callback=None, dist_shift=None, **extra_args ): """Draws samples from a model given starting noise. 4th-order Runge-Kutta""" assert steps is not None or sigmas is not None, "Either steps or sigmas must be provided" # Make tensor of ones to broadcast the single t values ts = x.new_ones([x.shape[0]]) if sigmas is None: # Create the noise schedule t = torch.linspace(sigma_max, 0, steps + 1) if dist_shift is not None: t = dist_shift.time_shift(t, x.shape[-1]) else: t = sigmas # alphas, sigmas = 1-t, t for i, (t_curr, t_prev) in enumerate(tqdm(zip(t[:-1], t[1:]))): # Broadcast the current timestep to the correct shape t_curr_tensor = t_curr * ts dt = t_prev - t_curr # we solve backwards in our formulation k1 = model(x, t_curr_tensor, **extra_args) k2 = model(x + dt / 2 * k1, (t_curr + dt / 2) * ts, **extra_args) k3 = model(x + dt / 2 * k2, (t_curr + dt / 2) * ts, **extra_args) k4 = model(x + dt * k3, t_prev * ts, **extra_args) x = x + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4) if callback is not None: denoised = x - t_prev * k4 callback({"x": x, "t": t_curr, "sigma": t_curr, "i": i + 1, "denoised": denoised}) # If we are on the last timestep, output the denoised data return x @torch.no_grad() def sample_flow_dpmpp( model, x, steps=None, sigma_max=1, sigmas=None, callback=None, dist_shift=None, disable_tqdm=False, **extra_args, ): """Draws samples from a model given starting noise. DPM-Solver++ for RF models""" assert steps is not None or sigmas is not None, "Either steps or sigmas must be provided" # Make tensor of ones to broadcast the single t values ts = x.new_ones([x.shape[0]]) if sigmas is None: # Create the noise schedule t = torch.linspace(sigma_max, 0, steps + 1) if dist_shift is not None: t = dist_shift.time_shift(t, x.shape[-1]) else: t = sigmas old_denoised = None log_snr = lambda t: ((1 - t) / t).log() for i in trange(len(t) - 1, disable=disable_tqdm): t_curr, t_next = t[i], t[i + 1] denoised = x - t_curr * model(x, t_curr * ts, **extra_args) if callback is not None: callback({"x": x, "i": i, "t": t_curr, "sigma": t_curr, "denoised": denoised}) alpha_t = 1 - t_next h = log_snr(t_next) - log_snr(t_curr) if old_denoised is None or t_next == 0: x = (t_next / t_curr) * x - alpha_t * (-h).expm1() * denoised else: h_last = log_snr(t_curr) - log_snr(t[i - 1]) r = h_last / h denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised x = (t_next / t_curr) * x - alpha_t * (-h).expm1() * denoised_d old_denoised = denoised return x @torch.no_grad() def sample_flow_pingpong( model, x, steps=None, sigma_max=1, sigmas=None, callback=None, dist_shift=None, **extra_args ): """Draws samples from a model given starting noise. Ping-pong sampling for distilled models""" assert steps is not None or sigmas is not None, "Either steps or sigmas must be provided" # Make tensor of ones to broadcast the single t values ts = x.new_ones([x.shape[0]]) if sigmas is None: # Create the noise schedule t = torch.linspace(sigma_max, 0, steps + 1) if dist_shift is not None: t = dist_shift.time_shift(t, x.shape[-1]) else: t = sigmas for i in trange(len(t) - 1, disable=False): denoised = x - t[i] * model(x, t[i] * ts, **extra_args) if callback is not None: callback({"x": x, "i": i, "t": t[i], "sigma": t[i], "sigma_hat": t[i], "denoised": denoised}) t_next = t[i + 1] x = (1 - t_next) * denoised + t_next * torch.randn_like(x) return x # init_data is init_audio as latents (if this is latent diffusion) # For sampling, set both init_data and mask to None # For variations, set init_data def sample_rf( model_fn, noise, init_data=None, steps=100, sampler_type="euler", sigma_max=1, device="cuda", callback=None, cond_fn=None, **extra_args, ): if sigma_max > 1: sigma_max = 1 if cond_fn is not None: denoiser = make_cond_model_fn(denoiser, cond_fn) if init_data is not None: # VARIATION # Interpolate the init data and the noise for init audio x = init_data * (1 - sigma_max) + noise * sigma_max else: # SAMPLING # set the initial latent to noise x = noise logsnr_max = math.log(((1 - sigma_max) / sigma_max) + 1e-6) if sigma_max < 1 else -6 logsnr = torch.linspace(logsnr_max, 2, steps + 1) t = torch.sigmoid(-logsnr) t[0] = sigma_max t[-1] = 0 if sampler_type == "euler": return sample_discrete_euler( model_fn, x, sigmas=t, sigma_max=sigma_max, callback=callback, **extra_args ) elif sampler_type == "rk4": return sample_rk4(model_fn, x, steps, sigma_max, callback=callback, **extra_args) elif sampler_type == "dpmpp": return sample_flow_dpmpp( model_fn, x, sigmas=t, sigma_max=sigma_max, callback=callback, **extra_args ) elif sampler_type == "pingpong": return sample_flow_pingpong( model_fn, x, sigmas=t, sigma_max=sigma_max, callback=callback, **extra_args )