DDIM和DDPM代码

### 关于DDIM和DDPM的代码实现

对于典型的扩散模型（DDPM），其训练过程涉及逐步向数据添加噪声并学习逆转这一过程。而在推理阶段，由于每一步都需要采样新的随机噪声，因此速度相对较慢[^1]。

相比之下，去噪扩散隐式模型（DDIM）通过采用确定性的前向传播方式，在不牺牲质量的前提下显著提高了推理效率，通常能达到DDPM十倍以上的速度。

#### DDPM 实现示例
以下是基于PyTorch框架的一个简化版DDPM实现：

import torch
from torch import nn
import numpy as np

class DDPM(nn.Module):
    def __init__(self, model, betas):
        super().__init__()
        self.model = model
        
        # 定义beta参数表
        self.betas = betas
        self.num_timesteps = len(betas)

        alphas = 1. - betas
        alpha_cumprod = np.cumprod(alphas, axis=0)
        
        to_torch = lambda x: torch.tensor(x).float()
        self.register_buffer('sqrt_alpha_cumprod', to_torch(np.sqrt(alpha_cumprod)))
        self.register_buffer('sqrt_one_minus_alpha_cumprod', to_torch(np.sqrt(1. - alpha_cumprod)))

    @torch.no_grad()
    def p_sample(self, xt, t, noise=None):    
        beta_t = extract(self.betas, t, xt.shape)
        sqrt_one_minus_at = extract(self.sqrt_one_minus_alpha_cumprod, t, xt.shape)
        pred_noise = self.model(xt, t)
        
        if isinstance(noise,bool) and not noise:
            z = 0.
        elif noise is None:
            z = torch.randn_like(xt)
            
        return (
                (xt - ((beta_t ** 0.5)*pred_noise)) / 
                 (1-beta_t)**0.5 +
                  beta_t*z
               )

此段代码展示了如何定义一个基本的DDPM类及其核心函数`p_sample()`用于生成样本[^2]。

#### DDIM 实现示例
下面是一个简单的DDIM版本，它利用了更少的时间步数来加速抽样流程：

def ddim_step(model_output, timestep, eta=0., shape=(1,)):
    """
    执行单次DDIM更新步骤
    
    参数:
      model_output : 来自预训练好的UNet或其他架构预测的结果
      timestep     : 当前时间戳索引
      eta          : 控制方差项的比例，默认设为零表示完全决定论模式
      shape        : 输出张量形状
      
    返回值:
      噪声减少后的图像估计
    """

    b = shape[0]

    # 获取当前时刻对应的alpha累积乘积以及其它必要变量...
    a_prev = ... # 上一时刻的alpha累计乘积
    a_next = ... # 下一刻的alpha累计乘积
    sigma_squared = ...
    
    e_tilde = torch.randn(shape, device=model_output.device) * math.sqrt(sigma_squared)
    
    x_0_pred = ...

    c1 = ((1-a_prev)/(1-a_next))*eta*sigma_squared/(a_next*(np.sqrt(a_next)-np.sqrt(a_prev)))
    c2 = (((1-a_prev)*(np.sqrt(a_next)+np.sqrt(a_prev)))/((1-a_next)*np.sqrt(a_next)))**2
    
    return x_0_pred*c2 + e_tilde*c1

上述代码片段提供了执行一次DDIM迭代所需的逻辑，其中包含了计算新状态的关键公式。