(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,
unconditional_guidance_scale=1.0, unconditional_conditioning=None, callback=None)
| 256 | |
| 257 | @torch.no_grad() |
| 258 | def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None, |
| 259 | unconditional_guidance_scale=1.0, unconditional_conditioning=None, callback=None): |
| 260 | num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0] |
| 261 | |
| 262 | assert t_enc <= num_reference_steps |
| 263 | num_steps = t_enc |
| 264 | |
| 265 | if use_original_steps: |
| 266 | alphas_next = self.alphas_cumprod[:num_steps] |
| 267 | alphas = self.alphas_cumprod_prev[:num_steps] |
| 268 | else: |
| 269 | alphas_next = self.ddim_alphas[:num_steps] |
| 270 | alphas = torch.tensor(self.ddim_alphas_prev[:num_steps]) |
| 271 | |
| 272 | x_next = x0 |
| 273 | intermediates = [] |
| 274 | inter_steps = [] |
| 275 | for i in tqdm(range(num_steps), desc='Encoding Image'): |
| 276 | t = torch.full((x0.shape[0],), i, device=self.model.device, dtype=torch.long) |
| 277 | if unconditional_guidance_scale == 1.: |
| 278 | noise_pred = self.model.apply_model(x_next, t, c) |
| 279 | else: |
| 280 | assert unconditional_conditioning is not None |
| 281 | e_t_uncond, noise_pred = torch.chunk( |
| 282 | self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)), |
| 283 | torch.cat((unconditional_conditioning, c))), 2) |
| 284 | noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond) |
| 285 | |
| 286 | xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next |
| 287 | weighted_noise_pred = alphas_next[i].sqrt() * ( |
| 288 | (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred |
| 289 | x_next = xt_weighted + weighted_noise_pred |
| 290 | if return_intermediates and i % ( |
| 291 | num_steps // return_intermediates) == 0 and i < num_steps - 1: |
| 292 | intermediates.append(x_next) |
| 293 | inter_steps.append(i) |
| 294 | elif return_intermediates and i >= num_steps - 2: |
| 295 | intermediates.append(x_next) |
| 296 | inter_steps.append(i) |
| 297 | if callback: callback(i) |
| 298 | |
| 299 | out = {'x_encoded': x_next, 'intermediate_steps': inter_steps} |
| 300 | if return_intermediates: |
| 301 | out.update({'intermediates': intermediates}) |
| 302 | return x_next, out |
| 303 | |
| 304 | @torch.no_grad() |
| 305 | def stochastic_encode(self, x0, t, use_original_steps=False, noise=None): |
no test coverage detected