Diffuse the data (t == 0 means diffused for 1 step)
(self, x_start, t, noise=None)
| 111 | return mean, variance, log_variance |
| 112 | |
| 113 | def q_sample(self, x_start, t, noise=None): |
| 114 | """ |
| 115 | Diffuse the data (t == 0 means diffused for 1 step) |
| 116 | """ |
| 117 | if noise is None: |
| 118 | noise = torch.randn(x_start.shape, device=x_start.device) |
| 119 | assert noise.shape == x_start.shape |
| 120 | return ( |
| 121 | self._extract(self.sqrt_alphas_cumprod.to(x_start.device), t, x_start.shape) * x_start + |
| 122 | self._extract(self.sqrt_one_minus_alphas_cumprod.to(x_start.device), t, x_start.shape) * noise |
| 123 | ) |
| 124 | |
| 125 | |
| 126 | def q_posterior_mean_variance(self, x_start, x_t, t): |