(
self,
x,
t,
clip_denoised=True,
repeat_noise=False,
condition_x=None,
self_cond=None,
get_interm_fm=False,
)
| 416 | |
| 417 | @torch.no_grad() |
| 418 | def p_sample( |
| 419 | self, |
| 420 | x, |
| 421 | t, |
| 422 | clip_denoised=True, |
| 423 | repeat_noise=False, |
| 424 | condition_x=None, |
| 425 | self_cond=None, |
| 426 | get_interm_fm=False, |
| 427 | ): |
| 428 | b, *_, device = *x.shape, x.device |
| 429 | model_mean, _, model_log_variance, _ = self.p_mean_variance( |
| 430 | x=x, |
| 431 | t=t, |
| 432 | clip_denoised=clip_denoised, |
| 433 | condition_x=condition_x, |
| 434 | self_cond=self_cond, |
| 435 | get_interm_fm=get_interm_fm, |
| 436 | ) |
| 437 | |
| 438 | noise = noise_like(x.shape, device, repeat_noise) |
| 439 | # no noise when t == 0 |
| 440 | # t shape is (b, ) -> (b, 1, 1, 1) |
| 441 | nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) |
| 442 | return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise |
| 443 | |
| 444 | @torch.no_grad() |
| 445 | def p_sample_loop(self, x_in, continous=False, get_interm_fm=False): |
no test coverage detected