源码解析（三）：Stable Diffusion

原文 技术博客

Stable Diffusion 是由 CompVis 与 Stability AI 和 Runway 合作开发的潜在文本到图像扩散模型。它通过在压缩的潜在空间（而非像素空间）中执行扩散过程，根据文本描述生成高质量图像，使其在保持高保真输出的同时，比以往的方法具有更高的计算效率。

主要特点：

• 在具有 8 倍下采样因子的压缩潜在空间中运行
• 使用冻结的 CLIP ViT-L/14 文本编码器根据文本提示调节模型
• 具有 860M 参数 UNet 主干，用于扩散过程
• 在至少具有 10GB VRAM 的消费级硬件上运行
• 支持文本到图像的生成和图像到图像的修改

模型架构

稳定扩散的工作原理是将扩散过程转移到压缩的潜在空间，而不是直接作用于像素。这种方法显著降低了计算需求，同时又不牺牲生成质量。

该图展示了稳定扩散流程中的数据流。文本提示被编码到嵌入向量中，这些嵌入向量通过交叉注意力机制来调节扩散过程。对于图像到图像的生成，输入图像被编码为潜在表征。扩散过程在此潜在空间中运行，逐渐将随机噪声张量去噪为有意义的潜在表征，然后将其解码为图像。

安装与使用

硬件要求

硬件	最低要求
图形处理器	NVIDIA GPU，至少配备 10GB VRAM
CUDA	CUDA 11.3 兼容驱动程序
磁盘空间	10GB+（更多用于存储生成的图像）
内存	建议使用 16GB
操作系统	Linux、Windows（带有 WSL）或 macOS

软件安装

Stable Diffusion 需要特定的 Python 环境以及一些依赖项。建议使用 Conda 进行环境管理。

创建新的 Conda 环境

conda env create -f environment.yaml
conda activate ldm

这将创建并激活一个名为 (Latent Diffusion Models) 的 Conda 环境，**ldm**其中包含 environment.yaml 文件中指定的所有必要依赖项。

更新现有环境

condainstallpytorch torchvision -c pytorch
pipinstalltransformers==4.19.2diffusersinvisible-watermark
pipinstall-e .

模型使用

手动下载模型

首先需要在Hugging Face下载模型

下面是官方的比较旧的基础模型，其他新的开源模型可以在开源社区直接下载使用。

模型	描述	训练步骤
sd-v1-1.ckpt	在 laion2B-en 和 laion-highresolution 上训练的基础模型	237k（256×256）、194k（512×512）
sd-v1-2.ckpt	从 v1-1 恢复，在 laion-aesthetics v2 5+ 上进行训练	515k（512×512）
sd-v1-3.ckpt	从 v1-2 恢复，文本条件丢失 10%	195k（512×512）
sd-v1-4.ckpt	从 v1-2 恢复，文本条件丢失 10%	225k（512×512）

获取权重后，需要将它们放置在正确的位置：

mkdir -p models/ldm/stable-diffusion-v1/
ln -s <path/to/model.ckpt> models/ldm/stable-diffusion-v1/model.ckpt

完成安装过程后，您可以通过运行简单的文本到图像生成来验证一切是否设置正确：

python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms

直接通过函数自动下载

对于喜欢更简单安装过程的用户，也可以通过 Hugging Face 扩散器库使用 Stable Diffusion：

# Requires registration at Hugging Face and acceptance of model terms
from torch import autocast
from diffusers import StableDiffusionPipelinepipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=True
).to("cuda")prompt = "a photo of an astronaut riding a horse on mars"
with autocast("cuda"):image = pipe(prompt)["sample"][0]  image.save("astronaut_rides_horse.png")

核心代码详解

扩散模型原理简介

核心思想

扩散模型的主要步骤：

1. 前向过程（Forward Process）：逐步向图像添加高斯噪声，最终将图像转化为纯噪声。
2. 反向过程（Reverse Process）：从纯噪声开始，通过神经网络逐步去噪，重构出目标图像。

Stable Diffusion 的独特之处在于：

• 潜空间操作：不是直接在像素空间中进行操作，而是将图像映射到潜在空间（通过预训练的 VAE 等模型），在潜空间中执行扩散过程，从而提高计算效率和生成质量。

核心公式

前向过程定义为将图像逐步添加噪声：

$$q(xt\mid xt-1)=N(xt;\alpha txt-1,(1-\alpha t)I)$$

• $xt$: 第 t 步的图像。
• $\alpha t$: 控制噪声量的超参数（通常表示为递减的序列）。
• $N(\cdot )$: 正态分布。

经过 $T$ 次迭代，初始图像 $x0$ 会变成纯噪声 $xT$。

直接从初始图像到任意时刻 t 的分布为：

$q(xt\mid x0)=N(xt;\alpha ˉtx0,(1-\alpha ˉt)I)$

其中：

$\alpha ˉt={\prod }_{i=1}^t{\alpha }_i$

反向过程

反向过程的目标是从噪声恢复图像：

$p\theta (xt-1\mid xt)=N(xt-1;\mu \theta (xt,t),\Sigma \theta (t))$

其中：

• $\mu \theta (xt,t)$: 神经网络预测的均值。
• $\Sigma \theta (t)$: 预测的不确定性（通常可以设为固定值）。

通过逐步迭代，将纯噪声 $xT$ 转化为图像 $x0$。

损失函数

模型使用的目标是重构数据分布，通常基于变分下界 (Variational Lower Bound, VLB)：

$L_{VLB}=Eq[\sum D_{KL}(q(xt-1\mid xt,x0)\Vert p\theta (xt-1\mid xt))]$

实际实现中，简化为对噪声的预测：

$L=E_{t,x0,ϵ}[\Vert ϵ-ϵ\theta (xt,t){\Vert }^2]$

其中：

• $ϵ$: 真实的噪声。
• $ϵ\theta (xt,t)$: 神经网络预测的噪声。

Stable Diffusion的改进

Stable Diffusion 的改进在于：

1. 使用潜在空间表示 $z_0=Encoder(x_0)$，在潜空间中进行扩散过程。
2. 结合条件输入（如文本描述）来控制生成结果，通过条件扩散模型生成潜在表示。

核心组件

LDM

LDM 是稳定扩散的核心架构，它使稳定扩散能够根据文本提示生成高质量图像或修改现有图像。与在像素空间中运行的传统扩散模型不同，LDM 在压缩的潜在空间中运行，从而显著降低了计算需求，同时保持了生成质量。

LatentDiffusionLDM 通过中的类实现ddpm.py，该类扩展了基类**DDPM**（去噪扩散概率模型）。

class LatentDiffusion(DDPM):"""main class"""def __init__(self,first_stage_config,cond_stage_config,num_timesteps_cond=None,cond_stage_key="image",cond_stage_trainable=False,concat_mode=True,cond_stage_forward=None,conditioning_key=None,scale_factor=1.0,scale_by_std=False,*args, **kwargs):self.num_timesteps_cond = default(num_timesteps_cond, 1)self.scale_by_std = scale_by_stdassert self.num_timesteps_cond <= kwargs['timesteps']# for backwards compatibility after implementation of DiffusionWrapperif conditioning_key is None:conditioning_key = 'concat' if concat_mode else 'crossattn'if cond_stage_config == '__is_unconditional__':conditioning_key = Noneckpt_path = kwargs.pop("ckpt_path", None)ignore_keys = kwargs.pop("ignore_keys", [])super().__init__(conditioning_key=conditioning_key, *args, **kwargs)self.concat_mode = concat_modeself.cond_stage_trainable = cond_stage_trainableself.cond_stage_key = cond_stage_keytry:self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1except:self.num_downs = 0if not scale_by_std:self.scale_factor = scale_factorelse:self.register_buffer('scale_factor', torch.tensor(scale_factor))self.instantiate_first_stage(first_stage_config)self.instantiate_cond_stage(cond_stage_config)self.cond_stage_forward = cond_stage_forwardself.clip_denoised = Falseself.bbox_tokenizer = None  self.restarted_from_ckpt = Falseif ckpt_path is not None:self.init_from_ckpt(ckpt_path, ignore_keys)self.restarted_from_ckpt = True

该类**LatentDiffusion**保留了几个重要的属性：

• first_stage_model：用于潜在空间操作的自动编码器
• cond_stage_model：文本编码器（通常为 CLIP）
• scale_factor：控制潜在表示的缩放比例
• conditioning_key"crossattn""concat"：确定调节机制（"crossattn"或"concat"）

扩散过程

前向扩散过程会根据预先定义的方案添加噪声以清理潜在表示。这在训练过程中使用，并通过以下**q_sample()**方法实现：

def q_sample(self, x_start, t, noise=None):# Adds noise to x_start according to timestep t# Formula: sqrt(alpha_cumprod) * x_start + sqrt(1-alpha_cumprod) * noise

去噪过程

逆扩散过程在条件反射的引导下，从纯噪声开始逐渐消除噪声。这在生成过程中使用，并通过**p_sample()和p_sample_loop()**方法实现。

对于每个去噪步骤：

1. UNet 预测当前潜在
2. 该模型使用此预测来计算下一个（噪声较小的）潜在
3. 该过程重复，直到达到干净潜伏期（t=0）

文本控制

文本控制是 LDM 的关键环节，能够实现文本到图像的生成。其主要机制是文本嵌入与 UNet 中间表示之间的交叉注意力机制：

交叉注意力机制的工作原理如下：

1. 文本嵌入在注意力机制中充当“键”和“值”
2. UNet 特征作为“查询”
3. 这使得 UNet 能够在去噪过程中选择性地关注文本嵌入的相关部分

采样方法

采样方法是一种算法，用于确定如何逐步去噪潜在噪声以生成图像。采样器的选择会影响图像质量和生成速度，并针对不同的用例提供不同的权衡。

代码文件位置

ldm/models/diffusion/ddim.py
ldm/models/diffusion/dpm_solver/sampler.py

Stable Diffusion 实现了多种采样策略，每种策略都有其独特的特点：

1. DDIM（去噪扩散隐式模型）：一种确定性采样方法，可以用更少的步骤实现更快的采样。
2. PLMS（伪线性多步）：一种改进 DDIM 的高阶求解器。
3. DPM-Solver：具有高阶收敛的快速求解器，提供最佳的速度/质量权衡

性能比较

采样器	速度	质量	建议步骤	用例
DDIM	缓和	好的	50-100	速度与质量的良好平衡
PLMS	缓和	更好的	50-70	当质量比速度更重要时
DPM求解器	快速地	出色的	15-30	适合大多数用例的最佳整体选择

DDIM采样器

DDIM（去噪扩散隐式模型）采样器是一种确定性采样方法，与原始 DDPM（去噪扩散概率模型）方法相比，其采样步骤明显减少

1. 调度：该**make_schedule**方法为采样过程创建一系列时间步长，在完整扩散时间步长和减少的子集之间进行转换，以实现更快的采样。
2. 采样过程：
   • 该**sample**方法就是主要入口点
   • 它准备参数并调用**ddim_sampling**
   • 这将执行迭代去噪，**p_sample_ddim**在每一步调用
3. 单步去噪：
   • **p_sample_ddim**执行单个去噪步骤
   • **unconditional_guidance_scale**它通过参数支持无分类器引导
   • 参数**eta**控制随机性 **eta=0.0**过程是确定性的

范围	描述	影响
S	采样步数	步骤越少，速度越快，但可能会降低质量
eta	控制随机性	0.0 表示确定性，1.0 表示随机性
unconditional_guidance_scale	引导强度	更高的价值会增加对提示的遵守

DPM求解器

DPM-Solver 是一种先进的快速采样器，可为扩散模型提供高阶收敛，只需很少的步骤即可获得高质量的结果。

DPM-Solver 由多个组件组成：

1. DPMSolverSampler（）:

   • 与稳定扩散模型接口的类
   • 实现**sample**管道期望的方法

   ldm/models/diffusion/dpm_solver/sampler.py
   

2. 噪音时间表 VP

   • 处理方差保持 (VP) SDE 的噪声调度
   • 支持离散时间和连续时间扩散模型
   • 提供计算 alpha、sigma 和 lambda 值的方法

   ldm/models/diffusion/dpm_solver/dpm_solver.py
   

3. DPM_求解器

   • DPM-Solver算法的核心实现
   • 支持单步和多步求解器
   • 实现一阶、二阶和三阶方法

   ldm/models/diffusion/dpm_solver/dpm_solver.py
   

主要特点

高阶求解器：

• 一阶：相当于DDIM
• 二阶：精度更高，计算开销更小
• 三阶：精度最高，尤其在步数较少的情况下

解决方法：

• 单步：直接应用高阶方法
• 多步骤：使用前面步骤的信息来提高效率

求解器类型：

• dpm_solver：默认求解器类型
• taylor：具有略微不同特征的替代实现

范围	描述	典型值
steps（S）	求解器步骤数	10-50（少于 DDIM）
skip_type	间隔时间步长的方法	‘time_uniform’（高分辨率默认）
method	解决方法	“多步”或“单步”
order	求解器的顺序	1、2 或 3

采样器使用

在入口点脚本（txt2img.py/img2img.py）中使用采样器时，采样器的选择和配置如下：

# Example from usage in scripts
if opt.sampler == 'dpm_solver':sampler = DPMSolverSampler(model)
elif opt.sampler == 'ddim':sampler = DDIMSampler(model)
# ...additional samplers...# Sampling
samples, _ = sampler.sample(S=opt.steps,batch_size=opt.n_samples,shape=shape,conditioning=c,unconditional_guidance_scale=opt.scale,unconditional_conditioning=uc,
)

Unet框架

UNet 架构高度可配置。以下是关键参数：

范围	描述	默认值
image_size	输入图像的大小	64（潜在空间）
in_channels	输入通道数	4
model_channels	基本通道数	320
out_channels	输出通道数	4
num_res_blocks	每个下采样级别的 ResBlocks	2
attention_resolutions	需要关注的决议	（4，2，1）
dropout	辍学率	0.0
channel_mult	每级通道乘数	（1、2、4、8）
conv_resample	使用学习到的卷积进行采样	真的
dims	信号的维度	2
num_heads	注意力头数量	8
num_head_channels	每个注意力头的通道数	64
use_scale_shift_norm	使用类似 FiLM 的调理	真的
use_spatial_transformer	使用 CrossAttention 进行文本调节	真的
transformer_depth	Transformer 交叉注意力深度	1
context_dim	上下文（文本）嵌入的维度	768 或 1024

核心配置

Stable Diffusion 使用 YAML 配置文件来定义模型架构、训练参数和数据处理流程。这些文件提供了一种声明式的方式来指定整个实验设置，而无需修改代码。

文件位置：

configs/latent-diffusion/ffhq-ldm-vq-4.yaml
configs/latent-diffusion/lsun_churches-ldm-kl-8.yaml

模型结构配置

该**model**部分是配置文件中最重要的部分。它定义了稳定扩散模型的核心架构和超参数。

model:base_learning_rate: 2.0e-06target: ldm.models.diffusion.ddpm.LatentDiffusionparams:linear_start: 0.0015linear_end: 0.0195num_timesteps_cond: 1log_every_t: 200timesteps: 1000first_stage_key: imageimage_size: 64channels: 3monitor: val/loss_simple_ema

关键参数包括：

• base_learning_rate：优化的基础学习率
• target：模型实现的 Python 类路径
• params：特定于模型的参数
  • **linear_start和linear_end**定义噪音时间表
  • timesteps：扩散步骤数
  • image_size：潜在表征的大小
  • channels：潜在空间中的通道数

UNet配置

UNet 是扩散模型的骨干，负责在去噪过程中预测噪声。

unet_config:target: ldm.modules.diffusionmodules.openaimodel.UNetModelparams:image_size: 64in_channels: 3out_channels: 3model_channels: 224attention_resolutions:- 8- 4- 2num_res_blocks: 2channel_mult:- 1- 2- 3- 4num_head_channels: 32

UNet 关键参数：

• image_size：输入的空间维度
• in_channel和sout_channels：输入和输出通道维度
• model_channels：模型的基本通道数
• attention_resolutions：在哪个分辨率级别应用注意力
• num_res_blocks：每个分辨率级别的残差块数量
• channel_mult：每个分辨率级别的通道乘数

数据配置

本**data**节定义数据集加载、预处理和批处理参数。

data:target: main.DataModuleFromConfigparams:batch_size: 48num_workers: 5wrap: falsetrain:target: taming.data.faceshq.CelebAHQTrainparams:size: 256validation:target: taming.data.faceshq.CelebAHQValidationparams:size: 256

关键参数：

• batch_size：每批次样品数量
• num_workers：数据加载的工作进程数
• wrap：是否包装数据集
• train和validation：训练和验证集的数据集配置
  • target：数据集实现的 Python 类路径
  • paramssize：数据集特定参数（通常包括图像分辨率）