【Hugging Face 开源库】Diffusers 库 —— 扩散模型

Diffusers 是 Hugging Face 开源的 Python 库，专门用于加载、训练和推理扩散模型（Diffusion Models）。

扩散模型是一类生成式模型，它们通过添加和去除噪声来生成高质量图像、音频和视频。

《从零开始学扩散模型》

Diffusers 的三个主要组件

1. DiffusionPipeline：端到端推理工具

DiffusionPipeline 是 Diffusers 库的核心组件之一，它提供了一个高层 API，帮助用户快速从预训练的扩散模型中生成样本，而无需深入了解底层实现。

示例：使用 Stable Diffusion 生成图像

from diffusers import StableDiffusionPipeline
import torch# 加载预训练的 Stable Diffusion 模型
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipeline.to("cuda")  # 使用 GPU 加速# 生成图像
prompt = "a futuristic city at sunset, high detail, digital painting"
image = pipeline(prompt).images[0]# 显示图像
image.show()

• 通过 from_pretrained() 加载 Hugging Face Hub 上的 Stable Diffusion 预训练模型。

  unwayml/stable-diffusion-v1-5 是 Stable Diffusion v1.5 预训练模型的 权重（weights），它被托管在 Hugging Face Hub 上，供用户下载并进行推理或微调。
  在 Diffusers 库中，from_pretrained("runwayml/stable-diffusion-v1-5") 其实是加载该模型的预训练参数，包括：

  • UNet（去噪网络）
  • VAE（变分自编码器，用于图像编码和解码）
  • Text Encoder（如 CLIP，用于处理文本输入）
  • 调度器（Scheduler，用于指导去噪过程）

  这些组件的权重都是从 runwayml/stable-diffusion-v1-5 仓库中下载的。

• 只需输入 prompt（文本描述），就能生成相应的图像。

__call__ 函数

在 Python 中，__call__ 是一个特殊的方法，它 允许一个对象像函数一样被调用。当你调用一个对象时，Python 实际上是调用了这个对象的 __call__ 方法。

在 diffusers 库中，所有的管道对象（如 StableDiffusionPipeline）都实现了一个 __call__ 方法，用于处理图像生成任务，所以说 管道（pipeline）对象可以像函数一样被调用。

让我们实现一个 简单的管道对象（Pipeline），用来模拟 Diffusers 的 __call__ 方法是如何工作的。这个管道将接受一个文本 prompt，然后通过一个简单的 UNet 模型 生成一个伪图像（这里只是模拟，不是实际的图像生成）。

示例：实现一个简单的 DiffusionPipeline

import torch
import torch.nn as nnclass SimpleUNet(nn.Module):""" 一个简单的 UNet 模型模拟去噪过程 """def __init__(self):super().__init__()self.fc = nn.Linear(100, 100)  # 简化的全连接层def forward(self, x):return self.fc(x)  # 这里只是简单的线性变换class SimplePipeline:""" 一个简单的管道对象，模拟 DiffusionPipeline 的 __call__ 方法 """def __init__(self):self.unet = SimpleUNet()  # 预训练的去噪模型self.device = "cuda" if torch.cuda.is_available() else "cpu"self.unet.to(self.device)def __call__(self, prompt: str):""" 模拟调用管道进行图像生成 """print(f"Processing prompt: {prompt}")# 1. 生成随机噪声作为输入noise = torch.randn(1, 100).to(self.device)# 2. 通过 UNet 进行处理output = self.unet(noise)# 3. 模拟图像输出return output.detach().cpu().numpy()# 使用管道
pipeline = SimplePipeline()
generated_image = pipeline("A beautiful sunset over the ocean")  # 通过 __call__ 触发
print("Generated image shape:", generated_image.shape)

1. SimpleUNet:

   • 这里用一个简单的 全连接层 代替真正的 UNet（通常是 CNN）。
   • 这个网络用于处理随机噪声，模拟去噪过程。

2. SimplePipeline:

   • __init__ 方法：创建一个 UNet 模型并加载到 GPU（如果可用）。
   • __call__ 方法：
     • ① 接收文本提示 prompt（但这里的代码没有真正解析文本，仅模拟处理）。
     • ② 生成随机噪声，作为输入。
     • ③ 通过 UNet 处理，得到输出。
     • ④ 返回最终“生成的图像”（其实只是一个数值数组）。

3. 如何使用 __call__ 方法:

   • pipeline("A beautiful sunset over the ocean") 直接调用 实例，会自动触发 __call__ 方法。
   • 这样 对象本身就像一个函数一样可以调用，符合 Diffusers 设计风格。

可以在 __call__ 方法中 添加真正的 VAE、文本编码器、调度器 来让它更接近 Diffusers 的 DiffusionPipeline。

这样，pipeline("prompt") 的行为就类似于 StableDiffusionPipeline(prompt) 了！ 🚀

在实际的 diffusers 库中，管道对象的 __call__ 方法会处理各种输入嵌入、噪声调度器、生成模型等，最终生成高质量的图像。例如，在 StableDiffusionPipeline 中，__call__ 方法会接受提示、图像嵌入等，并通过扩散模型逐步生成图像。

callback_on_step_end 管道回调函数

callback_on_step_end 允许我们在 扩散管道的每一步去噪迭代结束时 执行 自定义回调函数。

这样，可以 动态修改管道的属性或调整张量，而 无需修改 Diffusers 库的底层代码。

举个栗子，使用回调函数 在去噪的不同阶段动态调整 guidance_scale（引导比例），让模型在去噪的前几步加强条件引导（更遵循 prompt），后几步减少 guidance_scale 以生成更自然的图像。

import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler# 加载 Stable Diffusion 管道
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config) # 切换 DDIMScheduler 作为调度器
pipeline.to("cuda")# 定义回调函数
def dynamic_guidance_callback(pipe, i, latents):"""在去噪过程的每一步，动态修改 guidance_scale:param pipe: 当前管道对象:param i: 当前去噪步数:param latents: 当前的潜变量"""total_steps = pipe.scheduler.config.num_train_timestepsif i < total_steps * 0.3:  # 在前 30% 的步数里，增加 guidance_scalepipe.guidance_scale = 10.0  elif i < total_steps * 0.6:  # 在 30% - 60% 的步数里，降低 guidance_scalepipe.guidance_scale = 7.5  else:  # 在最后 40% 的步数里，进一步减少pipe.guidance_scale = 5.0  print(f"Step {i}: guidance_scale set to {pipe.guidance_scale}")# 生成图像
prompt = "A futuristic city with neon lights at night"# 在 pipeline() 调用时传递 callback_on_step_end
image = pipeline(prompt, callback_on_step_end=dynamic_guidance_callback).images[0]# 显示图像
image.show()

这个回调函数在 每次去噪步骤结束后执行，并动态调整 guidance_scale：

• 前 30% 的步数：使用 更高的 guidance_scale = 10.0，让生成的图像更符合 prompt 描述。

• 30% - 60% 步数：降低 guidance_scale 到 7.5，让图像稍微放松对 prompt 的严格约束。

• 最后 40% 步数：进一步降低到 5.0，让图像更自然，减少过度引导导致的“过拟合”问题。

Pipeline callbacks
除了动态调整 guidance_scale，还可以用 callback_on_step_end 进行：

• 添加自定义去噪步骤（比如在中间步骤插入额外的图像操作）
• 修改 latents 变量（例如，在某些步数中加入额外的噪声或调整颜色分布）
• 记录或可视化去噪过程（比如，每隔 10 步保存当前的潜变量图像，观察去噪演化）

2. 预训练模型架构和模块

Diffusers 提供了许多 预训练的模型组件，可以用来构建新的扩散系统，例如：

• UNet（去噪神经网络）
• VAE（Variational Autoencoder）（用于图像编码和解码）
• Text Encoder（例如 CLIP，用于理解文本提示）

示例：使用 UNet 作为去噪模型

from diffusers import UNet2DModel# 定义一个 UNet 模型
unet = UNet2DModel(sample_size=64,  # 图像大小in_channels=3,    # RGB 颜色通道out_channels=3,layers_per_block=2,block_out_channels=(64, 128, 256),
)# 查看模型参数
print(unet)

• UNet2DModel 是扩散模型的核心组件之一，负责在训练和推理过程中去噪。
• 这里的 UNet 结构可以自定义，如通道数、块的层数等。

UNet

U-Net: Convolutional Networks for Biomedical Image Segmentation

Unet 最初设计用于生物医学图像分割。

UNet 是一种 卷积神经网络 架构，结构类似于一个对称的 U 字形，由 编码器（下采样）和解码器（上采样） 组成。

• 编码器逐步提取图像特征并缩小空间维度，
• 解码器则将这些特征还原到原始的空间维度，同时逐步增加分辨率。
  

UNet 的关键特性：

• 对称结构：编码器和解码器对称分布。
• 跳跃连接：直接将编码器的中间层输出传递到解码器的对应层，保留了高分辨率特征。
• 多尺度特征提取：在不同尺度上提取特征，提升了网络对细节的捕捉能力。

VAE（Variational AutoEncoder）

VAE（Variational AutoEncoder） 变分自编码器是一种生成模型，通过学习输入数据的潜在表示来生成新数据。

VAE 由编码器和解码器组成：

• 编码器：将 输入图像 转换为 潜在空间的分布（均值和方差）。
• 解码器：从潜在空间的采样生成 新图像。
  
  

VAE 的关键特性：

• 概率模型：VAE 学习输入数据的概率分布，从而生成多样化的样本。
• 连续潜在空间：潜在空间中的小变化会导致生成图像的小变化，具有很好的连续性。

图像尺寸与 UNet 和 VAE 的关系

在图像生成任务中，输入图像的尺寸需要匹配 UNet 和 VAE 的预期输入输出尺寸。

在 diffusers 库的 MimicBrushPipeline（或类似的图像生成管道）中，默认的输入图像尺寸是通过以下代码计算的：

height = height or self.unet.config.sample_size * self.vae_scale_factor
width = width or self.unet.config.sample_size * self.vae_scale_factor

Stable Diffusion 生成图像时，涉及 VAE（变分自编码器） 和 UNet（去噪网络）：

• VAE 作用：将高清图像 压缩 成一个 低维潜空间（latent space），然后再 解码 回原始尺寸。

• UNet 作用：在潜空间中 去噪，逐步优化潜变量，使其接近真实图像的潜变量。

关键点：VAE 会对图像进行 vae_scale_factor 倍缩放。举个栗子吧，

• 输入 VAE 的图像： 512×512

• 经过 VAE 编码后： 512/8 = 64×64（缩小 8 倍）

• UNet 处理的就是 64 × 64 的潜变量。

所以：

• height=64×8=512
• width=64×8=512

这确保了：

• UNet 处理 64 × 64 潜变量时尺寸正确。

• VAE 进行解码时，最终输出的是 512 × 512 的图像。

EMA（Exponential Moving Average）

EMA（指数移动平均）是一种 平滑技术，在深度学习中，常用于 存储模型可学习参数的局部平均值。

可以把它看作一个“影子模型”，它的参数不是简单地复制原模型，而是随着训练 以指数衰减的方式 逐步向原模型靠拢。

为什么要使用 EMA？

• 提高模型稳定性：在训练过程中，模型参数可能会剧烈波动，EMA 平均化了参数，使其更稳定。
• 提升泛化能力：直接使用 EMA 计算的参数进行推理，通常比原始参数表现更好，尤其是在 少量训练步数 下。
• 适用于生成模型（如 Diffusion Models）：Diffusers 库中的 Stable Diffusion 训练时 使用 EMA 来平滑 UNet 权重，使生成的图像更加稳定。
• 在半监督学习中常用：如 Mean Teacher 方法，使用 EMA 计算的模型作为“教师”模型指导学生模型学习。

EMA 在累积历史信息的同时，更关注最近的更新，从而对新数据变化更敏感，而不会受太早的参数扰动。

假设：

• ${\theta }_t$ 是第 $t$ 轮训练的模型参数
• ${\theta }_{\text{EMA},t}$ 是第 $t$ 轮的 EMA 计算的影子参数
• $\alpha$ 是 EMA 衰减系数（通常取 0.99 ~ 0.999）

EMA 参数的更新方式：
$${\theta }_{\text{EMA},t}=\alpha \cdot {\theta }_{\text{EMA},t-1}+(1-\alpha )\cdot {\theta }_t$$
这意味着：

• 较早的参数影响力逐渐减弱（因为乘以了 $\alpha$）。
• 最近的参数更新权重更大（乘以 $1-\alpha$）。
• 选择 较大的 $\alpha$（如 0.999），EMA 更新较慢，适用于平滑长时间的变化。

为什么较早的参数影响力逐渐减弱？

我们可以将 EMA 当前参数展开，看看它是如何由历史所有参数的加权平均组成的：

$${\theta }_{\text{EMA},t}=(1-\alpha )\cdot {\theta }_t+\alpha (1-\alpha )\cdot {\theta }_{t-1}+{\alpha }^2(1-\alpha )\cdot {\theta }_{t-2}+{\alpha }^3(1-\alpha )\cdot {\theta }_{t-3}+\dots$$

这说明：

• 最近的参数 ${\theta }_t$ 乘以 $1-\alpha$（即 0.01），虽然数值小，但它是最新的更新，影响直接而强烈。
• 较早的参数 ${\theta }_{t-1},{\theta }_{t-2}$ 乘以 $\alpha ,{\alpha }^2$ 等次幂，影响力随着时间推移呈指数级衰减。
• 老的参数贡献依然存在，但比重越来越小，这使得 EMA 更关注近期变化，而不会被早期的不稳定训练步骤影响太多。

💡直觉理解 EMA 的本质是一种带有“记忆衰减”的平滑机制：

• 老的参数不会立刻丢失，但它的影响会随着时间逐步减弱，让新数据有更大的话语权。
• 虽然最近参数的权重（1 - α = 0.01）看似小，但它不会被 EMA 继续削弱，因此它的相对影响力更大。
• 较早的参数影响力会随着 ${\alpha }^t$ 指数级减少，长期来看其贡献会趋近于 0。

如果 $\alpha =0.99$，那么过去 5 个时间步的参数贡献依次为：
$$\begin{array}{rl}\text{Step }t:& (1-\alpha )=0.01\\ \text{Step }t-1:& 0.99\times 0.01=0.0099\\ \text{Step }t-2:& 0.99^2\times 0.01=0.009801\\ \text{Step }t-3:& 0.99^3\times 0.01=0.00970299\\ \text{Step }t-4:& 0.99^4\times 0.01=0.0096059601\end{array}$$

下面是一个简单的 PyTorch EMA 代码示例，展示如何在训练过程中维护一个 EMA 版本的模型参数。

import torch
import torch.nn as nnclass EMA:"""指数移动平均（EMA），用于平滑模型参数"""def __init__(self, model, decay=0.999):self.model = modelself.decay = decay  # EMA 影子参数衰减系数self.shadow = {name: param.clone().detach() for name, param in model.named_parameters()}def update(self):"""更新 EMA 影子模型参数"""for name, param in self.model.named_parameters():if param.requires_grad:self.shadow[name] = self.decay * self.shadow[name] + (1 - self.decay) * param.detach()def apply_shadow(self):"""使用 EMA 参数更新原模型（推理时调用）"""for name, param in self.model.named_parameters():if param.requires_grad:param.data.copy_(self.shadow[name])# 创建简单的神经网络
class SimpleModel(nn.Module):def __init__(self):super().__init__()self.fc = nn.Linear(10, 1)def forward(self, x):return self.fc(x)# 初始化模型和 EMA 影子模型
model = SimpleModel()
ema = EMA(model, decay=0.99)# 模拟训练过程
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for step in range(100):# 训练步骤（假设 x 是输入数据）x = torch.randn(16, 10)loss = model(x).mean()optimizer.zero_grad()loss.backward()optimizer.step()# 更新 EMA 影子模型ema.update()if step % 10 == 0:print(f"Step {step}: loss={loss.item():.4f}")# 在推理时应用 EMA 参数
ema.apply_shadow()

1. EMA 类

   • 维护了 shadow（影子模型参数）。
   • 通过 update() 逐步更新 EMA 版本的参数。
   • apply_shadow() 用于推理时将 EMA 参数应用到原模型上。

2. 训练过程中

   • 每次模型参数更新后，调用 ema.update()，让影子模型参数缓慢跟随原模型更新。

3. 推理时

   • ema.apply_shadow() 把 EMA 版本的参数复制到模型，通常能获得 更好的性能。

在 diffusers 库中，EMA 主要用于 训练 UNet（去噪网络）：

• 训练过程中，EMA 版本的 UNet 逐步更新。
• 在推理时，使用 EMA 版本的 UNet 进行采样，以 提高图像质量。

Diffusers 使用 EMAModel 进行 EMA 计算：

from diffusers.models import EMAModel
# 初始化 EMA 模型
ema_unet = EMAModel(pipeline.unet.parameters(), decay=0.999)
# 在训练后更新 EMA 影子模型
ema_unet.step(pipeline.unet.parameters())
# 复制 EMA 参数到 UNet（推理时）
ema_unet.copy_to(pipeline.unet.parameters())

3. 调度器（Schedulers）

Scheduler，中文译为“调度器”，在扩散模型中负责控制噪声的添加和去除过程。

它定义了 在每个扩散步骤中，向数据添加多少噪声，以及在去噪过程中如何逐步恢复原始数据。

Diffusers 库提供了多种调度器，例如：

• DDIMScheduler（去噪扩散隐变量模型）
• PNDMScheduler（更快的推理）
• DPMSolverMultistepScheduler（更稳定的采样）

示例：使用不同调度器进行推理

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler# 加载 Stable Diffusion 并更换调度器
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)# 生成图像
prompt = "a magical forest with glowing trees"
image = pipeline(prompt).images[0]
image.show()

• pipeline.scheduler = DPMSolverMultistepScheduler.from_config(...) 切换不同的去噪调度器。
• 不同的调度器会影响生成速度和图像质量，比如 DPMSolver 可以加快采样，同时保持高质量输出。