(
self,
x,
t,
repeat_noise=False,
clip_denoised=False,
condition_x=None,
self_cond=None,
eta=0.0,
)
| 593 | |
| 594 | @torch.no_grad() |
| 595 | def ddim_sample( |
| 596 | self, |
| 597 | x, |
| 598 | t, |
| 599 | repeat_noise=False, |
| 600 | clip_denoised=False, |
| 601 | condition_x=None, |
| 602 | self_cond=None, |
| 603 | eta=0.0, |
| 604 | ): |
| 605 | *_, pred_x_0 = self.p_mean_variance(x, t, clip_denoised, condition_x, self_cond) |
| 606 | eps = self.predict_noise_from_start(x, t, pred_x_0) |
| 607 | alpha_cumprod = extract(self.alphas_cumprod, t, x.shape) |
| 608 | alpha_cumprod_prev = extract(self.alphas_cumprod_prev, t, x.shape) |
| 609 | sigma = ( |
| 610 | eta |
| 611 | * torch.sqrt((1 - alpha_cumprod_prev) / (1 - alpha_cumprod)) |
| 612 | * torch.sqrt(1 - alpha_cumprod / alpha_cumprod_prev) |
| 613 | ) |
| 614 | noise = noise_like(x.shape, device=x.device, repeat=repeat_noise) |
| 615 | mean_pred = ( |
| 616 | pred_x_0 * torch.sqrt(alpha_cumprod_prev) |
| 617 | + torch.sqrt(1 - alpha_cumprod_prev - sigma ** 2) * eps |
| 618 | ) |
| 619 | nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) |
| 620 | sample = mean_pred + nonzero_mask * sigma * noise |
| 621 | return sample |
| 622 | |
| 623 | @torch.no_grad() |
| 624 | def ddim_sample_loop(self, x_in, section_counts="ddim300", eta=0.0): |
no test coverage detected