| 34 | |
| 35 | |
| 36 | class ForwardDiffusion(nn.Module): |
| 37 | def __init__(self, |
| 38 | im_size: int = 64, |
| 39 | n_diffusion_timesteps: int = 1000): |
| 40 | super().__init__() |
| 41 | self.n_diffusion_timesteps = n_diffusion_timesteps |
| 42 | cos_alpha_bar_t = shifted_cosine_alpha_bar( |
| 43 | np.linspace(0, 1, n_diffusion_timesteps), |
| 44 | im_size=im_size |
| 45 | ).astype(np.float32) |
| 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 """ |
no outgoing calls
no test coverage detected