【扩散模型】12、Stable Diffusion | 使用 Diffusers 库来看看 Stable Diffusion 的结构_diffusers库

参考：HuggingFace

参考：https://jalammar.github.io/illustrated-stable-diffusion/

一、什么是 Stable Diffusion

Stable Diffusion 这个模型架构是由 Stability AI 公司推于2022年8月由 CompVis、Stability AI 和 LAION 的研究人员在 Latent Diffusion Model 的基础上创建并推出的。

其原型是（Latent Diffusion Model），一般的扩散模型都需要直接在像素空间训练和运行，训练可能需要上百个 GPUs，在测试推理的时候也需要很多硬件支持。所以就有了 LDM 这个模型的提出，通过平衡【降低复杂度】和【保持图像细节】，能在保持保真度的同时实现模型的加速。

Stable Diffusion 是什么：

• Stable Diffusion 是一种基于扩散过程的图像生成模型，可以生成高质量、高分辨率的图像
• 有较强的稳定性和可控性，可以生成具有多样化效果且良好视觉效果的图片

Stable Diffusion 和 Midjourney 的对比：

• 开源情况：Stable Diffusion 开源，Midjourney 闭源
• 质量情况：Midjourney 生成的图片的质量更高
• 可控情况：Stable Diffusion 生成图片的可控性更高

当图片尺寸变大时，需要的计算能力也随之增加。这种现象在自注意力机制（self-attention）这种操作的影响下尤为突出，因为操作数随着输入量的增大呈平方地增加。一个 128px 的正方形图片有着四倍于 64px 正方形图片的像素数量，所以在自注意力层就需要16倍的内存和计算量。这是高分辨率图片生成任务的普遍问题。

隐式扩散致力于克服这一难题，它使用一个独立的模型 Variational Auto-Encoder（VAE）压缩图片到一个更小的空间维度。这背后的原理是，图片通场都包含了大量的冗余信息。因此，我们可以训练一个 VAE，并通过大量足够的图片数据训练，使得它可以将图片映射到一个较小的隐式空间，并将这个较小的特征映射回原图片。SD 模型中的 VAE 接收一个三通道图片输入，生成出一个四通道的隐式表征，同时每一个空间维度上都减少为原来的八分之一。比如，一个 512px 的正方形图片将会被压缩到一个 4×64×64 的隐式表征上。

通过在隐式表征上（而不是完整图像上）进行扩散过程，我们可以使用更少内存、减少 UNet 层数、加速生成。同时我们仍能把结果输入 VAE 的解码器，解码得到高分辨率图片。这一创新点极大地降低了训练和推理成本。

这里再回顾一下基础模型 LDM：

• 输入 Image 为 $x$，输出生成的 Image 为 $\tilde{x}$
• 先看上面向右传播的这行，$ϵ$ 为编码器，将输入图片编码到潜在空间，然后进行前向扩散加噪，得到 $z_T$。这里的编码器使用的 VAE
• 然后看最右侧的黑线框，展示的不同的输入，包括语义图、文本、图像等等，经过 domain-space 的编码器 ${\tau }_{\theta }$ 的处理后，输入去噪网络
• 最后看下面这行逆向扩散的这行，在去噪网络中，级联了多个 U-Net 网络，具体数量取决于 $T$ 的值，在每个时刻都使用 U-Net 的结构，U-Net 的每个中间层的输入都是 ${\tau }_{\theta }(y)$ 和 $z_T$。D 是从潜在空间又解码到图像空间，最后得到了生成的图片

Stable Diffusion模型中使用变分自编码器（VAE）作为编码器的原因主要有以下几点：

• 生成模型：VAE是一种生成模型，能够学习数据的潜在表示，并且可以从这个潜在空间中抽样来生成新的数据。这与Stable Diffusion的目标，即学习复杂数据分布并从中抽样，非常契合。

• 变分推理：VAE利用变分推理进行参数估计和优化。这使得它能够有效地处理大规模和高维度的数据，这对于许多实际应用（如图像和文本）来说是非常重要的。

• 损失函数：VAE具有结构化的损失函数，包括重构损失和KL散度。这使得我们可以明确地控制重构质量与潜在表示之间的平衡。

• 鲁棒性：由于其随机性质，VAE通常比确定性方法更鲁棒，在面对噪声或缺失值时表现更好。

• 连续潜在空间：VAE假设一个连续的潜在空间，并试图将输入映射到此空间上。连续性使得我们可以通过插值和外推等操作来探索未知区域，并可能帮助改善生成结果的质量。

VAE 是什么：

是一种生成模型，它使用深度学习技术来提供数据的紧凑连续潜在表示。VAE由编码器和解码器两部分组成。

• 编码器：编码器将输入数据（例如图像）映射到一个潜在空间。每个输入都被映射为该空间中的一个点，或者更准确地说，是一个分布。这个分布通常假设为高斯分布，由均值和方差定义。

• 解码器：解码器则执行相反的操作，它从潜在空间中取样，并将这些样本映射回原始数据空间（例如生成图像）。解码过程也是随机的，在给定潜在变量的情况下，输出是原始数据空间上的一个条件分布。

VAE通过最大化下界（ELBO, Evidence Lower BOund）进行训练。ELBO包含两部分：

• 重构损失：测量解码后的输出与原始输入之间的差异。
• KL散度：测量编码后得到的高斯分布与标准正态分布之间的差异。
  通过优化这两项内容使得模型能够有效地学习出数据背后复杂且有用的特征，并且保证了潜在空间具有良好结构性质以便于采样和推理。

总结一句话就是, VAE 是一种利用深度神经网络进行参数化并利用变分推理进行训练以学习复杂数据集隐含结构并能从中生成新样本 的生成模型。

二、Diffusers 库

Diffusers 库如何使用：

git clone https://github.com/huggingface/diffusers.git
#可以使用 diffusers 这个 image 起容器
nvcr.io/nvidia/pytorch:22.08-py3
# 然后在容器中安装，为了更方便看内部，我使用了 -e . 的安装方式，每次调用库都会进入我自己的 diffusers 库
cd diffusers
pip install -e .
1
2
3
4
5
6

Diffusers 库是什么：

• diffusers 是 Hugging Face 推出的一个扩散模型库，它提供了简单方便的推理、训练 pipeline，同时拥有一个模型和数据社区，代码可以像 torchhub 一样直接从指定的仓库去调用别人上传的数据集和 pretrain checkpoint
• diffusers 库包括了 SOTA 的扩散模型的管道，只需几行代码就可以在推理中运行，生成想要的图像
• 可使用不同的噪声调度器，用于平衡不同的扩散速度和输出质量。
• 有很多预训练好的模型，可以作为构建模块使用，并与调度器结合，以创建自己的端到端扩散系统。

扩散模型的主要步骤：

• 获取一批数据
• 添加随机噪声
• 将数据输入模型
• 将模型预测与干净图像进行比较，以计算loss
• 更新模型的参数

一个简单的示例如下：
# Dataloader (you can mess with batch size)
batch_size = 128
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# How many runs through the data should we do?
n_epochs = 3

# Create the network
net = BasicUNet()
net.to(device)

# Our loss finction
loss_fn = nn.MSELoss()

# The optimizer
opt = torch.optim.Adam(net.parameters(), lr=1e-3) 

# Keeping a record of the losses for later viewing
losses = []

# The training loop
for epoch in range(n_epochs):

    for x, y in train_dataloader:

        # Get some data and prepare the corrupted version
        x = x.to(device) # Data on the GPU
        noise_amount = torch.rand(x.shape[0]).to(device) # Pick random noise amounts
        noisy_x = corrupt(x, noise_amount) # Create our noisy x

        # Get the model prediction
        pred = net(noisy_x)

        # Calculate the loss
        loss = loss_fn(pred, x) # How close is the output to the true 'clean' x?

        # Backprop and update the params:
        opt.zero_grad()
        loss.backward()
        opt.step()

        # Store the loss for later
        losses.append(loss.item())

    # Print our the average of the loss values for this epoch:
    avg_loss = sum(losses[-len(train_dataloader):])/len(train_dataloader)
    print(f'Finished epoch {epoch}. Average loss for this epoch: {avg_loss:05f}')

# View the loss curve
plt.plot(losses)
plt.ylim(0, 0.1);
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52