(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
unconditional_guidance_scale=1., unconditional_conditioning=None,
dynamic_threshold=None)
| 179 | |
| 180 | @torch.no_grad() |
| 181 | def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False, |
| 182 | temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, |
| 183 | unconditional_guidance_scale=1., unconditional_conditioning=None, |
| 184 | dynamic_threshold=None): |
| 185 | b, *_, device = *x.shape, x.device |
| 186 | |
| 187 | if unconditional_conditioning is None or unconditional_guidance_scale == 1.: |
| 188 | model_output = self.model.apply_model(x, t, c) |
| 189 | else: |
| 190 | model_t = self.model.apply_model(x, t, c) |
| 191 | model_uncond = self.model.apply_model(x, t, unconditional_conditioning) |
| 192 | model_output = model_uncond + unconditional_guidance_scale * (model_t - model_uncond) |
| 193 | |
| 194 | if self.model.parameterization == "v": |
| 195 | e_t = self.model.predict_eps_from_z_and_v(x, t, model_output) |
| 196 | else: |
| 197 | e_t = model_output |
| 198 | |
| 199 | if score_corrector is not None: |
| 200 | assert self.model.parameterization == "eps", 'not implemented' |
| 201 | e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs) |
| 202 | |
| 203 | alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas |
| 204 | alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev |
| 205 | sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas |
| 206 | sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas |
| 207 | # select parameters corresponding to the currently considered timestep |
| 208 | a_t = torch.full((b, 1, 1, 1), alphas[index], device=device) |
| 209 | a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device) |
| 210 | sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device) |
| 211 | sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device) |
| 212 | |
| 213 | # current prediction for x_0 |
| 214 | if self.model.parameterization != "v": |
| 215 | pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt() |
| 216 | else: |
| 217 | pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output) |
| 218 | |
| 219 | if quantize_denoised: |
| 220 | pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0) |
| 221 | |
| 222 | if dynamic_threshold is not None: |
| 223 | raise NotImplementedError() |
| 224 | |
| 225 | # direction pointing to x_t |
| 226 | dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t |
| 227 | noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature |
| 228 | if noise_dropout > 0.: |
| 229 | noise = torch.nn.functional.dropout(noise, p=noise_dropout) |
| 230 | x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise |
| 231 | return x_prev, pred_x0 |
| 232 | |
| 233 | @torch.no_grad() |
| 234 | def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None, |
no test coverage detected