| 340 | |
| 341 | @torch.no_grad() |
| 342 | def p_sample_loop(self, shape, return_intermediates=False): |
| 343 | device = self.betas.device |
| 344 | b = shape[0] |
| 345 | img = torch.randn(shape, device=device) |
| 346 | intermediates = [img] |
| 347 | for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps): |
| 348 | img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long), |
| 349 | clip_denoised=self.clip_denoised) |
| 350 | if i % self.log_every_t == 0 or i == self.num_timesteps - 1: |
| 351 | intermediates.append(img) |
| 352 | if return_intermediates: |
| 353 | return img, intermediates |
| 354 | return img |
| 355 | |
| 356 | @torch.no_grad() |
| 357 | def sample(self, batch_size=16, return_intermediates=False): |