【Hugging Face 开源库】Diffusers 库 ——扩散模型
- Diffusers 的三个主要组件
- 1. DiffusionPipeline:端到端推理工具
- `__call__` 函数
- `callback_on_step_end` 管道回调函数
- 2. 预训练模型架构和模块
- UNet
- VAE(Variational AutoEncoder)
- 图像尺寸与 UNet 和 VAE 的关系
- EMA(Exponential Moving Average)
- 3. 调度器(Schedulers)
《从零开始学扩散模型》
Diffusers 是 Hugging Face 开源的 Python 库,专门用于加载、训练和推理扩散模型(Diffusion Models)。
扩散模型是一类生成式模型,它们通过添加和去除噪声来生成高质量图像、音频和视频。
深入浅出去噪扩散概率模型 DDPM
- 从简单空间开始的好处是 我们知道如何从这个简单分布中获取/采样一个点(可以将其视为数据子空间外的任何点)。
- 我们的目标是 弄清楚如何返回到数据子空间。
- 问题在于,从这个“简单”空间中的起点可以走无限多条路径,但 只有其中的一小部分能带我们到达“数据”子空间。
NOTE:
- 扩散模型不是一个神经网络,而是一套定义了加噪、去噪公式的模型。
- 扩散模型中需要一个去噪模型来去噪,去噪模型一般是个神经网络。
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)
-
SimpleUNet
:- 这里用一个简单的 全连接层 代替真正的 UNet(通常是 CNN)。
- 这个网络用于处理随机噪声,模拟去噪过程。
-
SimplePipeline
:__init__
方法:创建一个 UNet 模型并加载到 GPU(如果可用)。__call__
方法:- ① 接收文本提示
prompt
(但这里的代码没有真正解析文本,仅模拟处理)。 - ② 生成随机噪声,作为输入。
- ③ 通过 UNet 处理,得到输出。
- ④ 返回最终“生成的图像”(其实只是一个数值数组)。
- ① 接收文本提示
-
如何使用
__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 在累积历史信息的同时,更关注最近的更新,从而对新数据变化更敏感,而不会受太早的参数扰动。
假设:
- θ t \theta_t θt 是第 t t t 轮训练的模型参数
- θ EMA , t \theta_{\text{EMA},t} θEMA,t 是第 t t t 轮的 EMA 计算的影子参数
- α \alpha α 是 EMA 衰减系数(通常取 0.99 ~ 0.999)
EMA 参数的更新方式:
θ EMA , t = α ⋅ θ EMA , t − 1 + ( 1 − α ) ⋅ θ t \theta_{\text{EMA},t} = \alpha \cdot \theta_{\text{EMA},t-1} + (1 - \alpha) \cdot \theta_t θEMA,t=α⋅θEMA,t−1+(1−α)⋅θt
这意味着:
- 较早的参数影响力逐渐减弱(因为乘以了 α \alpha α)。
- 最近的参数更新权重更大(乘以 1 − α 1 - \alpha 1−α)。
- 选择 较大的 α \alpha α(如
0.999
),EMA 更新较慢,适用于平滑长时间的变化。
为什么较早的参数影响力逐渐减弱?
我们可以将 EMA 当前参数展开,看看它是如何由历史所有参数的加权平均组成的:
θ EMA , t = ( 1 − α ) ⋅ θ t + α ( 1 − α ) ⋅ θ t − 1 + α 2 ( 1 − α ) ⋅ θ t − 2 + α 3 ( 1 − α ) ⋅ θ t − 3 + … \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 θEMA,t=(1−α)⋅θt+α(1−α)⋅θt−1+α2(1−α)⋅θt−2+α3(1−α)⋅θt−3+…
这说明:
- 最近的参数 θ t \theta_t θt 乘以 1 − α 1 - \alpha 1−α(即 0.01),虽然数值小,但它是最新的更新,影响直接而强烈。
- 较早的参数 θ t − 1 , θ t − 2 \theta_{t-1}, \theta_{t-2} θt−1,θt−2 乘以 α , α 2 \alpha, \alpha^2 α,α2 等次幂,影响力随着时间推移呈指数级衰减。
- 老的参数贡献依然存在,但比重越来越小,这使得 EMA 更关注近期变化,而不会被早期的不稳定训练步骤影响太多。
💡直觉理解 EMA 的本质是一种带有“记忆衰减”的平滑机制:
- 老的参数不会立刻丢失,但它的影响会随着时间逐步减弱,让新数据有更大的话语权。
- 虽然最近参数的权重(
1 - α = 0.01
)看似小,但它不会被 EMA 继续削弱,因此它的相对影响力更大。- 较早的参数影响力会随着 α t \alpha^t αt 指数级减少,长期来看其贡献会趋近于 0。
如果 α = 0.99 \alpha = 0.99 α=0.99,那么过去 5 个时间步的参数贡献依次为:
Step t : ( 1 − α ) = 0.01 Step t − 1 : 0.99 × 0.01 = 0.0099 Step t − 2 : 0.9 9 2 × 0.01 = 0.009801 Step t − 3 : 0.9 9 3 × 0.01 = 0.00970299 Step t − 4 : 0.9 9 4 × 0.01 = 0.0096059601 \begin{aligned} \text{Step } t: & \quad (1 - \alpha) = 0.01 \\ \text{Step } t-1: & \quad 0.99 \times 0.01 = 0.0099 \\ \text{Step } t-2: & \quad 0.99^2 \times 0.01 = 0.009801 \\ \text{Step } t-3: & \quad 0.99^3 \times 0.01 = 0.00970299 \\ \text{Step } t-4: & \quad 0.99^4 \times 0.01 = 0.0096059601 \\ \end{aligned} Step t:Step t−1:Step t−2:Step t−3:Step t−4:(1−α)=0.010.99×0.01=0.00990.992×0.01=0.0098010.993×0.01=0.009702990.994×0.01=0.0096059601
下面是一个简单的 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()
-
EMA
类- 维护了
shadow
(影子模型参数)。 - 通过
update()
逐步更新 EMA 版本的参数。 apply_shadow()
用于推理时将 EMA 参数应用到原模型上。
- 维护了
-
训练过程中
- 每次模型参数更新后,调用
ema.update()
,让影子模型参数缓慢跟随原模型更新。
- 每次模型参数更新后,调用
-
推理时
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 可以加快采样,同时保持高质量输出。