(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,
unconditional_guidance_scale=1.0, unconditional_conditioning=None, callback=None)
| 232 | |
| 233 | @torch.no_grad() |
| 234 | def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None, |
| 235 | unconditional_guidance_scale=1.0, unconditional_conditioning=None, callback=None): |
| 236 | timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps |
| 237 | num_reference_steps = timesteps.shape[0] |
| 238 | |
| 239 | assert t_enc <= num_reference_steps |
| 240 | num_steps = t_enc |
| 241 | |
| 242 | if use_original_steps: |
| 243 | alphas_next = self.alphas_cumprod[:num_steps] |
| 244 | alphas = self.alphas_cumprod_prev[:num_steps] |
| 245 | else: |
| 246 | alphas_next = self.ddim_alphas[:num_steps] |
| 247 | alphas = torch.tensor(self.ddim_alphas_prev[:num_steps]) |
| 248 | |
| 249 | x_next = x0 |
| 250 | intermediates = [] |
| 251 | inter_steps = [] |
| 252 | for i in tqdm(range(num_steps), desc='Encoding Image'): |
| 253 | t = torch.full((x0.shape[0],), timesteps[i], device=self.model.device, dtype=torch.long) |
| 254 | if unconditional_guidance_scale == 1.: |
| 255 | noise_pred = self.model.apply_model(x_next, t, c) |
| 256 | else: |
| 257 | assert unconditional_conditioning is not None |
| 258 | e_t_uncond, noise_pred = torch.chunk( |
| 259 | self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)), |
| 260 | torch.cat((unconditional_conditioning, c))), 2) |
| 261 | noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond) |
| 262 | |
| 263 | xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next |
| 264 | weighted_noise_pred = alphas_next[i].sqrt() * ( |
| 265 | (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred |
| 266 | x_next = xt_weighted + weighted_noise_pred |
| 267 | if return_intermediates and i % ( |
| 268 | num_steps // return_intermediates) == 0 and i < num_steps - 1: |
| 269 | intermediates.append(x_next) |
| 270 | inter_steps.append(i) |
| 271 | elif return_intermediates and i >= num_steps - 2: |
| 272 | intermediates.append(x_next) |
| 273 | inter_steps.append(i) |
| 274 | if callback: callback(i) |
| 275 | |
| 276 | out = {'x_encoded': x_next, 'intermediate_steps': inter_steps} |
| 277 | if return_intermediates: |
| 278 | out.update({'intermediates': intermediates}) |
| 279 | return x_next, out |
| 280 | |
| 281 | @torch.no_grad() |
| 282 | def stochastic_encode(self, x0, t, use_original_steps=False, noise=None): |
nothing calls this directly
no test coverage detected