Sample from the model
(self, denoise_fn, data, t, noise_fn, clip_denoised=False, return_pred_xstart=False)
| 231 | ''' samples ''' |
| 232 | |
| 233 | def p_sample(self, denoise_fn, data, t, noise_fn, clip_denoised=False, return_pred_xstart=False): |
| 234 | """ |
| 235 | Sample from the model |
| 236 | """ |
| 237 | model_mean, _, model_log_variance, pred_xstart = self.p_mean_variance(denoise_fn, data=data, t=t, clip_denoised=clip_denoised, |
| 238 | return_pred_xstart=True) |
| 239 | noise = noise_fn(size=data.shape, dtype=data.dtype, device=data.device) |
| 240 | assert noise.shape == data.shape |
| 241 | # no noise when t == 0 |
| 242 | nonzero_mask = torch.reshape(1 - (t == 0).float(), [data.shape[0]] + [1] * (len(data.shape) - 1)) |
| 243 | |
| 244 | sample = model_mean + nonzero_mask * torch.exp(0.5 * model_log_variance) * noise |
| 245 | assert sample.shape == pred_xstart.shape |
| 246 | return (sample, pred_xstart) if return_pred_xstart else sample |
| 247 | |
| 248 | |
| 249 | def p_sample_loop(self, denoise_fn, shape, device, |
no test coverage detected