零基础看懂免费开源的Stable Diffusion

前言

前面一篇文章主要讲了扩散模型的理论基础，还没看过上篇的小伙伴可以点击查看：DDPM理论基础。这篇我们主要讲一下一经推出，就火爆全网的Stable Diffusion模型。Stable Diffusion因其免费，开源，生成图像质量高等优点，一经推出，就火爆全网，后面stable-diffusion-webui的推出，更是降低了使用Stable Diffusion模型作画的门槛，一时刷爆了整个社区。今天笔者详细的带大家看一下Stable Diffusion背后的原理。

Diffusion模型

这里为了让大家更好的理解Stable Diffusion模型，我们先来简单介绍一下Diffusion模型。下图展示了diffusion模型在训练和推理的过程。从图中可以看出，扩散过程主要包括几个主要的模型，分别是text encoder(文本编码器)，unet，image decoder(图像解码器)。其中，text encoder的作用主要是将输入的文本，即prompt，编码为token embeddings，这个token embeddings就是代表文本的一个个向量。这一个个文本向量会通过某种方式注入到unet中，用来控制unet生成符合文本描述的图像。

推理过程

在推理阶段，扩散过程是一个多步去噪的过程，主要就是一个unet网络结构，其输入和输出具有相同的形状，输入为含噪声的图像和时间戳$t$，输出为图像上添加的噪声，进而得到去掉该噪声的图像。就这样经过unet的一步步去噪，逐步生成一个不含噪声的，符合文本描述的图像。有很多人会问，为什么不直接一步预测出噪声，然后直接将该噪声去掉，生成不含噪声的图像呢？其实这样的话，噪声很大，网络很难预测出准确的噪声分布。上图中的N即扩散过程执行的步数，该参数可以由我们自己指定，一般步数设置的越大，生成的图像会越精细。经过扩散过程后会生成低分辨率的，不含噪声的图像，为了生成更高分辨率的图像，这时就会在后面再接一个image decoder，用来扩大图像的分辨率，image decoder输出的图像即为最后我们想要的高分辨率图像。

文本特征主要是通过cross attention模块加入的，我们来大致拆解一下unet内部的网络结构。如下图，unet内部主要是由多个resnet block和attention模块组成的，两者交替出现。每一个attention模块接受resnet block输出的图像特征和文本特征向量作为输入，将两种特征进行融合，从而达到以文本为条件，控制图像生成的目的。

attention内部的计算过程如下图。图像特征和文本特征分别通过三个参数矩阵映射到Q，K，V，然后Q与K的转置点乘除以scale因子后经过softmax计算，最后点乘V，得到最后的特征。当然现在都是基于multi-head的多头attention操作，multi-head只是多次执行下面的操作，得到多个$Z^i$，最后再将$Z^i$拼接在一起，经过最后一个参数矩阵映射得到最终的$Z$。

训练过程

diffusion模型的训练过程主要涉及unet网络的学习，需要让unet具备能力：
给它输入一张含噪图像，unet能够预测出含噪图像上的噪声。
这样我们就可以去掉含噪图像上的噪声，得到一张干净的、不含噪声的图像。训练数据的构造如下图：

首先第一步，选择一张图片；第二步，随机生成一个基础噪声；第三步从0到$T$的时间范围内，随机选择一个时间戳$t$，通过$t$和基础噪声计算出最终要添加的噪声，时间戳$t$越大，代表噪声添加的次数越多，也即添加噪声的强度越大。第四步就是将第三步生成的噪声加到图像上，得到一个含噪声的图像。此时，步骤四中得到的含噪图像作为unet网络的输入，步骤三生成的噪声作为unet学习的目标，用来训练unet网络。

通过上面的1,2,3,4步，我们可以生成很多训练数据，训练过程中就是不断将训练数据喂给unet，让其自主学习如何预测出含噪图像上的噪声，以达到去噪的目的。

Stable Diffusion模型

stable diffusion的最大贡献就是没有直接在像素空间进行图像的加噪和去噪，而是先将图像进行压缩（下采样），压缩到一个图像表征维度更低的隐空间(latent)，然后在隐空间中进行扩散过程，这不仅加快了扩散过程的速度，同时减少了计算资源的消耗，而且在隐空间中操作依然能够保证生成图像的质量。举个例子，如果原图像的分辨率是256x256，现在将它下采样8倍，到32x32，那么在32x32分辨率的图像上操作肯定比直接在256x256分辨率的图像上操作更快且节省资源。256x256分辨率的原图即为像素空间特征，压缩后的32x32分辨率的图就是隐空间特征。

stable diffusion模型在推理和训练阶段的流程图如上图所示。与diffusion模型相比，最大的变化就是在推理阶段，扩散过程的输入由原来的随机噪声图像image变成了随机噪声latent，其实两者本质上都是纯噪声，只不过latent的分辨率比image的分辨率低，所以经过扩散过程生成的latent(Generated low resolution latent)也要比生成的image(Generated low resolution image)分辨率低，如果stable diffusion最终要生成与diffusion模型相同分辨率的图像的话，这里image decoder的放大倍率就要更大。
那么在训练阶段，主要进行前向过程，也就是给图像加噪声。下图就是stable diffusion生成训练数据的过程。可见，和diffusion相比，只是多了一步将原图压缩到latent的过程，后面的添加噪声都是在latent上进行的。

下面我们放一张stable diffusion论文中的原图。

其中$\mathcal{E}$为图像编码器，用来压缩图像尺寸，$\mathcal{D}$为图像解码器用来恢复图像尺寸。图中的Diffusion Process即为前向过程，原图像$x$经过编码器$\mathcal{E}$压缩到隐空间$z$后，在$z$上进行加噪，生成$z_T$。生成过程为逆向过程，给定隐空间噪声，经过多步去噪，生成不含噪声的压缩后的$z$，再经过图像解码器恢复原始图像的尺寸。

参考

https://www.cnblogs.com/gczr/p/14693829.html
https://jalammar.github.io/illustrated-stable-diffusion/
https://readpaper.com/pdf-annotate/note?pdfId=4665140328076951553&noteId=1834381375833065728