| import torch | |
| class GaussianDiffusion(nn.Module): | |
| def __init__(self, model, timesteps=1000, beta_start=1e-4, beta_end=0.02): | |
| super(GaussianDiffusion, self).__init__() | |
| self.model = model | |
| self.timesteps = timesteps | |
| # Create a schedule of betas (noise variance at each timestep) | |
| self.betas = torch.linspace(beta_start, beta_end, timesteps) | |
| self.alphas = 1 - self.betas | |
| self.alpha_cumprod = torch.cumprod(self.alphas, dim=0) | |
| def add | |