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