"一.前言上次的文章王建周：AIGC 基础，从VAE到DDPM 原理、代码详解对图像生成的基本原理已经做了较为详细的介绍，但是在真实的产业界，要实现可控的图像生成，还是较多实际的问题需要解决比如以下 1.控制条件是如何注入到模型中来控制生成的图像 2.如何生成细节足够逼真的高清大图 3.如何在模型 ...."

Fork me on GitHub

AIPreacher
AI发布小助手
• 0 回帖 • 1.7K 浏览 • 1 年前

基于Diffusion的典型可控图片生成模型

以下文章来源于 https://zhuanlan.zhihu.com/p/610952389

一.前言

上次的文章王建周：AIGC 基础，从VAE到DDPM 原理、代码详解对图像生成的基本原理已经做了较为详细的介绍，但是在真实的产业界，要实现可控的图像生成，还是较多实际的问题需要解决比如以下

1.控制条件是如何注入到模型中来控制生成的图像

2.如何生成细节足够逼真的高清大图

3.如何在模型设计中，减少参数和计算量

4.除了文本，如何用其他形式的条件控制

带着这些典型问题，本文试图尝试分享目前文本控制图像生成模型最具有代表性几个模型：OpenAI的DALL-E-2【1】、Google的ImaGen【2】、Stability.AI的Latent Stable Diffusion【3】，最新斯坦福的ControlNet 【4】这些模型都用到上文介绍diffusion技术。因为这些生成模型都是由多个典型的技术模块组成，不是一个单一的模型，所以在分享以上模型的的时候，也会对这些模型中涉及的几个其他典型模型做简单介绍，如CLIP【5】，T5【6】，基于Diffusion的超分辨率模型。

二.DALL-E-2

DALL-E-2本质上属于上文提到的Classifier-free diffusion model，其中condition 就是输入的文本，模型整体架构如下图，可以先忽略绿框部分（这部分属于CLIP）

整体上模型分成4大部分（一部分图上没有呈现，属于文图生成的非必要部分），整个生成过程也是这4个部分依次执行

1.Text encoder。将输入的文本编码成向量。

2.Prior 模块。将文本向量转换为图片向量，转换后的图片向量作为解码器的控制条件。

3.Decoder。通过diffusion model 将图片向量解码为图片。

4.SR超分模块。将第3步生成的图像通过超分技术变成更大分辨率的图片。

这4部分在DALL-E-2中分开训练而不是端到端的联合训练，接下来将详细介绍这4部分内容（因为DALL-E-2的原始论文缺失较多细节，所以有一些个人的主观猜测）

1.Text encoder

Text encoder的目的是将文本转化为数学表示的向量，希望在这个过程可以捕获文本语义、同时可以把图片描述的文本语义和图像语义联系起来。在DALL-E-2模型中，使用了CLIP的模型。CLIP模型一样来自OpenAI，出发点很直白,通过对比学习将图片和图片描述文本链接起来，希望一个图片的表示向量 $I_i$ 和这个图片描述文本 $T_i$ 的表示向量相近，与此同时希望 $I_i$ 和其他图片文本向量 $T_{j \ne i}$ 足够远，如下图。

伪代码实现如下，对比学习最常用是用InfoNCE 作为Loss

def clip_training(imgs, texts):
    # imgs 和texts的batch size 尽可能大，对比学习bs越大，效果越好
    # 图像编码器和文本编码器将图像和文本编码成维度相当的embedding
    img_embedding = img_encoder(imgs)
    txt_embedding = text_encoder(texts)
    norm_img_embedding = tf.nn.l2_normalize(img_embedding, -1)
    norm_txt_embedding = tf.nn.l2_normalize(txt_embedding, -1)
    logits = tf.matmul(norm_txt_embedding,norm_img_embedding,transpose_b=True)
    batch_size = tf.range(tf.shape(imgs)[0])
    label = tf.range(batch_size)
    # infonce 的简化实现，认为相似度矩阵的对角线部分为label=1,其他部分label=0
    loss1 = tf.keras.losses.sparse_categorical_crossentropy(
            label, logits, from_logits=True
        )
    # 非对角部分交换，实现t和i,i和t的计算
    loss2 = tf.keras.losses.sparse_categorical_crossentropy(
            label, tf.transpose(logits), from_logits=True
        )
    return (loss1+loss2)/2.0

CLIP模型使用40亿的图文对训练的，可以很好的关联文本语义和图像语义，通过构造prompt在分类任务上，在一些抽象的图像上可以超越SOTA的图像分类模型。

CLIP训练好了后进行冻结，后边会使用CLIP里边训练好的img_encoder和text_ecoder这两个模型。

2.Prior模块

对于训练集合中图片x和描述y(控制条件），以及x和y分别经过CLIP得到图像和文本特征向量 $z_i,z_t$ ，我们期望的似然模型为 $p(x|y)$ ，可以经过以下推导

$\begin{align} p(x|y)=p(x,z_i|y)=\frac{p(x,z_i,y)}{p(y)}=\frac{p(x|z_i,y)p(z_i,y)}{p(y)}=p(x|z_i,y)p(z_i|y) \end{align}$

说明引入由y得到的 $z_i$ ，模型是等价的。

如何由y得到 $z_i$ ，论文中作者尝试了自回归和扩散模型，经过效果和效率对比，最终使用了扩散模型，这个扩散模型的输入由以下组成：

1.y，也就是图片的描述

$z_t$ ，也就是图片描述经过CLIP文本编码器得到的向量

3.t，也就是扩散模型中的时间步

常见情况下，扩散模型预测噪音效果会比较好，但是在DALL-E-2的Prior模块中，作者发现直接预测 $z_i$ 效果会比较好，公式推导如下，从上文的公式的88可以得到

$\begin{align} \tilde\mu(x_t,x_0) =& \frac{\sqrt{\alpha_t} \bar\beta_{t-1}}{\bar\beta_t}x_t+ \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{\bar\beta_t}x_0 \end{align}$

同预测噪音的逻辑如下，在已知 $x_t$ 的情况下，可以设计下边目标（认为 $z_i^t就是x_t,z_i^0就是x_0$ ）

$\begin{align} \mathcal f_\theta(z_i^t,t,y) \sim=z_i^0\quad\quad \mathcal f_\theta(z_i^t,t,y)为目标网络\ \mu(z_i^t,t,y)=\frac{1}{\bar{\beta_t}}( \sqrt{\alpha_t} \bar\beta_{t-1}z_i^t+\sqrt{\bar\alpha_{t-1}}\beta_t*\mathcal f_\theta(z_i^t,t,y) ) \end{align}$

经过类似上文130公式的推导，最终的Loss变成以下

$\begin{align} \mathcal L= \parallel z_i^0 -\mathcal f_\theta(z_i^t,t,y) \parallel^2 \end{align}$

在Prior最终输出量时还多了一步操作，Prior会生成2个 $z_i$ ，选取和 $z_t$ 最相似的作为Decoder的控制条件。注意 Prior模块不是必须的，从模型架构图上可以看到，我们也可以直接用文本向量 $z_t$ ，论文中发现这样生成图形的效果不够好；还有论文中为了加速diffusion的生成，采用DDIM跳跃采样方式来减少采样步数。

3.Decoder 模块

Decoder 模块就是上文提到的Classifier-free diffusion model，其中Condition 就是Prior模块的输出，这里不做赘述

4.SR超分辨率模块

从上文看到扩散模型要在输入图像的尺寸上计算，而且要经过非常多步的计算，所以在大图上速度非常慢，为了加快生成速度，一般会生成较小的图片，再通过超分辨率恢复成尺寸的较大的图片，目前SOTA的SR模型也基本都采用了diffusion 技术，DALL-E-2就采用两个基于diffusion的超分模型，先将第3步生成的64x64的图片恢复成256x256，再恢复成1024x1024。

基于diffusion的SR的也是使用Classifier-free diffusion，其中低分辨率的图像是控制条件y，高分辨率图像为x也就是我们最终希望生成的目标，所以模型变为 $p(x|y)$ ，每一个训练step如下图

三.ImaGen

ImaGen来着Google Rearch ，本质上也属于Classifier-free diffusion model，整体模型的架构图如下

和DALL-E-2非常像，大的模块上只是缺少了Prior模块，主要分为以下几部分

1.Text encoder。输出text embedding。

2.Text to Image Diffusion Model，这个作用和DALL-E-2的Decoder一致，采用Classifier-free diffusion来生成图像，其中控制条件y为第一步的text embedding。

3.两个级联的超分辨率图像模块，不同DALL-E-2的超分模块，ImaGen的超分模块输入除了低分辨率图像，还有text embedding。

下边将简单介绍这些模块

1.Text encoder

在ImaGen中，使用的文本编码器为T5-XXL，一个纯由自然语言预训练的模型，因为训练文本的数据规模和多样性会远远多于图文对数据的图片描述，所以在文本语义理解让要好于CLIP Text encoder，不过没有用了跨模态的训练，所以缺乏文本语义和图像语义的对齐，不过从最终生成的图像效果上看不会因为这点而低于DALL-E-2，所以在文本到图像生成任务中，预训练的文本编码器只要能很好的提取文本语义就可以，不用关心这个特征是否和图像特征可以对齐。

在这篇论文中还发现，如果采用预训练的文本编码模型，这个模型越强大最终生成图像的效果越好（这个很容易理解，两个语义不同的文本经过文本编码器的向量如果没法区分，最终生成的图像也肯定也是视觉语义上没法区分的），比如text encoder从T5-small到T5-XXL，生成图片人的主观评分在提高。这里简单介绍下T5。

T5 是google为了统一建模自然语言任务（比如以前的分开建模方式，自然语言的NER是预测目标字符的位置；分类是模型需要接入分类头；翻译是序列到序列任务），通过序列到序列生成的方式来统一建模。模型的输入是：任务描述（prompt）+ 输入文本，模型通过自回归的方式逐步解码生成答案，可以通过下图说明（左侧是模型输入、冒号前的是prompt；右侧是模型的生成输出）

这种方法除了可以统一建模自然语言的不同任务，还有一个优点可以统一自监督预训练和最终目标任务，从而zero-shot的效果会更好。

生成式的模型架构有以下三种（现在模型架构上基本每层都是用transformer)

1.最左侧，encoder-decoder，encoder可以看到全部的文本信息得到编码向量；decoder是自回归生成，所以训练时通过causal mask，让decoder只能看到当前时间步前边的文本信息和encoder的输出向量

2.中间的decoder，和左侧的decoder类似，只是输入缺少encoder的输出向量

3.PrefixLM ，混合了前两者，模型的部分层能看到全部文本，部分只能看到当前时间步前边的文本

最终T5选择的是第一种（GPT系列都是第二种，在生成的任务上效果会更好）

2.Text to Image Diffusion Model

ImaGen的图像生成模块采用Classifier-free diffusion model来生成图像，其中控制条件y为第一步的text embedding，这块不做过多赘述。

不过在这个环节有一个细节值得学习，在diffusion模型中，输入图像会被归一化到[-1,1]之间，不过在多步的扩散过程中，模型的输出特征会越来越超过范围，导致很难训练以及训练和最终采样生成的不一致，所以在ImaGen针对这个问题用了截断的处理方式，分别尝试了阈值和比例截断

1.阈值截断，简单的将模型每一个t时刻输出大于1缩放到1，小于-1缩放到-1，从而保证模型的输出一直是[-1,1]之间

2.比例截断，统计模型输出的像素值，然后选取一个确定值s，让[-s,s]范围内的像素满足一个确定比例，超过[-s,s]的阈值用-s和s截断，最终将截断后的值除以s整体归一化到[-1,1]。

经过效果分析，发现比例截断的效果会更好，也能容易理解原因，随着t变大，模型的输出在[-1,1]这个静态阈值外的像素会越来越多，将这些像素直接变成-1和1，会导致较多的细节损失。

3.级联超分辨率模块

Text to Image Diffusion Mode输出的图像大小是64X64，ImaGen使用两级的超分模块，将图像放大至1024x1024。的图像生成模块采用，超分模块也是采用Classifier-free diffusion model，不同于DALL-E-2的超分模块，控制条件y除了用低分辨率图片，还用了输入文本的编码向量也就是text encoder的输出。

论文中提到用了cross-attention的方式来融入了文本向量，个人猜测和Stable Diffusion是一致的，后文会提到

四.Latent Stable Diffusion Model

Latent Stable Diffusion Model (LDM)是Stability.AI 开发的，从模型权重和代码全部开源，基于这个开源的预训练模型已经有很多的创业产品如Lensa，个人像照片编辑工具

LDM在模型本质上也是采用了Classifier-free diffusion model，这个模型最显著两个特点

1.相比DALL-E-2和ImaGen直接在图像上做扩散，LDM是在一个隐空间上做的扩散，这个隐空间是由上文的VAE或者VQVAE类似模型学习得到的，因为隐特征的特征数远远小于原始的图像像素，这样就能极大的减少了多步扩散过程的计算量，所以模型可以运行在低算力的设备上。当然这么做也有一些仔细的思考，图像上大多数像素点对于语义信息是冗余的，而我们的可控扩散模型主要是控制图片语义，下边图片可以看到左边语义级别像素可以极大的压缩

2.相比其他的Classifier-free diffusion model，condition采用简单的相加或者concat融入到模型中，LDM采用了cross-attention做了融合，这样对于不同类型的控制条件y，模型可以更好捕获y和图像的相关性。

整个模型架构如下图

整个模型由3部分组成

1.最左侧的Latent编解码模型

通过上一节提到的AutoEncoder，将原始的图像x压缩到一个低维的隐特征 $\mathcal z$ ，同时得到编码器 $\mathcal \epsilon$ ，解码器 $\mathcal D$ ，为了避免的 $\mathcal z$ 方差太大扩散模型难以学习，所以对加了正则项进行约束，一种做法是上文VAE提到的KL散度（KL散度loss的权重不能太大）。还有一种是VQVAE，不同VAE假设的 $\mathcal z$ 符合高斯连续分布，VQVAE假设的 $\mathcal z$ 是一个确定的离散矩阵，模型架构如下

做简单的解释

1.期望Encoder编码出来的全部在图中紫色的Embedding Space，Embedding Space 是一个 $e_{K \times d}$ 的矩阵。

2.Decoder生成图片用的 $\mathcal z_q(x)_{ \in m\times m\times d}$ ，而不是Encoder编码得到 $\mathcal z_e(x)_{ \in m\times m\times d}$ 。 $\mathcal z_q(x)$ 全部来自 $e_{K \times d}$ 。具体的做法可以把 $\mathcal z_e(x)_{ \in m\times m\times d}$ 看成mxm个d维的向量，然后我们依次从 $e_{K \times d}$ 中找到和这mxm个距离最近的向量，找到这些向量就是复原成mxmxd的矩阵就是 $\mathcal z_q(x)_{ \in m\times m\times d}$ 。

3.在训练过程中除了图像的重建Loss，还要更新 $e_{K \times d}$ ，让 $\mathcal z_e(x)_{ \in m\times m\times d}$ 的每一维向量和 $e_{K \times d}$ 中离得最近的逐渐靠近。

2.最右边的控制条件conditioning模块

可以把任意的输入作为控制条件，如图片、语义图，文本等等，这些原始条件通过一个预训练的神经网络得到特征表示 $\tau_\theta(y)$ ，比如Stable Diffusion V1的text encoder用的是CLIP的text encoder，而Stable Diffusion V2 用的是OpenCLIP的text encoder。

3.中间的Diffusion Model 模块

这块采用的是一个Classifier-free diffusion，本质上和其他扩散模型一致，在正向扩散过程从 $\mathcal z$ 得到 $\mathcal z_T$ ，模型也是预测噪音 $\epsilon$ 。在反向过程生成隐特征 $\mathcal z$ ，再通过AutoEncoder 的解码器 $\mathcal D$ 得到图像。

和前边提到的DALL-E-2和ImaGen的的控制条件y做简单的相加融入不同，这里用了transformer 中的cross-attention进行融合，将y变化的特征 $\tau_\theta(y)$ 作为attention中的K和V，扩散每步的变换后特征作为attention中Q，用下图很容易看明白（context embedding 就是 $\tau_\theta(y)$ ，image就是 $\mathcal z$ ）

五.ControlNet

前边提到DALL-E-2，ImaGen，LDM等虽然可以通过文本等条件可以控制生成图像，但是因为文本语义的自由度很高，没法做的精细的控制，如线条、形状、图像深度等，基于这种缺陷最终新出了ControlNet。ControlNet 可以在原来的LDM模型上通过插件的方式引入额外的控制条件，只用微调级别的数据量就可以对生成图片进行精细的控制，如下图是利用边缘线精确的控制生成图像。

1.基本原理

ControlNet的基本原理如下图，引入类似残差的模块来学习新的控制条件。

左侧的训练好原始模型是冻结的、参数不会更新；右侧ControlNet中的trainable copy是左侧的模型和参数的拷贝，核心是两个zero conv，用的是weight、bias全部为0的1x1的卷积，第一个zero conv接受的是新的控制条件（通过神经网络或其他技术提取的控制特征，如边缘线，景深等）。模型开始训练时，因为zero conv的原因，右侧训练部分传递给y的的特征是0，相当于原始LDM模型的输出。

2.模型架构

基于上边的原理，对原始的LDM的改进后的架构图如下：