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