Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models

class GaussianDiffusion(nn.Module):
    def __init__(
        self,
        model,
        *,
        image_size,
        timesteps = 1000,
        sampling_timesteps = None,
        ddim_sampling_eta = 0.,
    ):
        super().__init__()
        self.model = model # 用于拟合\epsilon(x_t,t)的神经网络
        self.channels = self.model.channels
        self.image_size = image_size

        betas = linear_beta_schedule(timesteps) # \beta_t
        alphas = 1. - betas # \alpha_t
        alphas_cumprod = torch.cumprod(alphas, dim=0) # \bar{\alpha}_t
        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.) # \bar{\alpha}_{t-1}

        self.num_timesteps = int(timesteps)
        self.alphas_cumprod = alphas_cumprod

        # sampling related parameters
        self.sampling_timesteps = default(sampling_timesteps, timesteps) # default num sampling timesteps to number of timesteps at training
        assert self.sampling_timesteps <= timesteps
        self.ddim_sampling_eta = ddim_sampling_eta
        
        # helper function to register buffer from float64 to float32
        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
        register_buffer('betas', betas)
        
        # calculations for diffusion q(x_t | x_{t-1}) and others
        register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod)) # \sqrt{\bar{\alpha}_t}
        register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod)) # \sqrt{1-\bar{\alpha}_t}
        register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod)) # \log{1-\bar{\alpha}_t}
        register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod)) # \sqrt{1/\bar{\alpha}_t}
        register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1)) # \sqrt{1/(\bar{\alpha}_t-1)}

        # calculations for posterior q(x_{t-1} | x_t, x_0)
        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) # \sigma_t
        register_buffer('posterior_variance', posterior_variance)
        # log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
        register_buffer('posterior_log_variance_clipped', torch.log(posterior_variance.clamp(min =1e-20)))
        register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
        register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod))

    """
    Training
    """

    # 计算x_t=\sqrt{\bat{\alpha}_t}x_0+\sqrt{1-\bat{\alpha}_t}\epsilon
    def q_sample(self, x_start, t, noise=None):
        return (
            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
            extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
        )
    
    # 计算损失函数L_t=||\epsilon-\epsilon(x_t, t)||_1
    def p_losses(self, x_start, t, noise = None):
        b, c, h, w = x_start.shape
        noise = torch.randn_like(x_start)
        target = noise

        # 计算 x_t
        x = self.q_sample(x_start = x_start, t = t, noise = noise)

        # 计算 \epsilon(x_t, t)
        model_out = self.model(x, t)

        loss = F.l1_loss(model_out, target, reduction = 'none')
        loss = reduce(loss, 'b ... -> b (...)', 'mean')
        return loss.mean()

    # 训练过程
    def forward(self, img, *args, **kwargs):
        b, c, h, w, device = *img.shape, img.device
        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()

        img = img * 2 - 1 # data [0, 1] -> [-1, 1]
        return self.p_losses(img, t, *args, **kwargs)

    """
    Sampling
    """
    
    def data_generator(self, batch, folder='./images/official-artwork'):
        transform = T.Compose([
            T.Resize(self.image_size),
            T.RandomHorizontalFlip(),
            T.ToTensor()
        ])

        selected_imgs = random.sample(os.listdir(folder), k=batch)
        batch_imgs = []
        for img in selected_imgs:
            img = Image.open(os.path.join(folder, img))
            batch_imgs.append(transform(img).unsqueeze(0))
        return torch.cat(batch_imgs, dim=0) * 2 - 1
            
    # 计算 x_0 = \sqrt{1/\bat{\alpha}_t}x_t-\sqrt{1/(1-\bat{\alpha}_t)}\epsilon_t
    def predict_start_from_noise(self, x_t, t, noise):
        return (
            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
            extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
        )

    @torch.no_grad()
    def ddim_sample(self, shape, return_all_timesteps = False):
        batch, device, total_timesteps, sampling_timesteps, eta = shape[0], self.betas.device, self.num_timesteps, self.sampling_timesteps, self.ddim_sampling_eta
        
        # 用(batch_size * steps)个样本去估计方差修正项
        T = torch.linspace(0, total_timesteps-1, steps=total_timesteps, device=device)
        T = list(T.int().tolist())
        factors = []
        for t in tqdm(T, ncols=0, desc = 'computing unbiased term'):
            batch_imgs = self.data_generator(batch).to(device)
            t = torch.tensor(t, device=device).long().repeat(batch)
            factor_t = torch.mean((self.model(self.q_sample(x_start = batch_imgs, t = t, noise = torch.randn(shape, device = device)), t))**2)
            factors.append(factor_t)
        
        times = torch.linspace(-1, total_timesteps - 1, steps = sampling_timesteps + 1)   # [-1, 0, 1, 2, ..., T-1] when sampling_timesteps == total_timesteps
        times = list(reversed(times.int().tolist()))
        time_pairs = list(zip(times[:-1], times[1:])) # [(T-1, T-2), (T-2, T-3), ..., (1, 0), (0, -1)]

        img = torch.randn(shape, device = device)
        imgs = [img]

        x_start = None

        for time, time_next in tqdm(time_pairs, desc = 'sampling loop time step'):
            time_cond = torch.full((batch,), time, device = device, dtype = torch.long)
            pred_noise = self.model(img, time_cond) # 计算 \epsilon(x_t, t)
            x_start = self.predict_start_from_noise(img, time_cond, pred_noise) # x_0
            if time_next < 0:
                img = x_start
                imgs.append(img)
                continue

            alpha = self.alphas_cumprod[time]
            alpha_next = self.alphas_cumprod[time_next]
            factor = factors[time]

            sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
            c = (1 - alpha_next - sigma ** 2).sqrt()
            
            unbias = (alpha_next.sqrt() - c * (alpha / (1-alpha)).sqrt()) ** 2 * (1-alpha) / alpha * factor

            noise = torch.randn_like(img)
            img = x_start * alpha_next.sqrt() + \
                  c * pred_noise + \
                  (sigma**2 + unbias).sqrt() * noise
            imgs.append(img)

        ret = img if not return_all_timesteps else torch.stack(imgs, dim = 1)
        ret = (ret + 1) * 0.5
        return ret
    
    # 采样过程
    @torch.no_grad()
    def sample(self, batch_size = 16, img_channel = 3, return_all_timesteps = False):
        image_size, channels = self.image_size, img_channel
        return self.ddim_sample((batch_size, channels, image_size, image_size), return_all_timesteps = return_all_timesteps)