import torch                              # 导入 PyTorch 库，用于张量操作
import torch.nn.functional as F           # 导入 PyTorch 的函数式 API，包含各种神经网络函数
import numpy as np                        # 导入 NumPy 库，用于数值计算
import matplotlib as mpl                  # 导入 Matplotlib 库的主模块
import matplotlib.pyplot as plt           # 导入 Matplotlib 的绘图模块，用于显示图像
import imageio.v2 as imageio                            # 导入 imageio 库，用于读取图片文件

mpl.rcParams['figure.figsize'] = (12, 8)  # 设置 Matplotlib 图表默认尺寸为 12x8 英寸

img = torch.FloatTensor(imageio.imread('imgs/hills_2.png')/255)  # 读取图片，像素值归一化到[0,1]，转为 PyTorch 浮点张量
print(type(img))
plt.imshow(img)                           # 使用 Matplotlib 显示图片
plt.show()

import torch                              # 导入 PyTorch 库，用于张量操作
import torch.nn.functional as F           # 导入 PyTorch 的函数式 API，包含各种神经网络函数
import numpy as np                        # 导入 NumPy 库，用于数值计算
import matplotlib as mpl                  # 导入 Matplotlib 库的主模块
import matplotlib.pyplot as plt           # 导入 Matplotlib 的绘图模块，用于显示图像
import imageio.v2 as imageio                            # 导入 imageio 库，用于读取图片文件

mpl.rcParams['figure.figsize'] = (12, 8)  # 设置 Matplotlib 图表默认尺寸为 12x8 英寸

img = torch.FloatTensor(imageio.imread('imgs/hills_2.png')/255)  # 读取图片，像素值归一化到[0,1]，转为 PyTorch 浮点张量
print(type(img))
                           # 使用 Matplotlib 显示图片
def input_T(input):
    # [0,1] -> [-1,+1]
    return 2*input-1
    
def output_T(input):
    # [-1,+1] -> [0,1]
    return (input+1)/2

def show(input):
    plt.imshow(output_T(input).clip(0,1))
    
img_=input_T(img)

plt.imshow(img_)

#show(img_)

plt.show()

num_timesteps=1000
betas=torch.linspace(1e-4,2e-2,num_timesteps)

alphas=1-betas
alphas_sqrt=alphas.sqrt()
alphas_cumprod=torch.cumprod(alphas,0)
alphas_cumprod_sqrt=alphas_cumprod.sqrt()

1. 公式整体是什么意思

$q(x_t|x_{t-1})$

表示一个条件概率分布，读作：

“在给定 $x_{t-1}$ 的条件下，$x_t$ 的分布”。

这里：

• $q$：表示前向扩散过程
• $x_{t-1}$：第 $t-1$ 步的样本
• $x_t$：第 $t$ 步的样本

也就是说，这个式子定义了：

如何从上一步的样本 $x_{t-1}$ 随机生成当前步的样本 $x_t$。

2. 符号 := 的含义

$q(x_t|x_{t-1}) := \mathcal{N}(\cdots)$

这里的 $:=$ 表示：

“定义为”

也就是说，不是在推导结论，而是在规定：

前向过程的单步转移分布就定义成右边这个高斯分布。

右边的 $\mathcal{N}$ 是什么

$\mathcal{N}(x_t;\mu,\Sigma)$ 表示一个高斯分布，其中：

• 第一个位置 $x_t$：表示随机变量
• 第二个位置 $\mu$：表示均值
• 第三个位置 $\Sigma$：表示协方差矩阵

所以：

$\mathcal{N}(x_t;\sqrt{1-\beta_t}x_{t-1},\beta_t I)$

表示：

$x_t$ 服从一个均值为 $\sqrt{1-\beta_t}x_{t-1}$、协方差为 $\beta_t I$ 的高斯分布。

4. 公式中每个位置分别表示什么

公式为：

$q(x_t|x_{t-1}) := \mathcal{N}(x_t;\sqrt{1-\beta_t}x_{t-1},\beta_t I)$

可以按位置拆成三部分：

第 1 个位置：$x_t$

这是高斯分布里的随机变量。

意思是：

我们要描述的是 $x_t$ 这个量的分布。

也就是“当前步生成出来的样本是多少”。

第 2 个位置：$\sqrt{1-\beta_t}x_{t-1}$

这是高斯分布的均值。

意思是：

$x_t$ 大致会围绕这个中心值波动。

其中：

• $x_{t-1}$：上一时刻的样本
• $\sqrt{1-\beta_t}$：对上一时刻样本的缩放系数

所以均值部分表示：

当前步会保留上一步的大部分信息，但会稍微缩小一点。

如果 $\beta_t$ 很小，那么：

• $1-\beta_t$ 很接近 $1$
• $\sqrt{1-\beta_t}$ 也很接近 $1$

说明当前步和上一步会很接近，只是稍微被扰动一点。

第 3 个位置：$\beta_t I$

这是高斯分布的协方差矩阵。

其中：

• $\beta_t$：控制噪声强度
• $I$：单位矩阵

$\beta_t I$ 的意思是：

• 每个维度上的噪声方差都是 $\beta_t$
• 不同维度之间相互独立

如果是在图像里理解，就是：

每个像素或每个特征维度都被加上独立的高斯噪声。

5. 这个公式展开后是什么意思

这个分布形式可以等价地写成采样形式：

$x_t = \sqrt{1-\beta_t}x_{t-1} + \sqrt{\beta_t}\,\epsilon_t$

其中：

$\epsilon_t \sim \mathcal{N}(0,I)$

这就是这个公式最常见的“展开形式”。

从 $x_0$ 直接跳到 $x_t$ 的形式：

$q(x_t|x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t) I) \tag{2}$

定义一个函数 `forward_step()`

下面我们定义一个函数 `forward_step()`，它可以让我们使用单步前向过程 $q(x_t|x_{t-1})$；同时再定义 `forward_jump()`，用于表示从 $x_0$ 直接到 $x_t$ 的过程 $q(x_t|x_0)$。

import torch                              # 导入 PyTorch 库，用于张量操作
import torch.nn.functional as F           # 导入 PyTorch 的函数式 API，包含各种神经网络函数
import numpy as np                        # 导入 NumPy 库，用于数值计算
import matplotlib as mpl                  # 导入 Matplotlib 库的主模块
import matplotlib.pyplot as plt           # 导入 Matplotlib 的绘图模块，用于显示图像
import imageio.v2 as imageio              # 导入 imageio 库，用于读取图片文件

mpl.rcParams['figure.figsize'] = (12, 8)  # 设置 Matplotlib 图表默认尺寸为 12x8 英寸

# 读取图片，像素值归一化到[0,1]，转为 PyTorch 浮点张量
img = torch.FloatTensor(imageio.imread('imgs/hills_2.png')/255)
print(type(img))

# 值域转换函数：将图像从 [0,1] 转换到 [-1,+1]
def input_T(input):
    # [0,1] -> [-1,+1]
    return 2*input-1
    
# 值域转换函数：将图像从 [-1,+1] 转换到 [0,1]
def output_T(input):
    # [-1,+1] -> [0,1]
    return (input+1)/2

# 显示图像函数，自动转换值域并裁剪到 [0,1]
def show(input):
    plt.imshow(output_T(input).clip(0,1))

# ========== 扩散模型超参数设置 ==========
T = 100  # 时间步总数

# 定义 beta 调度（线性调度）
betas = torch.linspace(0.0001, 0.02, T)  # 从 0.0001 线性增加到 0.02

# 计算 alpha 相关参数
alphas = 1 - betas                       # alpha_t = 1 - beta_t
alphas_cumprod = torch.cumprod(alphas, dim=0)  # 累积乘积：ᾱ_t = ∏(1-β_s)
alphas_cumprod_sqrt = torch.sqrt(alphas_cumprod)  # √ᾱ_t
alphas_sqrt = torch.sqrt(alphas)          # √α_t

# ========== 加噪过程函数 ==========

def forward_step(t, condition_img, return_noise=False):
    """
        forward step: t-1 -> t
        单步加噪：从前一时刻 t-1 加噪到当前时刻 t
    """    
    assert t >= 0

    mean = alphas_sqrt[t] * condition_img    
    std = betas[t].sqrt()
      
    # sampling from N
    if not return_noise:
        return mean + std * torch.randn_like(img)
    else:
        noise = torch.randn_like(img)
        return mean + std * noise, noise
    
def forward_jump(t, condition_img, condition_idx=0, return_noise=False):
    """
        forward jump: 0 -> t
        直接跳转加噪：从初始时刻 0 直接加噪到时刻 t
    """   
    assert t >= 0
    
    mean = alphas_cumprod_sqrt[t] * condition_img
    std = (1 - alphas_cumprod[t]).sqrt()
      
    # sampling from N
    if not return_noise:
        return mean + std * torch.randn_like(img)
    else:
        noise = torch.randn_like(img)
        return mean + std * noise, noise


# ========== 测试加噪过程 ==========
if __name__ == "__main__":
    # 将图像转换到 [-1, 1] 值域
    img_ = input_T(img)
    
    # 测试 forward_jump：展示不同时间步的加噪效果
    fig, axes = plt.subplots(2, 5, figsize=(15, 6))
    
    # 显示原图
    axes[0, 0].imshow(img)
    axes[0, 0].set_title('Original (t=0)')
    axes[0, 0].axis('off')
    
    # 显示不同时间步的加噪结果
    timesteps = [10, 30, 50, 70, 90, 20, 40, 60, 80, 99]
    for idx, t in enumerate(timesteps):
        row = idx // 5
        col = idx % 5
        if row == 0 and col == 0:
            continue  # 跳过原图位置
        noisy_img = forward_jump(t, img_)
        axes[row, col].imshow(output_T(noisy_img).clip(0, 1))
        axes[row, col].set_title(f't={t}')
        axes[row, col].axis('off')
    
    plt.tight_layout()
    plt.suptitle('Forward Diffusion Process (Jump)', y=1.02)
    plt.show()
    
    # 测试 forward_step：逐步加噪过程
    fig, axes = plt.subplots(2, 5, figsize=(15, 6))
    
    current_img = img_
    for t in range(10):
        row = t // 5
        col = t % 5
        if t > 0:
            current_img = forward_step(t, current_img)
        axes[row, col].imshow(output_T(current_img).clip(0, 1))
        axes[row, col].set_title(f'Step t={t}')
        axes[row, col].axis('off')
    
    plt.tight_layout()
    plt.suptitle('Forward Diffusion Process (Step by Step)', y=1.02)
    plt.show()
    
    print("加噪过程测试完成！")

反向过程

反向过程 $p$ 的目的是：根据扩散链中的当前样本 $x_t$，去近似恢复前一步的样本 $x_{t-1}$。在实际情况下，这种近似 $p(x_{t-1}|x_t)$ 必须在不知道 $x_0$ 的前提下完成。

通常会使用一个带参数 $\theta$ 的可参数化预测模型，来估计 $p_\theta(x_{t-1}|x_t)$。

如果扩散的每一步足够小，那么反向过程也可以近似看作高斯分布：

$$p_\theta(x_{t-1}|x_t) := \mathcal{N}(x_{t-1};\mu_\theta(x_t),\Sigma_\theta(x_t))\tag{3}$$

在很多工作中，通常假设这个分布的方差不应强烈依赖于 $x_0$ 或 $x_t$，而应主要依赖于扩散过程所处的阶段 $t$。这一点可以从真实分布 $q(x_{t-1}|x_t, x_0)$ 中观察到，因为该分布的方差等于 $\tilde{\beta}_t$。

对 $\mu_\theta$ 的参数化

对于反向单步分布 $p_\theta(x_{t-1}|x_t)$ 的均值 $\mu_\theta$，至少有 3 种参数化方式：

• 直接预测：由神经网络直接估计 $\mu_\theta$
• 通过 $x_0$ 参数化：由神经网络先估计 $x_0$

$$\tilde{\mu}_\theta = \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}x_0 + \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t\tag{4}$$

• 通过从 $x_0$ 中减去噪声 $\epsilon$ 来参数化：由神经网络估计 $\epsilon$

$$x_0=\frac{1}{\sqrt{\bar{\alpha}_t}}(x_t-\sqrt{1-\bar{\alpha}_t}\epsilon)\tag{5}$$

其中，最常用的方法是近似预测高斯噪声 $\epsilon$。

下面我们来看一个 $\epsilon$ 的示例可能长什么样：

import torch                              # 导入 PyTorch 库，用于张量操作
import torch.nn.functional as F           # 导入 PyTorch 的函数式 API
import numpy as np                        # 导入 NumPy 库，用于数值计算
import matplotlib as mpl                  # 导入 Matplotlib 库的主模块
import matplotlib.pyplot as plt           # 导入 Matplotlib 的绘图模块
import imageio.v2 as imageio              # 导入 imageio 库，用于读取图片文件

mpl.rcParams['figure.figsize'] = (12, 8)  # 设置 Matplotlib 图表默认尺寸

# 读取图片，像素值归一化到[0,1]，转为 PyTorch 浮点张量
img = torch.FloatTensor(imageio.imread('imgs/hills_2.png') / 255)

# 值域转换函数：将图像从 [0,1] 转换到 [-1,+1]（扩散模型标准输入值域）
def input_T(input):
    # [0,1] -> [-1,+1]
    return 2 * input - 1

# 值域转换函数：将图像从 [-1,+1] 转换到 [0,1]（用于显示）
def output_T(input):
    # [-1,+1] -> [0,1]
    return (input + 1) / 2

# 显示图像函数，自动转换值域并裁剪到 [0,1]
def show(input):
    plt.imshow(output_T(input).clip(0, 1))


# ========== 扩散模型超参数设置 ==========
T = 100  # 时间步总数

# 定义 beta 调度（线性调度）
betas = torch.linspace(0.0001, 0.02, T)  # 从 0.0001 线性增加到 0.02

# 计算 alpha 相关参数
alphas = 1 - betas                       # alpha_t = 1 - beta_t
alphas_cumprod = torch.cumprod(alphas, dim=0)  # 累积乘积：ᾱ_t = ∏(1-β_s)
alphas_cumprod_sqrt = torch.sqrt(alphas_cumprod)  # √ᾱ_t

# 设置前向跳转的时间步
t_step = 50


def forward_jump(t, condition_img, condition_idx=0, return_noise=False):
    """
        forward jump: 0 -> t
        直接跳转加噪：从初始时刻 0 直接加噪到时刻 t
        
        根据扩散模型公式：x_t = √ᾱ_t * x_0 + √(1-ᾱ_t) * ε
        其中 ε ~ N(0, I)
    """
    assert t >= 0

    mean = alphas_cumprod_sqrt[t] * condition_img
    std = (1 - alphas_cumprod[t]).sqrt()

    # 从正态分布采样
    if not return_noise:
        return mean + std * torch.randn_like(condition_img)
    else:
        noise = torch.randn_like(condition_img)
        return mean + std * noise, noise


# ========== 主程序 ==========
if __name__ == "__main__":
    # 将图像转换到 [-1, 1] 值域
    img_ = input_T(img)

    # 执行前向跳转加噪，获取加噪后的图像和噪声
    x_t, noise = forward_jump(t_step, img_, return_noise=True)

    # 可视化结果
    plt.subplot(1, 3, 1)
    show(img_)
    plt.title(r'$x_0$')
    plt.axis('off')

    plt.subplot(1, 3, 2)
    show(x_t)
    plt.title(r'$x_t$')
    plt.axis('off')

    plt.subplot(1, 3, 3)
    show(noise)
    plt.title(r'$\epsilon$')
    plt.axis('off')

    plt.suptitle(f'Forward Diffusion Process (t={t_step})', y=0.98)
    plt.tight_layout()
    plt.show()

如果噪声 $\epsilon$ 能够被正确预测，那么我们就可以利用公式 $x_0=\frac{1}{\sqrt{\bar{\alpha}_t}}(x_t-\sqrt{1-\bar{\alpha}_t}\epsilon)$来预测 $x_0$：

import torch                              # 导入 PyTorch 库，用于张量操作
import torch.nn.functional as F           # 导入 PyTorch 的函数式 API
import numpy as np                        # 导入 NumPy 库，用于数值计算
import matplotlib as mpl                  # 导入 Matplotlib 库的主模块
import matplotlib.pyplot as plt           # 导入 Matplotlib 的绘图模块
import imageio.v2 as imageio              # 导入 imageio 库，用于读取图片文件
 
mpl.rcParams['figure.figsize'] = (12, 8)  # 设置 Matplotlib 图表默认尺寸
 
# 读取图片，像素值归一化到[0,1]，转为 PyTorch 浮点张量
img = torch.FloatTensor(imageio.imread('imgs/hills_2.png') / 255)
 
# 值域转换函数：将图像从 [0,1] 转换到 [-1,+1]（扩散模型标准输入值域）
def input_T(input):
    # [0,1] -> [-1,+1]
    return 2 * input - 1
 
# 值域转换函数：将图像从 [-1,+1] 转换到 [0,1]（用于显示）
def output_T(input):
    # [-1,+1] -> [0,1]
    return (input + 1) / 2
 
# 显示图像函数，自动转换值域并裁剪到 [0,1]
def show(input):
    plt.imshow(output_T(input).clip(0, 1))
 
 
# ========== 扩散模型超参数设置 ==========
T = 100  # 时间步总数
 
# 定义 beta 调度（线性调度）
betas = torch.linspace(0.0001, 0.02, T)  # 从 0.0001 线性增加到 0.02
 
# 计算 alpha 相关参数
alphas = 1 - betas                       # alpha_t = 1 - beta_t
alphas_cumprod = torch.cumprod(alphas, dim=0)  # 累积乘积：ᾱ_t = ∏(1-β_s)
alphas_cumprod_sqrt = torch.sqrt(alphas_cumprod)  # √ᾱ_t
 
# 设置前向跳转的时间步
t_step = 50
 
 
def forward_jump(t, condition_img, return_noise=False):
    """
        forward jump: 0 -> t
        直接跳转加噪：从初始时刻 0 直接加噪到时刻 t
 
        根据扩散模型公式：x_t = √ᾱ_t * x_0 + √(1-ᾱ_t) * ε
        其中 ε ~ N(0, I)
    """
    assert t >= 0
 
    mean = alphas_cumprod_sqrt[t] * condition_img
    std = (1 - alphas_cumprod[t]).sqrt()
 
    # 从正态分布采样
    if not return_noise:
        return mean + std * torch.randn_like(condition_img)
    else:
        noise = torch.randn_like(condition_img)
        return mean + std * noise, noise
 
 
# ========== 主程序 ==========
if __name__ == "__main__":
    # 将图像转换到 [-1, 1] 值域
    img_ = input_T(img)
 
    # 执行前向跳转加噪，获取加噪后的图像和噪声
    x_t, noise = forward_jump(t_step, img_, return_noise=True)
 
    # 根据 x_t 和噪声预测 x_0
    # 由 x_t = √ᾱ_t * x_0 + √(1-ᾱ_t) * ε 可得：
    # x_0 = (x_t - √(1-ᾱ_t) * ε) / √ᾱ_t
    x_0_pred = (x_t - (1 - alphas_cumprod[t_step]).sqrt() * noise) / alphas_cumprod_sqrt[t_step]
 
    # 可视化结果：x_t（加噪图像）、x_0 预测值、原始 x_0
    plt.subplot(1, 3, 1)
    show(x_t)
    plt.title(r'$x_t$ ($\ell_1$: {:.3f})'.format(F.l1_loss(x_t, img_)))
    plt.axis('off')
 
    plt.subplot(1, 3, 2)
    show(x_0_pred)
    plt.title(r'$x_0$ prediction ($\ell_1$: {:.3f})'.format(F.l1_loss(x_0_pred, img_)))
    plt.axis('off')
 
    plt.subplot(1, 3, 3)
    show(img_)
    plt.title(r'$x_0$')
    plt.axis('off')
 
    plt.suptitle(f'x_0 Prediction from x_t (t={t_step})', y=0.98)
    plt.tight_layout()
    plt.show()

这里的含义是：

• 第一幅图显示当前带噪样本 $x_t$
• 第二幅图显示根据预测噪声反推出的 $x_0$，记作 $x_0\_pred$
• 第三幅图显示真实的干净样本 $x_0$

如果预测的噪声足够准确，那么 $x_0\_pred$ 就会接近真实的 $x_0$。

近似或已知 $x_0$ 后，可以进一步近似 $t-1$ 步的均值

一旦我们得到了对 $x_0$ 的近似（或者直接知道真实的 $x_0$），就可以利用公式 (4) 来近似第 $t-1$ 步对应分布的均值：

# estimate mean
mean_pred = x_0_pred * (alphas_cumprod_sqrt[t_step-1] * betas[t_step]) / (1 - alphas_cumprod[t_step]) \
          + x_t * (alphas_sqrt[t_step] * (1 - alphas_cumprod[t_step-1])) / (1 - alphas_cumprod[t_step])
 
# let's compare it to ground truth mean of the previous step (requires knowledge of x_0)
mean_gt = alphas_cumprod_sqrt[t_step-1] * img_

其中：

• `mean_pred`：根据预测得到的 $x_0$ 和当前样本 $x_t$ 计算出的反向过程均值近似
• `mean_gt`：真实的前一步均值，它需要知道真实的 $x_0$

这里的 `mean_gt` 本质上对应的是前向过程中从干净图像得到的“真实均值”。

---

为什么 `mean_pred` 的误差通常会更大

由于公式 (4) 中的反向过程均值估计 $\tilde{\mu}_\theta$，本质上是对带噪的 $x_t$ 和估计出来的 $x_0$ 做一种线性组合，因此它通常会比前向过程直接得到的均值误差更大。

原因是：

• $x_t$ 本身仍然含有噪声
• $x_0\_pred$ 只是估计值，不是真实值
• 因此两者线性组合得到的均值仍然会带来额外误差

而前向过程中的均值通常只是对干净样本乘上一个标量，因此更精确。

下面的代码展示了这种比较：

plt.subplot(1,3,1)
show(x_t)
plt.title('$x_t$   ($\ell_1$: {:.3f})'.format(F.l1_loss(x_t, img_)))
 
plt.subplot(1,3,2)
show(mean_pred)
plt.title(r'$\tilde{\mu}_{t-1}$' + '  ($\ell_1$: {:.3f})'.format(F.l1_loss(mean_pred, img_)))
 
plt.subplot(1,3,3)
show(mean_gt)
plt.title(r'$\mu_{t-1}$' + '  ($\ell_1$: {:.3f})'.format(F.l1_loss(mean_gt, img_)))

这三幅图分别表示：

• 左图：当前带噪图像 $x_t$
• 中图：反向过程估计的均值 $\tilde{\mu}_{t-1}$
• 右图：真实均值 $\mu_{t-1}$

比较它们的 $\ell_1$ 误差，就可以看出估计值和真实值之间的差别。

---

得到 `mean_pred` 之后，就可以定义前一步的分布

当我们获得了 `mean_pred`，也就是 $\tilde{\mu}_t$ 之后，就可以将前一步的条件分布定义为：

$$\tilde{\beta}_t = \beta_t \tag{6}$$

$$p_\theta(x_{t-1}|x_t) := \mathcal{N}(x_{t-1}; \tilde{\mu}_\theta(x_t,t), \tilde{\beta}_t I) \tag{7}$$

这个定义表示：

• 均值使用网络估计出来的 $\tilde{\mu}_\theta(x_t,t)$
• 方差使用 $\tilde{\beta}_t$
• 从而构造出反向过程中的高斯分布

重要说明

下面的实验应被看作一种模拟。
在实际训练和推理中，网络必须预测以下三者之一：
* $\epsilon$
* $x_0$
* $\tilde{\mu}_\theta$

而这里的 $\epsilon$ 值只是直接代入使用的，并不是真实模型预测出来的结果。

`reverse_step` 函数

下面这个函数实现了一次反向扩散步骤：

def reverse_step(epsilon, x_t, t_step, return_noise=False):
 
    # estimate x_0 based on epsilon
    x_0_pred = (x_t - (1 - alphas_cumprod[t_step]).sqrt() * epsilon) / (alphas_cumprod_sqrt[t_step])
 
    if t_step == 0:
        sample = x_0_pred
        noise = torch.zeros_like(x_0_pred)
    else:
        # estimate mean
        mean_pred = x_0_pred * (alphas_cumprod_sqrt[t_step-1] * betas[t_step]) / (1 - alphas_cumprod[t_step]) \
                  + x_t * (alphas_sqrt[t_step] * (1 - alphas_cumprod[t_step-1])) / (1 - alphas_cumprod[t_step])
 
        # compute variance
        beta_pred = betas[t_step].sqrt() if t_step != 0 else 0
 
        sample = mean_pred + beta_pred * torch.randn_like(x_t)
 
        # this noise is only computed for simulation purposes (since x_0_pred is not known normally)
        noise = (sample - x_0_pred * alphas_cumprod_sqrt[t_step-1]) / (1 - alphas_cumprod[t_step-1]).sqrt()
 
    if return_noise:
        return sample, noise
    else:
        return sample

---

`reverse_step()` 的逻辑讲解

1. 根据 $\epsilon$ 先恢复 $x_0$

x_0_pred = (x_t - (1 - alphas_cumprod[t_step]).sqrt() * epsilon) / (alphas_cumprod_sqrt[t_step])

这一行对应公式 (5)：

$$x_0=\frac{1}{\sqrt{\bar{\alpha}_t}}(x_t-\sqrt{1-\bar{\alpha}_t}\epsilon)$$

作用是：

• 输入当前带噪样本 $x_t$
• 输入预测得到的噪声 $\epsilon$
• 反推出干净图像的估计值 $x_0$

---

2. 如果已经到达 $t=0$

if t_step == 0:
    sample = x_0_pred
    noise = torch.zeros_like(x_0_pred)

当 `t_step == 0` 时，说明已经到达反向过程的最后一步，不需要再继续去噪。

因此：

• 直接把 $x_0\_pred$ 当作最终结果
• 此时噪声设为全零

---

3. 如果 $t > 0$，先估计反向均值

mean_pred = x_0_pred * (alphas_cumprod_sqrt[t_step-1] * betas[t_step]) / (1 - alphas_cumprod[t_step]) \
          + x_t * (alphas_sqrt[t_step] * (1 - alphas_cumprod[t_step-1])) / (1 - alphas_cumprod[t_step])

这一行对应公式 (4)：

$$\tilde{\mu}_\theta = \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}x_0 + \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t$$

只不过代码里用的是预测值 `x_0_pred` 来代替真实的 $x_0$。

意思是：

• 综合当前带噪样本 $x_t$
• 再结合恢复出的干净样本估计 $x_0\_pred$
• 得到前一步分布的均值估计 `mean_pred`

---

4. 计算标准差

beta_pred = betas[t_step].sqrt() if t_step != 0 else 0

这里使用：

• 方差 $\tilde{\beta}_t = \beta_t$
• 所以标准差为 $\sqrt{\beta_t}$

这是为了后面从高斯分布中采样：

$$x_{t-1} \sim \mathcal{N}(\tilde{\mu}_\theta, \tilde{\beta}_t I)$$

---

5. 从反向分布中采样

sample = mean_pred + beta_pred * torch.randn_like(x_t)

这一步就是从高斯分布中真正采样出 $x_{t-1}$。

也就是：

• 均值为 `mean_pred`
• 标准差为 `beta_pred`
• 加上一份标准高斯噪声

对应采样形式可以写成：

$$x_{t-1} = \tilde{\mu}_\theta + \sqrt{\tilde{\beta}_t}\,z,\quad z\sim\mathcal{N}(0,I)$$

---

6. 计算 `noise` 仅用于模拟分析

noise = (sample - x_0_pred * alphas_cumprod_sqrt[t_step-1]) / (1 - alphas_cumprod[t_step-1]).sqrt()

这里的 `noise` 只是为了实验分析而额外算出来的。

因为在真实的反向扩散推理中：

• 我们并不知道真实的 $x_0$
• 这里只是因为做实验，才利用 `x_0_pred` 反推一个噪声值

所以注释里才特别说明：

• 这个噪声只是为了模拟目的而计算
• 并不是实际系统中能直接获得的量

7. 是否返回噪声

if return_noise:
    return sample, noise
else:
    return sample

如果 `return_noise=True`：

• 返回当前反向步骤采样结果 `sample`
• 同时返回额外计算的 `noise`

否则：

• 只返回 `sample`

这个 `reverse_step()` 函数做的事情是：

• 先根据预测噪声 $\epsilon$ 还原出 $x_0$
• 再根据 $x_t$ 和 $x_0$ 的估计值计算反向分布均值
• 然后从这个高斯分布中采样得到 $x_{t-1}$

也就是说，它实现了扩散模型中的一次“去噪反推”。