| import math |
|
|
| import torch |
| from torch import nn |
|
|
| from . import sampling, utils |
|
|
|
|
| class VDenoiser(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. |
|
|
| 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 loss(self, input, noise, sigma, **kwargs): |
| c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] |
| noised_input = input + noise * utils.append_dims(sigma, input.ndim) |
| model_output = self.inner_model(noised_input * c_in, self.sigma_to_t(sigma), **kwargs) |
| target = (input - c_skip * noised_input) / c_out |
| return (model_output - target).pow(2).flatten(1).mean(1) |
|
|
| def forward(self, input, sigma, **kwargs): |
| c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] |
| return self.inner_model(input * c_in, self.sigma_to_t(sigma), **kwargs) * c_out + input * c_skip |
|
|
|
|
| class DiscreteSchedule(nn.Module): |
| """A mapping between continuous noise levels (sigmas) and a list of discrete noise |
| levels.""" |
|
|
| def __init__(self, sigmas, quantize): |
| super().__init__() |
| self.register_buffer('sigmas', sigmas) |
| self.register_buffer('log_sigmas', sigmas.log()) |
| self.quantize = quantize |
|
|
| @property |
| def sigma_min(self): |
| return self.sigmas[0] |
|
|
| @property |
| def sigma_max(self): |
| return self.sigmas[-1] |
|
|
| def get_sigmas(self, n=None): |
| if n is None: |
| return sampling.append_zero(self.sigmas.flip(0)) |
| t_max = len(self.sigmas) - 1 |
| t = torch.linspace(t_max, 0, n, device=self.sigmas.device) |
| return sampling.append_zero(self.t_to_sigma(t)) |
|
|
| def sigma_to_discrete_timestep(self, sigma): |
| log_sigma = sigma.log() |
| dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None] |
| return dists.abs().argmin(dim=0).view(sigma.shape) |
|
|
| def sigma_to_t(self, sigma, quantize=None): |
| quantize = self.quantize if quantize is None else quantize |
| if quantize: |
| return self.sigma_to_discrete_timestep(sigma) |
| log_sigma = sigma.log() |
| dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None] |
| low_idx = dists.ge(0).cumsum(dim=0).argmax(dim=0).clamp(max=self.log_sigmas.shape[0] - 2) |
| high_idx = low_idx + 1 |
| low, high = self.log_sigmas[low_idx], self.log_sigmas[high_idx] |
| w = (low - log_sigma) / (low - high) |
| w = w.clamp(0, 1) |
| t = (1 - w) * low_idx + w * high_idx |
| return t.view(sigma.shape) |
|
|
| def t_to_sigma(self, t): |
| t = t.float() |
| low_idx = t.floor().long() |
| high_idx = t.ceil().long() |
| w = t-low_idx if t.device.type == 'mps' else t.frac() |
| log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx] |
| return log_sigma.exp() |
|
|
| def predict_eps_discrete_timestep(self, input, t, **kwargs): |
| if t.dtype != torch.int64 and t.dtype != torch.int32: |
| t = t.round() |
| sigma = self.t_to_sigma(t) |
| input = input * ((utils.append_dims(sigma, input.ndim) ** 2 + 1.0) ** 0.5) |
| return (input - self(input, sigma, **kwargs)) / utils.append_dims(sigma, input.ndim) |
|
|
| class DiscreteEpsDDPMDenoiser(DiscreteSchedule): |
| """A wrapper for discrete schedule DDPM models that output eps (the predicted |
| noise).""" |
|
|
| def __init__(self, model, alphas_cumprod, quantize): |
| super().__init__(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, quantize) |
| self.inner_model = model |
| self.sigma_data = 1. |
|
|
| def get_scalings(self, sigma): |
| c_out = -sigma |
| c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5 |
| return c_out, c_in |
|
|
| def get_eps(self, *args, **kwargs): |
| return self.inner_model(*args, **kwargs) |
|
|
| def loss(self, input, noise, sigma, **kwargs): |
| c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] |
| noised_input = input + noise * utils.append_dims(sigma, input.ndim) |
| eps = self.get_eps(noised_input * c_in, self.sigma_to_t(sigma), **kwargs) |
| return (eps - noise).pow(2).flatten(1).mean(1) |
|
|
| def forward(self, input, sigma, **kwargs): |
| c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] |
| eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs) |
| return input + eps * c_out |
|
|
|
|
| class OpenAIDenoiser(DiscreteEpsDDPMDenoiser): |
| """A wrapper for OpenAI diffusion models.""" |
|
|
| def __init__(self, model, diffusion, quantize=False, has_learned_sigmas=True, device='cpu'): |
| alphas_cumprod = torch.tensor(diffusion.alphas_cumprod, device=device, dtype=torch.float32) |
| super().__init__(model, alphas_cumprod, quantize=quantize) |
| self.has_learned_sigmas = has_learned_sigmas |
|
|
| def get_eps(self, *args, **kwargs): |
| model_output = self.inner_model(*args, **kwargs) |
| if self.has_learned_sigmas: |
| return model_output.chunk(2, dim=1)[0] |
| return model_output |
|
|
|
|
| class CompVisDenoiser(DiscreteEpsDDPMDenoiser): |
| """A wrapper for CompVis diffusion models.""" |
|
|
| def __init__(self, model, quantize=False, device='cpu'): |
| super().__init__(model, model.alphas_cumprod, quantize=quantize) |
|
|
| def get_eps(self, *args, **kwargs): |
| return self.inner_model.apply_model(*args, **kwargs) |
|
|
|
|
| class DiscreteVDDPMDenoiser(DiscreteSchedule): |
| """A wrapper for discrete schedule DDPM models that output v.""" |
|
|
| def __init__(self, model, alphas_cumprod, quantize): |
| super().__init__(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, quantize) |
| self.inner_model = model |
| self.sigma_data = 1. |
|
|
| 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 get_v(self, *args, **kwargs): |
| return self.inner_model(*args, **kwargs) |
|
|
| def loss(self, input, noise, sigma, **kwargs): |
| c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] |
| noised_input = input + noise * utils.append_dims(sigma, input.ndim) |
| model_output = self.get_v(noised_input * c_in, self.sigma_to_t(sigma), **kwargs) |
| target = (input - c_skip * noised_input) / c_out |
| return (model_output - target).pow(2).flatten(1).mean(1) |
|
|
| def forward(self, input, sigma, **kwargs): |
| c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] |
| return self.get_v(input * c_in, self.sigma_to_t(sigma), **kwargs) * c_out + input * c_skip |
|
|
|
|
| class CompVisVDenoiser(DiscreteVDDPMDenoiser): |
| """A wrapper for CompVis diffusion models that output v.""" |
|
|
| def __init__(self, model, quantize=False, device='cpu'): |
| super().__init__(model, model.alphas_cumprod, quantize=quantize) |
|
|
| def get_v(self, x, t, cond, **kwargs): |
| return self.inner_model.apply_model(x, t, cond) |
|
|
