目录
模型简介 模型结构 VAE变分编码器 VAE结构 VAE模型训练与损失函数 U-Net结构 CLIP Text Encoder 模型工作流程 参考模型简介
Stable Diffusion（SD）模型是一种基于Latent Diffusion Models（LDMs）的生成式模型，总共有1B左右的参数量，可以用于文生图、图生图、等任务中。
文生图任务是将文本输入到SD模型中，输出符合文本描述的图片；图生图任务是在输入文本的基础上，再输入一张图片，模型根据文本的提示，将输入的图片重建为符合文本描述的图片。
对于之前的扩散模型diffusion models，其包括前向过程和反向过程，其中前向过程又成为扩散过程，即对数据逐渐加入高斯噪声直至数据变成随机噪声的过程，是在原始图像中对图片直接进行加噪和去噪处理。
而本文的stable diffusion是一种latent diffusion，是图像经过VAE编码器对图像进行压缩后，再对潜空间（latent space）进行diffusion处理，然后通过解码器将压缩后的latent编码还原为图像。两种扩散模型都使用U-Net网络进行预测噪声。
Latent Diffusion Models（LDMs）论文：《High-Resolution Image Synthesis with Latent Diffusion Models 》ps:stabilityAI团队将LDMs在LAION-5B的一个子集上训练并开源得到stable diffusion
SD代码：https://github.com/CompVis/stable-diffusion
LDMs代码：https://github.com/CompVis/latent-diffusion
模型结构
stable diffusion模型是一个端到端的模型，主要由VAE变分编码器、CLIP text encoder及U-Net三个部分组成。其网络结构图如下。
x：RGB空间表示的图像
ϵ
\epsilon
ϵ:编码器encoder，将x压缩为低维表示z=
ϵ
\epsilon
ϵ(x)
D：解码器decoder，将低维表示z还原为原始图像空间
ϵ
θ
\epsilon_{\theta}
ϵθ:条件去噪自动编码器
τ
θ
\tau_{\theta}
τθ:特定领域编码器，将y（灰框内的控制生成的条件）投影到中间层
VAE变分编码器
VAE结构
VAE编码器是一个基于编码-解码结构的生成模型，其编码器能将输入图像映射到低维latent特征，解码器能将低维latent特征重构为像素级图像。VAE功能大致如下：
当输入尺寸为HWC时，VAE的encoder将其编码为尺寸为hwc的低维latent特征，其中f=W/h=W/w为VAE的下采样率，同时decoder有一个上采样率将低维特征重构为像素级别的图像。f较小时重构效果较好，压缩率较低，但模型收敛速度慢；f较大时重构效果差，压缩损失大，文中设置f=8。通过使用低维潜空间，将图片尺寸减小，模型的计算效率更高。
VAE模型训练与损失函数
主要采用L1回归损失和感知损失，同时采用基于patch的对抗训练策略。
L1回归损失用于衡量预测值与真实值之间的差异，公式如下。
感知损失的核心思想是比较原始图像和生成图像在传统深度学习模型（VGG、ResNet、ViT等）不同层中特征图之间的相似度，而不直接进行像素级别的对比。由于传统深度学习模型能够提取图像的高维语义信息特征，如果两个图像在高维语义信息特征上接近，那么在像素级别的语义上也应该是相似的。公式如下：
基于patch的对抗训练策略，使用patchGAN的判别器对VAE模型进行对抗训练，通过优化判别器损失，来提升生成图像的局部纹理和细节以及清晰度，该方法能够专注于优化生成图像的局部特征，生成更细腻更富有表现力的纹理。
同时为了防止在latent空间的缩放导致的标准差过大，还引入了正则化损失，包括KL正则化和VQ正则化等。
最终stable diffusion中的VAE模型选择KL正则化进行优化训练，设置下采样率f=8，特征维度c=4，当输入尺寸为768768时得到尺寸为9696*4的latent特征。
虽然VAE模型使用了KL正则化，权重系数非常小，实际生成的Latent特征的标准差依旧较大，所以Stable Diffusion论文中提出了一种rescaling方法强化正则效果。首先计算第一个batch数据中Latent特征的标准差
σ
\sigma
σ，然后采用的1/
σ
\sigma
σ系数来rescale后续所有的Latent特征使其标准差接近于1。同时在Decoder模块进行重建时，只需要将生成的Latent特征除以1/
σ
\sigma
σ，再进行像素级重建即可。在SD中，U-Net模型使用的是经过rescaling后的Latent特征，并且将rescaling系数设置为0.18215。
同时，StabilityAI官方也对VAE模型进行了优化。两个优化版本都只优化了VAE的Decoder部分，由于SD在微调训练中只需Encoder部分提供的Latent特征，所以优化训练后的VAE模型可以与开源社区的所有SD模型都兼容。
U-Net结构
U-Net是扩散模型的主体，用来实现条件引导下的图片生成。SD在生成图片时，需要输入prompt（文本信息），那么语义信息如何影响图片的生成呢？那便是注意力机制。
在SD模型的训练中，每个训练样本都会对应一个标签，将对应标签通过CLIP Text Encoder输出Text Embeddings，并将Text Embeddings以Cross Attention的形式与U-Net结构耦合，使得每次输入的图片信息与文字信息进行融合训练。
cross-attention的公式如下：
attention中content影响K、V矩阵，图像影响Q矩阵。
两个输入首先经过Attention机制（将Context Embedding对应的语义信息与图片中对应的语义信息相耦合），输出新的Latent Feature，再将新输出的Latent Feature与输入的Context Embedding再做一次Attention机制，从而使得SD模型学习到了文本与图片之间的特征对应关系。
从整体上看，不管是在训练过程还是前向推理过程，Stable Diffusion中的U-Net在每次循环迭代中Content Embedding部分始终保持不变，而Time Embedding每次都会发生变化。和传统深度学习时代的U-Net一样，Stable Diffusion中的U-Net也是不限制输入图片的尺寸，因为这是个基于Transformer和卷积的模型结构。
在进行Stable Diffusion模型训练时，VAE部分和CLIP部分都是冻结的，所以说官方在训练SD系列模型的时候，训练过程一般主要训练U-Net部分（预测并去除噪声）。
CLIP Text Encoder
CLIP是多模态领域的神器，它实现了文本和图像之间的连通，主要包含Text Encoder和Image Encoder两个模型。其中Text Encoder用来提取文本的特征，可以使用NLP中常用的text transformer模型作为Text Encoder；而Image Encoder主要用来提取图像的特征，可以使用CNN/vision transformer模型（ResNet和ViT）作为Image Encoder。与此同时，他直接使用4亿个图片与标签文本对数据集进行训练，来学习图片与本文内容的对应关系。
CLIP在训练时，从训练集中随机取出一张图片和标签文本。CLIP模型的任务主要是通过Text Encoder和Image Encoder分别将标签文本和图片提取embedding向量，然后用余弦相似度（cosine similarity）来比较两个embedding向量的相似性，以判断随机抽取的标签文本和图片是否匹配，并进行梯度反向传播，不断进行优化训练。
当把训练的Batch提高到N时，CLIP模型将N个标签文本和
N个图片的两两组合预测出
N
2
\ N^2
N2个可能的文本-图片对的余弦相似性，即下图所示的矩阵。CLIP模型的训练目标就是最大N个正样本的余弦相似性，同时最小化
N
2
\N^2
N2个负样本的余弦相似性。
完成CLIP的训练后，输入配对的图片和标签文本，则Text Encoder和Image Encoder可以输出相似的embedding向量，计算余弦相似度就可以得到接近1的结果。同时对于不匹配的图片和标签文本，输出的embedding向量计算余弦相似度则会接近0。
模型工作流程
大致流程如下：
1.从数据集中随机选择一个训练样本
2.从K个噪声量级随机抽样一个timestep t
3.将timestep t对应的高斯噪声添加到图片中
4.将加噪图片输入U-Net中预测噪声
5.计算真实噪声和预测噪声的L2损失
6.计算梯度并更新SD模型参数
对于图生图任务，输入图片通过VAE编码器后将x映射到Latent Feature，进入前向扩散过程，不断加入噪声直至变成随机噪音矩阵，每次迭代用random函数生成从强到弱不同强度的噪声，将随机一种噪声强度加入到图像中，通过time embedding嵌入到训练过程中。Time Embedding由Timesteps（时间步长）编码而来，引入Timesteps能够模拟一个随时间逐渐向图像加入噪声扰动的过程。每个Timestep代表一个噪声强度（较小的Timestep代表较弱的噪声扰动，而较大的Timestep代表较强的噪声扰动），通过多次增加噪声来逐渐改变干净图像的特征分布。
对于文生图任务，只需要输入文本信息，再用random函数生成一个高斯噪声矩阵作为Latent Feature替代输入到SD模型的“图像优化模块”中。
“图像优化模块”是由一个U-Net网络和一个Schedule算法共同组成，U-Net网络负责预测噪声，不断优化生成过程，在预测噪声的同时不断注入文本语义信息。而schedule算法对每次U-Net预测的噪声进行优化处理（动态调整预测的噪声，控制U-Net预测噪声的强度），从而统筹生成过程的进度。
通过U-Net预测噪声后，计算真实噪声和预测噪声的L2损失，计算梯度并更新SD模型的参数。
U-Net网络和Schedule算法的工作完成以后，SD模型会将优化迭代后的Latent Feature输入到图像解码器（VAE Decoder）中，将Latent Feature重建成像素级图像。
参考
https://zhuanlan.zhihu.com/p/640545463
https://blog.csdn.net/imwaters/article/details/127269368
https://zhuanlan.zhihu.com/p/632809634