01 - DDIM加速采样¶

# 原始1000步
timesteps = list(range(1000))

# DDIM 50步（均匀子采样）
timesteps = list(range(0, 1000, 20))  # [0, 20, 40, ..., 980]

# DDIM 10步（更稀疏）
timesteps = list(range(0, 1000, 100))  # [0, 100, 200, ..., 900]

import torch
import torch.nn as nn
import numpy as np
from typing import Optional, List

class DDIMSampler:
    """
    DDIM采样器
    """
    def __init__(
        self,
        diffusion,  # GaussianDiffusion对象
        ddim_steps: int = 50,
        ddim_eta: float = 0.0,  # 0=确定性, 1=随机性
    ):
        self.diffusion = diffusion
        self.ddim_steps = ddim_steps
        self.ddim_eta = ddim_eta

        # 选择时间步子集
        self.timestep_map = self._get_timestep_map()

    def _get_timestep_map(self) -> List[int]:
        """
        将DDPM的1000步映射到DDIM的ddim_steps步
        """
        c = self.diffusion.timesteps // self.ddim_steps
        timesteps = list(range(0, self.diffusion.timesteps, c))
        return timesteps[:self.ddim_steps]

    def ddim_sample(
        self,
        model: nn.Module,
        x_t: torch.Tensor,
        t: int,
        t_prev: int,
        clip_denoised: bool = True,
    ) -> torch.Tensor:
        """
        单步DDIM采样

        参数:
            model: 噪声预测模型
            x_t: 当前状态
            t: 当前时间步
            t_prev: 上一个时间步
            clip_denoised: 是否裁剪预测值

        返回:
            x_{t_prev}: 上一个时间步的状态
        """
        # 预测噪声
        t_tensor = torch.full((x_t.shape[0],), t, device=x_t.device, dtype=torch.long)
        eps = model(x_t, t_tensor)

        # 计算alpha值
        alpha_t = self.diffusion.alphas_cumprod[t]
        alpha_t_prev = self.diffusion.alphas_cumprod[t_prev] if t_prev >= 0 else torch.tensor(1.0)

        # 预测x_0
        x_0_pred = (x_t - torch.sqrt(1 - alpha_t) * eps) / torch.sqrt(alpha_t)

        if clip_denoised:
            x_0_pred = torch.clamp(x_0_pred, -1.0, 1.0)

        # 计算方差
        sigma_t = self.ddim_eta * torch.sqrt(
            (1 - alpha_t_prev) / (1 - alpha_t) * (1 - alpha_t_prev / alpha_t)
        )

        # 计算噪声系数
        noise_coeff = torch.sqrt(1 - alpha_t_prev - sigma_t**2)

        # DDIM采样公式
        x_prev = (
            torch.sqrt(alpha_t_prev) * x_0_pred +
            noise_coeff * eps
        )

        # 添加随机噪声（如果eta > 0）
        if self.ddim_eta > 0:
            noise = torch.randn_like(x_t)
            x_prev = x_prev + sigma_t * noise

        return x_prev

    def sample(
        self,
        model: nn.Module,
        shape: tuple,
        device: torch.device,
        noise: Optional[torch.Tensor] = None,
        progress: bool = True,
    ) -> torch.Tensor:
        """
        完整的DDIM采样

        参数:
            model: 噪声预测模型
            shape: 输出形状 (B, C, H, W)
            device: 计算设备
            noise: 可选的初始噪声
            progress: 是否显示进度条

        返回:
            生成的样本
        """
        model.eval()  # eval()评估模式

        # 初始化
        if noise is not None:
            x = noise
        else:
            x = torch.randn(shape, device=device)

        # 时间步列表（反向）
        timesteps = self.timestep_map[::-1]

        if progress:
            from tqdm import tqdm
            timesteps = tqdm(timesteps, desc="DDIM Sampling")

        # 采样循环
        for i, t in enumerate(timesteps):  # enumerate同时获取索引和元素
            t_prev = timesteps[i + 1] if i + 1 < len(timesteps) else -1
            x = self.ddim_sample(model, x, t, t_prev)

        return x

# 使用示例
if __name__ == "__main__":
    from diffusion import GaussianDiffusion

    # 创建扩散过程
    diffusion = GaussianDiffusion(timesteps=1000)

    # 创建DDIM采样器
    ddim_sampler = DDIMSampler(
        diffusion=diffusion,
        ddim_steps=50,  # 只用50步
        ddim_eta=0.0,   # 确定性采样
    )

    print(f"DDPM时间步: {diffusion.timesteps}")
    print(f"DDIM时间步: {ddim_sampler.ddim_steps}")
    print(f"加速比: {diffusion.timesteps / ddim_sampler.ddim_steps:.1f}x")

def uniform_timestep_schedule(total_steps, num_ddim_steps):
    """均匀选择时间步"""
    c = total_steps // num_ddim_steps
    return list(range(0, total_steps, c))[:num_ddim_steps]

# 示例
print(uniform_timestep_schedule(1000, 50))
# [0, 20, 40, 60, ..., 980]

def quadratic_timestep_schedule(total_steps, num_ddim_steps):
    """二次采样，早期更密集"""
    timesteps = []
    for i in range(num_ddim_steps):
        t = int((i / num_ddim_steps) ** 2 * total_steps)
        timesteps.append(t)
    return timesteps

print(quadratic_timestep_schedule(1000, 50))
# 早期步数更密集

def benchmark_sampling(model, diffusion, shape, device):
    """对比DDPM和DDIM的采样速度"""
    import time

    # DDPM采样
    start = time.time()
    x_ddpm = diffusion.p_sample_loop(model, shape, device, progress=False)
    ddpm_time = time.time() - start

    # DDIM采样（50步）
    ddim_sampler = DDIMSampler(diffusion, ddim_steps=50)
    start = time.time()
    x_ddim = ddim_sampler.sample(model, shape, device, progress=False)
    ddim_time = time.time() - start

    print(f"DDPM (1000步): {ddpm_time:.2f}s")
    print(f"DDIM (50步): {ddim_time:.2f}s")
    print(f"加速比: {ddpm_time / ddim_time:.1f}x")

    return ddpm_time, ddim_time

# 核心DDIM采样步骤
x_0_pred = (x_t - sqrt(1-alpha_t) * eps) / sqrt(alpha_t)
x_{t-1} = sqrt(alpha_{t-1}) * x_0_pred + sqrt(1-alpha_{t-1}) * eps