源码解析（三）：Stable Diffusion_stable diffusion源码

原文 技术博客

😀 Stable Diffusion是一种基于扩散模型（Diffusion Model）的生成式AI技术，通过逐步去噪过程将随机噪声转化为高质量图像。其核心优势在于开源免费、支持本地部署，且能通过文本提示（prompt）精准控制生成内容，还支持图像修改、超分辨率重建等功能。

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.yamlconda activate ldm

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

更新现有环境

condainstallpytorch torchvision -c pytorchpipinstalltransformers==4.19.2diffusersinvisible-watermarkpipinstall-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 termsfrom torch import autocastfrom 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)$

其中：

αˉt= ∏ i = 1 t α i αˉt = \\prod_{i=1}^t \\alpha_i αˉt=∏i=1t​α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 V L B =Eq[∑ D K L (q(xt−1∣xt,x0)∥pθ(xt−1∣xt))] L_{VLB}=Eq[∑D_{KL}(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))] LVLB​=Eq[∑DKL​(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))]

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

L= E t , x 0 , ϵ [∥ϵ−ϵθ(xt,t) ∥ 2 ] L=E_{t,x0,ϵ}[∥ϵ−ϵθ(xt,t)∥^2] L=Et,x0,ϵ​[∥ϵ−ϵθ(xt,t)∥2]

其中：

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

Stable Diffusion的改进

Stable Diffusion 的改进在于：

1. 使用潜在空间表示 z 0 = E n c o d e r ( x 0 ) z_0=Encoder(x_0) z0​=Encoder(x0​)，在潜空间中进行扩散过程。
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_std assert self.num_timesteps_cond <= kwargs[\'timesteps\'] # for backwards compatibility after implementation of DiffusionWrapper if conditioning_key is None: conditioning_key = \'concat\' if concat_mode else \'crossattn\' if cond_stage_config == \'__is_unconditional__\': conditioning_key = None ckpt_path = kwargs.pop(\"ckpt_path\", None) ignore_keys = kwargs.pop(\"ignore_keys\", []) super().__init__(conditioning_key=conditioning_key, *args, **kwargs) self.concat_mode = concat_mode self.cond_stage_trainable = cond_stage_trainable self.cond_stage_key = cond_stage_key try: self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1 except: self.num_downs = 0 if not scale_by_std: self.scale_factor = scale_factor else: 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_forward self.clip_denoised = False self.bbox_tokenizer = None self.restarted_from_ckpt = False if 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.pyldm/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 scriptsif opt.sampler == \'dpm_solver\': sampler = DPMSolverSampler(model)elif opt.sampler == \'ddim\': sampler = DDIMSampler(model)# ...additional samplers...# Samplingsamples, _ = 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.yamlconfigs/latent-diffusion/lsun_churches-ldm-kl-8.yaml

模型结构配置

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

model: base_learning_rate: 2.0e-06 target: ldm.models.diffusion.ddpm.LatentDiffusion params: linear_start: 0.0015 linear_end: 0.0195 num_timesteps_cond: 1 log_every_t: 200 timesteps: 1000 first_stage_key: image image_size: 64 channels: 3 monitor: 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.UNetModel params: image_size: 64 in_channels: 3 out_channels: 3 model_channels: 224 attention_resolutions: - 8 - 4 - 2 num_res_blocks: 2 channel_mult: - 1 - 2 - 3 - 4 num_head_channels: 32

UNet 关键参数：

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

数据配置

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

data: target: main.DataModuleFromConfig params: batch_size: 48 num_workers: 5 wrap: false train: target: taming.data.faceshq.CelebAHQTrain params: size: 256 validation: target: taming.data.faceshq.CelebAHQValidation params: size: 256

关键参数：

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

💡 有关Notion安装或者使用上的问题，欢迎您在底部评论区留言，一起交流~