stable-diffusion-webui
137 строк · 5.6 Кб
1import torch
2import tqdm
3import k_diffusion.sampling
4import numpy as np
5
6from modules import shared
7from modules.models.diffusion.uni_pc import uni_pc
8
9
10@torch.no_grad()
11def ddim(model, x, timesteps, extra_args=None, callback=None, disable=None, eta=0.0):
12alphas_cumprod = model.inner_model.inner_model.alphas_cumprod
13alphas = alphas_cumprod[timesteps]
14alphas_prev = alphas_cumprod[torch.nn.functional.pad(timesteps[:-1], pad=(1, 0))].to(torch.float64 if x.device.type != 'mps' and x.device.type != 'xpu' else torch.float32)
15sqrt_one_minus_alphas = torch.sqrt(1 - alphas)
16sigmas = eta * np.sqrt((1 - alphas_prev.cpu().numpy()) / (1 - alphas.cpu()) * (1 - alphas.cpu() / alphas_prev.cpu().numpy()))
17
18extra_args = {} if extra_args is None else extra_args
19s_in = x.new_ones((x.shape[0]))
20s_x = x.new_ones((x.shape[0], 1, 1, 1))
21for i in tqdm.trange(len(timesteps) - 1, disable=disable):
22index = len(timesteps) - 1 - i
23
24e_t = model(x, timesteps[index].item() * s_in, **extra_args)
25
26a_t = alphas[index].item() * s_x
27a_prev = alphas_prev[index].item() * s_x
28sigma_t = sigmas[index].item() * s_x
29sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_x
30
31pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
32dir_xt = (1. - a_prev - sigma_t ** 2).sqrt() * e_t
33noise = sigma_t * k_diffusion.sampling.torch.randn_like(x)
34x = a_prev.sqrt() * pred_x0 + dir_xt + noise
35
36if callback is not None:
37callback({'x': x, 'i': i, 'sigma': 0, 'sigma_hat': 0, 'denoised': pred_x0})
38
39return x
40
41
42@torch.no_grad()
43def plms(model, x, timesteps, extra_args=None, callback=None, disable=None):
44alphas_cumprod = model.inner_model.inner_model.alphas_cumprod
45alphas = alphas_cumprod[timesteps]
46alphas_prev = alphas_cumprod[torch.nn.functional.pad(timesteps[:-1], pad=(1, 0))].to(torch.float64 if x.device.type != 'mps' and x.device.type != 'xpu' else torch.float32)
47sqrt_one_minus_alphas = torch.sqrt(1 - alphas)
48
49extra_args = {} if extra_args is None else extra_args
50s_in = x.new_ones([x.shape[0]])
51s_x = x.new_ones((x.shape[0], 1, 1, 1))
52old_eps = []
53
54def get_x_prev_and_pred_x0(e_t, index):
55# select parameters corresponding to the currently considered timestep
56a_t = alphas[index].item() * s_x
57a_prev = alphas_prev[index].item() * s_x
58sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_x
59
60# current prediction for x_0
61pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
62
63# direction pointing to x_t
64dir_xt = (1. - a_prev).sqrt() * e_t
65x_prev = a_prev.sqrt() * pred_x0 + dir_xt
66return x_prev, pred_x0
67
68for i in tqdm.trange(len(timesteps) - 1, disable=disable):
69index = len(timesteps) - 1 - i
70ts = timesteps[index].item() * s_in
71t_next = timesteps[max(index - 1, 0)].item() * s_in
72
73e_t = model(x, ts, **extra_args)
74
75if len(old_eps) == 0:
76# Pseudo Improved Euler (2nd order)
77x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
78e_t_next = model(x_prev, t_next, **extra_args)
79e_t_prime = (e_t + e_t_next) / 2
80elif len(old_eps) == 1:
81# 2nd order Pseudo Linear Multistep (Adams-Bashforth)
82e_t_prime = (3 * e_t - old_eps[-1]) / 2
83elif len(old_eps) == 2:
84# 3nd order Pseudo Linear Multistep (Adams-Bashforth)
85e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
86else:
87# 4nd order Pseudo Linear Multistep (Adams-Bashforth)
88e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
89
90x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
91
92old_eps.append(e_t)
93if len(old_eps) >= 4:
94old_eps.pop(0)
95
96x = x_prev
97
98if callback is not None:
99callback({'x': x, 'i': i, 'sigma': 0, 'sigma_hat': 0, 'denoised': pred_x0})
100
101return x
102
103
104class UniPCCFG(uni_pc.UniPC):
105def __init__(self, cfg_model, extra_args, callback, *args, **kwargs):
106super().__init__(None, *args, **kwargs)
107
108def after_update(x, model_x):
109callback({'x': x, 'i': self.index, 'sigma': 0, 'sigma_hat': 0, 'denoised': model_x})
110self.index += 1
111
112self.cfg_model = cfg_model
113self.extra_args = extra_args
114self.callback = callback
115self.index = 0
116self.after_update = after_update
117
118def get_model_input_time(self, t_continuous):
119return (t_continuous - 1. / self.noise_schedule.total_N) * 1000.
120
121def model(self, x, t):
122t_input = self.get_model_input_time(t)
123
124res = self.cfg_model(x, t_input, **self.extra_args)
125
126return res
127
128
129def unipc(model, x, timesteps, extra_args=None, callback=None, disable=None, is_img2img=False):
130alphas_cumprod = model.inner_model.inner_model.alphas_cumprod
131
132ns = uni_pc.NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
133t_start = timesteps[-1] / 1000 + 1 / 1000 if is_img2img else None # this is likely off by a bit - if someone wants to fix it please by all means
134unipc_sampler = UniPCCFG(model, extra_args, callback, ns, predict_x0=True, thresholding=False, variant=shared.opts.uni_pc_variant)
135x = unipc_sampler.sample(x, steps=len(timesteps), t_start=t_start, skip_type=shared.opts.uni_pc_skip_type, method="multistep", order=shared.opts.uni_pc_order, lower_order_final=shared.opts.uni_pc_lower_order_final)
136
137return x
138