Method __init__

diff2flow/diffusion.py:37–46

(self,
                 im_size: int = 64,
                 n_diffusion_timesteps: int = 1000)

class ForwardDiffusion(nn.Module):
    def __init__(self,
                 im_size: int = 64,
                 n_diffusion_timesteps: int = 1000):
        super().__init__()
        self.n_diffusion_timesteps = n_diffusion_timesteps
        cos_alpha_bar_t = shifted_cosine_alpha_bar(
            np.linspace(0, 1, n_diffusion_timesteps),
            im_size=im_size
        ).astype(np.float32)
        self.register_buffer("alpha_bar_t", torch.from_numpy(cos_alpha_bar_t))
        
    def q_sample(self, x_start, t, noise=None):
        """