Diffuse the data for a given number of diffusion steps. In other words sample from q(x_t | x_0). Args: x_start: The initial data batch. t: The diffusion time-step (must be a single t). noise: If specified, the noise to use for the diffusi
(self, x_start, t, noise=None)
| 46 | self.register_buffer("alpha_bar_t", torch.from_numpy(cos_alpha_bar_t)) |
| 47 | |
| 48 | def q_sample(self, x_start, t, noise=None): |
| 49 | """ |
| 50 | Diffuse the data for a given number of diffusion steps. In other |
| 51 | words sample from q(x_t | x_0). |
| 52 | |
| 53 | Args: |
| 54 | x_start: The initial data batch. |
| 55 | t: The diffusion time-step (must be a single t). |
| 56 | noise: If specified, the noise to use for the diffusion. |
| 57 | Returns: |
| 58 | A noisy version of x_start. |
| 59 | """ |
| 60 | if noise is None: |
| 61 | noise = torch.randn_like(x_start) |
| 62 | |
| 63 | alpha_bar_t = self.alpha_bar_t[t] |
| 64 | |
| 65 | return torch.sqrt(alpha_bar_t) * x_start + torch.sqrt(1 - alpha_bar_t) * noise |
| 66 | |
| 67 | |
| 68 | """ Diffusion Wrapper adapted to FlowModel """ |