源码解析（三）：Stable Diffusion

原文 技术博客

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

主要特点：

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

模型架构

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

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

安装与使用

硬件要求

软件安装

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下载模型

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

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

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 StableDiffusionPipeline

pipe = 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

(

x

t

∣

x

t

−

1

)

=

N

(

x

t

;

α

t

x

t

−

1

,

(

1

−

α

t

)

I

)

q(xt∣xt−1)=N(xt;αtxt−1,(1−αt)I)

q(xt∣xt−1)=N(xt;αtxt−1,(1−αt)I)

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

经过

T

T

T 次迭代，初始图像

x

0

x0

x0 会变成纯噪声

x

T

xT

xT。

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

q

(

x

t

∣

x

0

)

=

N

(

x

t

;

α

ˉ

t

x

0

,

(

1

−

α

ˉ

t

)

I

)

q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)

q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)

其中：

α

ˉ

t

=

∏

i

=

1

t

α

i

αˉt = prod_{i=1}^t alpha_i

αˉt=∏i=1t​αi​

反向过程

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

p

θ

(

x

t

−

1

∣

x

t

)

=

N

(

x

t

−

1

;

μ

θ

(

x

t

,

t

)

,

Σ

θ

(

t

)

)

pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(t))

pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(t))

其中：

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

通过逐步迭代，将纯噪声

x

T

xT

xT 转化为图像

x

0

x0

x0。

损失函数

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

L

V

L

B

=

E

q

[

∑

D

K

L

(

q

(

x

t

−

1

∣

x

t

,

x

0

)

∥

p

θ

(

x

t

−

1

∣

x

t

)

)

]

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

,

ϵ

[

∥

ϵ

−

ϵ

θ

(

x

t

,

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=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_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.py
ldm/models/diffusion/dpm_solver/sampler.py

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

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

性能比较

DDIM采样器

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

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

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：具有略微不同特征的替代实现

采样器使用

在入口点脚本（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 架构高度可配置。以下是关键参数：

核心配置

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-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：数据集特定参数（通常包括图像分辨率）

文章来源于互联网:源码解析（三）：Stable Diffusion

相关推荐: 保姆级教程！如何用GPT-4o快速定制你的手办和视频？

前言 最近像这样的照片变手办真的很火，我也跑了一遍流程，再配合一些 3D 工具可以制作模型，还可以利用 AI 视频工具制作有趣的视频，让我们来一起看看吧。 更多GPT-4o教程： 案例一：GPT 生图 + Tripo3D 建模 + 3D 打印 案例二：GPT …