Generate samples keep_running: True if we run 2 x num_timesteps, False if we just run num_timesteps
(self, denoise_fn, shape, device,
noise_fn=torch.randn, clip_denoised=True, keep_running=False)
| 247 | |
| 248 | |
| 249 | def p_sample_loop(self, denoise_fn, shape, device, |
| 250 | noise_fn=torch.randn, clip_denoised=True, keep_running=False): |
| 251 | """ |
| 252 | Generate samples |
| 253 | keep_running: True if we run 2 x num_timesteps, False if we just run num_timesteps |
| 254 | |
| 255 | """ |
| 256 | |
| 257 | assert isinstance(shape, (tuple, list)) |
| 258 | img_t = noise_fn(size=shape, dtype=torch.float, device=device) |
| 259 | for t in reversed(range(0, self.num_timesteps if not keep_running else len(self.betas))): |
| 260 | t_ = torch.empty(shape[0], dtype=torch.int64, device=device).fill_(t) |
| 261 | img_t = self.p_sample(denoise_fn=denoise_fn, data=img_t,t=t_, noise_fn=noise_fn, |
| 262 | clip_denoised=clip_denoised, return_pred_xstart=False) |
| 263 | |
| 264 | assert img_t.shape == shape |
| 265 | return img_t |
| 266 | |
| 267 | def p_sample_loop_trajectory(self, denoise_fn, shape, device, freq, |
| 268 | noise_fn=torch.randn,clip_denoised=True, keep_running=False): |