stable-diffusion-webui
104 строки · 3.8 Кб
1import torch2
3from k_diffusion import utils, sampling4from k_diffusion.external import DiscreteEpsDDPMDenoiser5from k_diffusion.sampling import default_noise_sampler, trange6
7from modules import shared, sd_samplers_cfg_denoiser, sd_samplers_kdiffusion, sd_samplers_common8
9
10class LCMCompVisDenoiser(DiscreteEpsDDPMDenoiser):11def __init__(self, model):12timesteps = 100013original_timesteps = 50 # LCM Original Timesteps (default=50, for current version of LCM)14self.skip_steps = timesteps // original_timesteps15
16alphas_cumprod_valid = torch.zeros((original_timesteps), dtype=torch.float32, device=model.device)17for x in range(original_timesteps):18alphas_cumprod_valid[original_timesteps - 1 - x] = model.alphas_cumprod[timesteps - 1 - x * self.skip_steps]19
20super().__init__(model, alphas_cumprod_valid, quantize=None)21
22
23def get_sigmas(self, n=None,):24if n is None:25return sampling.append_zero(self.sigmas.flip(0))26
27start = self.sigma_to_t(self.sigma_max)28end = self.sigma_to_t(self.sigma_min)29
30t = torch.linspace(start, end, n, device=shared.sd_model.device)31
32return sampling.append_zero(self.t_to_sigma(t))33
34
35def sigma_to_t(self, sigma, quantize=None):36log_sigma = sigma.log()37dists = log_sigma - self.log_sigmas[:, None]38return dists.abs().argmin(dim=0).view(sigma.shape) * self.skip_steps + (self.skip_steps - 1)39
40
41def t_to_sigma(self, timestep):42t = torch.clamp(((timestep - (self.skip_steps - 1)) / self.skip_steps).float(), min=0, max=(len(self.sigmas) - 1))43return super().t_to_sigma(t)44
45
46def get_eps(self, *args, **kwargs):47return self.inner_model.apply_model(*args, **kwargs)48
49
50def get_scaled_out(self, sigma, output, input):51sigma_data = 0.552scaled_timestep = utils.append_dims(self.sigma_to_t(sigma), output.ndim) * 10.053
54c_skip = sigma_data**2 / (scaled_timestep**2 + sigma_data**2)55c_out = scaled_timestep / (scaled_timestep**2 + sigma_data**2) ** 0.556
57return c_out * output + c_skip * input58
59
60def forward(self, input, sigma, **kwargs):61c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]62eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)63return self.get_scaled_out(sigma, input + eps * c_out, input)64
65
66def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):67extra_args = {} if extra_args is None else extra_args68noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler69s_in = x.new_ones([x.shape[0]])70
71for i in trange(len(sigmas) - 1, disable=disable):72denoised = model(x, sigmas[i] * s_in, **extra_args)73
74if callback is not None:75callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})76
77x = denoised78if sigmas[i + 1] > 0:79x += sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1])80return x81
82
83class CFGDenoiserLCM(sd_samplers_cfg_denoiser.CFGDenoiser):84@property85def inner_model(self):86if self.model_wrap is None:87denoiser = LCMCompVisDenoiser88self.model_wrap = denoiser(shared.sd_model)89
90return self.model_wrap91
92
93class LCMSampler(sd_samplers_kdiffusion.KDiffusionSampler):94def __init__(self, funcname, sd_model, options=None):95super().__init__(funcname, sd_model, options)96self.model_wrap_cfg = CFGDenoiserLCM(self)97self.model_wrap = self.model_wrap_cfg.inner_model98
99
100samplers_lcm = [('LCM', sample_lcm, ['k_lcm'], {})]101samplers_data_lcm = [102sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: LCMSampler(funcname, model), aliases, options)103for label, funcname, aliases, options in samplers_lcm104]
105