| 176 | |
| 177 | @torch.no_grad() |
| 178 | def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False, |
| 179 | temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, |
| 180 | unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None, |
| 181 | dynamic_threshold=None): |
| 182 | b, *_, device = *x.shape, x.device |
| 183 | |
| 184 | def get_model_output(x, t): |
| 185 | if unconditional_conditioning is None or unconditional_guidance_scale == 1.: |
| 186 | e_t = self.model.apply_model(x, t, c) |
| 187 | else: |
| 188 | x_in = torch.cat([x] * 2) |
| 189 | t_in = torch.cat([t] * 2) |
| 190 | c_in = torch.cat([unconditional_conditioning, c]) |
| 191 | e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2) |
| 192 | e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond) |
| 193 | |
| 194 | if score_corrector is not None: |
| 195 | assert self.model.parameterization == "eps" |
| 196 | e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs) |
| 197 | |
| 198 | return e_t |
| 199 | |
| 200 | alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas |
| 201 | alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev |
| 202 | sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas |
| 203 | sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas |
| 204 | |
| 205 | def get_x_prev_and_pred_x0(e_t, index): |
| 206 | # select parameters corresponding to the currently considered timestep |
| 207 | a_t = torch.full((b, 1, 1, 1), alphas[index], device=device) |
| 208 | a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device) |
| 209 | sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device) |
| 210 | sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device) |
| 211 | |
| 212 | # current prediction for x_0 |
| 213 | pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt() |
| 214 | if quantize_denoised: |
| 215 | pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0) |
| 216 | if dynamic_threshold is not None: |
| 217 | pred_x0 = norm_thresholding(pred_x0, dynamic_threshold) |
| 218 | # direction pointing to x_t |
| 219 | dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t |
| 220 | noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature |
| 221 | if noise_dropout > 0.: |
| 222 | noise = torch.nn.functional.dropout(noise, p=noise_dropout) |
| 223 | x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise |
| 224 | return x_prev, pred_x0 |
| 225 | |
| 226 | e_t = get_model_output(x, t) |
| 227 | if len(old_eps) == 0: |
| 228 | # Pseudo Improved Euler (2nd order) |
| 229 | x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index) |
| 230 | e_t_next = get_model_output(x_prev, t_next) |
| 231 | e_t_prime = (e_t + e_t_next) / 2 |
| 232 | elif len(old_eps) == 1: |
| 233 | # 2nd order Pseudo Linear Multistep (Adams-Bashforth) |
| 234 | e_t_prime = (3 * e_t - old_eps[-1]) / 2 |
| 235 | elif len(old_eps) == 2: |