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, constrain_fn=lambda x, t:x,
clip_denoised=True, max_timestep=None, keep_running=False)
| 203 | |
| 204 | |
| 205 | def p_sample_loop(self, denoise_fn, shape, device, |
| 206 | noise_fn=torch.randn, constrain_fn=lambda x, t:x, |
| 207 | clip_denoised=True, max_timestep=None, keep_running=False): |
| 208 | """ |
| 209 | Generate samples |
| 210 | keep_running: True if we run 2 x num_timesteps, False if we just run num_timesteps |
| 211 | |
| 212 | """ |
| 213 | if max_timestep is None: |
| 214 | final_time = self.num_timesteps |
| 215 | else: |
| 216 | final_time = max_timestep |
| 217 | |
| 218 | assert isinstance(shape, (tuple, list)) |
| 219 | img_t = noise_fn(size=shape, dtype=torch.float, device=device) |
| 220 | for t in reversed(range(0, final_time if not keep_running else len(self.betas))): |
| 221 | img_t = constrain_fn(img_t, t) |
| 222 | t_ = torch.empty(shape[0], dtype=torch.int64, device=device).fill_(t) |
| 223 | img_t = self.p_sample(denoise_fn=denoise_fn, data=img_t,t=t_, noise_fn=noise_fn, |
| 224 | clip_denoised=clip_denoised, return_pred_xstart=False).detach() |
| 225 | |
| 226 | |
| 227 | assert img_t.shape == shape |
| 228 | return img_t |
| 229 | |
| 230 | def reconstruct(self, x0, t, denoise_fn, noise_fn=torch.randn, constrain_fn=lambda x, t:x): |
| 231 |
no test coverage detected