Stable diffusion 简介

news2024/11/22 21:03:00

Stable diffusion 是 CompVis、Stability AI、LAION、Runway 等公司研发的一个文生图模型，将 AI 图像生成提高到了全新高度，其效果和影响不亚于 Open AI 发布 ChatGPT。Stable diffusion 没有单独发布论文，而是基于 CVPR 2022 Oral —— 潜扩散模型 (Latent Diffusion Model, LDM)。

Stable diffusion 模型并不是直接在像素空间上操作图像，而是在隐空间中进行操作：通过将原始数据编码到更小的隐空间中，让 U-Net 可以在低维表示上添加和删除噪声，这样可以显著提高模型的训练和生成效率。此外，stable diffusion 还引入了 CLIP，将输入文本的特征与模型生成的图像做比对以输出最符合文本的图像。

在这里插入图片描述

一. 背景

在 Stable Diffusion 诞生之前，AIGC 最主要的模型是生成对抗网络 GAN，通过同时训练两个网络进行对抗训练，实现高质量的图像生成。GAN 让超越训练数据已有内容成为可能，从而打开了一个全新领域 —— 生成建模。¹ 然而，在经历了一段蓬勃发展后，GAN 开始暴露出一些瓶颈和弊病：图像生成缺乏多样性、模式崩溃、多模态分布学习困难、训练时间长，GAN 由此进入瓶颈期。

Diffusion 的出现打破了 AIGC 的僵局，通过向数据中添加噪声再去噪恢复数据，从而实现图像的自由生成。Diffusion 自 DDPM 的提出就迅速收获大量好评，现在大火的最 GLIDE、DALLE2、Imagen，以及一系列的 Image Editing 方法，都是基于 diffusion。然而，diffusion 的训练时间和经济成本都极其昂贵，于是出现了 stable diffusion。

Stable diffusion 模型不直接在像素空间上操作图像，而是在隐空间中进行操作。通过将原始数据编码到更小的隐空间中，让 U-Net 可以在低维表示上添加和删除噪声，这样能够加速 U-Net 处理数据的速度，实现更高的生成效率。²
在这里插入图片描述

为了实现文生图的效果，stable diffusion 还在 U-Net 中引入了 CLIP text encoder 提取文本特征，从而生成更加符合文本的图像。

二. Stable diffusion 模型

Stable diffusion 是一个基于 latent 的扩散模型，采用一个 autoencoder 将图像压缩到 latent 空间，然后用扩散模型来生成图像的 latents，最后送入 decoder 模块就可以得到生成的图像。Stable diffusion 还在 U-Net 中引入 text condition 实现了基于文本的生成图像。³
在这里插入图片描述

Stable diffusion 模型的主体结构如图所示，主要包括三个模型：

Autoencoder：将图像在 pixel 空间和 latent 空间之间压缩和解压缩；
CLIP text encoder：提取输入文本的 text embeddings 传入 U-Net；
U-Net：扩散模型的主体，用来实现文本引导下的图像 latents 生成。

1. Autoencoder

为了让 stable diffusion 能够在像素空间和隐空间之间做变换，引入了编码器和解码器：

编码器 $E$ ：将全尺寸图像编码为低维潜在数据；
解码器 $D$ ：将潜在数据解码回全尺寸图像；

Autoencoder 将 $\times W \times 3$ 大小的 RGB 图像 $x$ 压缩成 $\times w \times c$ 大小的隐空间表示， $f = H / h = W / w$ 称为下采样率 (downsampling factor)。LDM 尝试了 $f=2^m$ 的不同取值，权衡生成质量和收敛速度。

2. CLIP text encoder

为了实现文生图的效果，stable diffusion 使用预训练的 CLIP text encoder 提取文本特征 text embeddings 传入 U-Net，以指导其图像生成。Stable diffusion 采用的 CLIP text encoder 是 clip-vit-large-patch14 的 text encoder，是一个 12 层、768 特征维度的 transformer 模型，得到 77x768 特征维度的 text embeddings，将其送入 U-Net 的 Cross-Attention 中。
在这里插入图片描述

图中的开 switch 用于在不同类型的输入之间进行调节：

对于文本输入，首先使用语言模型 $\tau_{\theta}$ 将文本转换为嵌入向量，然后通过 Multi-Head Attention 映射到 U-Net；
对于其他模态的对齐输入（例如语义映射、图像、修复），使用连接来完成调节；

3. U-Net

U-Net 是扩散模型的主体，用来实现文本引导下的图像 latents 生成。为了实现文本指导，stable diffusion 通过使用 cross-attention 机制增强 U-Net，将内部扩散模型转变为条件图像生成器。

Stable Diffusion 的训练目标与 diffusion 的目标较为相似，需要改进的地方是 ⁴：

图像 $x_t$ 变为了隐空间中的变量 $z_t$ ；
U-Net 中添加了条件输入 $\tau_{\theta}(y)$ ；

因此 Stable Diffusion 的损失函数如下：
$L_{\mathrm{LDM}}=\mathbb{E}_{t, z_0, \epsilon, y}\left[\left\|\epsilon-\epsilon_\theta\left(z_t, t, \tau_\theta(y)\right)\right\|_2^2\right]$

更多介绍见 stable diffusion原理解读通俗易懂，史诗级万字爆肝长文，喂到你嘴里。

三. 训练方法

在这里插入图片描述

四. 图像生成

在这里插入图片描述

五. 实验

六. 总结

七. 复现

Stable diffusion 开源了预训练好的权重：the CompVis organization at Hugging Face，可以直接拿来做下游任务：

平台：AutoDL
显卡：RTX 3090 24G
镜像：PyTorch 1.11.0、Python 3.8(ubuntu20.04)、Cuda 11.3
源码：https://github.com/CompVis/stable-diffusion

Hugging Face 上提供了预训练好的不同版本的 checkpoints，以 stable-diffusion-v-1-1 为例，可以直接下载模型参数：CompVis/stable-diffusion-v-1-1-original。Hugging Face 上还将 stable diffusion 封装成 diffusers 库，可以直接调用完成下游任务，demo 如下：

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

model_id = "CompVis/stable-diffusion-v1-1"
device = "cuda"


pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe = pipe.to(device)

prompt = "a photo of an astronaut riding a horse on mars"
with autocast("cuda"):
    image = pipe(prompt).images[0]  
    
image.save("astronaut_rides_horse.png")

实验记录：

一开始在 Notebook 里运行，即使进行学术资源加速，还是会报错 SSLCertVerificationError。后来直接在终端执行 python 文件就可以正常连接了；
执行 image = pipe(prompt)["sample"][0] 时出现 KeyError: 'sample'，将其修改为 image = pipe(prompt).images[0] ⁵ 就可以正常使用了；
第一次执行需要下载模型文件和加载管道组件，费时较长，耐心等待即可。以后每次调用就可以直接加载；