ControlNet Adding Conditional Control to Text-to-Image Diffusion Models

news2024/10/6 14:32:40

ControlNet: Adding Conditional Control to Text-to-Image Diffusion Models

TL; DR：ControlNet 使得我们能通过输入额外的条件图（如 Canny 边缘、人体姿态、深度图等），对 SD 生成结果的空间位置有更准确的控制。它拷贝 SD 部分原权重作为一个新的分支进行微调训练，同时维持 SD 原权重分支不变，并在开始训练使用零卷积将二者相连。从而能利用已经在大规模图片数据上预训练的 SD 权重，尽量保持原模型能力，且训练高效。

方法

模型结构

ControlNet 对 SD 原模型的 adapt 如图所示。图中 x 是输入噪声图，y 是输出，c 是条件图输入。拷贝一份原参数权重作为另一分支，进行训练，原参数权重分支不动，两分支之间用零卷积链接。所谓零卷积，就是权重和偏置都初始化为 0 的 1x1 卷积。这样，在训练刚开始时，由于零卷积的输出都是 0，因此改动后的模型和原 SD 模型的输出是完全一致的。随着训练的进行，才会产生变化。这样的结构能保证不会对预训练的 SD 模型带来有害的噪声。

在这里插入图片描述

文章以 SD txt2img 为例，详细介绍了插入 ControlNet 的方法。原 SD 的 UNet 网络共有 12 层降采样，12 层上采样和 1 层中间层。其特征图共有 4 中空间分辨率大小（64、32、16、8）。ControlNet 对 UNet 的 encoder 部分（12 层降采样 + 1 层中间层）构建了可训练拷贝。其输出分别经过零卷积，输入给原 SD 模型的解码器部分。

在这里插入图片描述

此外，SD 是隐层扩散模型（LDM），其扩散过程时发生在 latent 空间的。输入的条件图自然也需要转换到 latent 空间，ControlNet 是使用了四层卷积层完成了这个转换，将 512x512 的真实图片转换到 64x64 的 latent 空间。

训练

ControlNet 是在预训练好的模型权重上，加入一些参数，然后进行微调。其训练的目标函数与 SD 是完全一致的，但是加入了各种条件图（如 Canny、人体姿势等）作为输入。在训练时 ControlNet 时，还会有 50% 的概率将文本 prompt 替换为空字符串，从而训练 ControlNet 直接识别条件图语义的能力。

作者在训练 ControlNet 时还观察到一个奇特的现象，即模型并不是逐渐学习到条件控制能力的，而是在某一步，突然就会了。如下图，模型在 6133 步，突然就能精准地按照输入的 Canny 边缘进行生成了。作者称这为 “突然收敛现象” （sudden convergence phenomenon）。

在这里插入图片描述

推理

训练结束之后，推理时就可以按照输入的条件图来控制空间位置，进行生成了。这里作者还交代了几点细节，用于更好地控制生成结果。

CFG-RW：SD 在训练时使用了 Classifier-Free Guidance 的方法，即同时训练模型条件生成和无条件生成的能力，然后在生成时，也同时生成条件结果和无条件结果，并根据指定的 CFG Scale 得到最终的生成结果： $\epsilon_{\text{prd}}=\epsilon_{\text{uc}}+\beta_{\text{cfg}}(\epsilon_{\text{c}}-\epsilon_{\text{uc}})$ 。ControlNet 中，条件生成的结果既可以加到 $\epsilon_{\text{uc}}$ 和 $\epsilon_{\text{c}}$ 中，也可以只添加到 $\epsilon_{\text{c}}$ 中。如果 prompt 为空（有 50% 概率），此时如果同时加到 $\epsilon_{\text{uc}}$ 和 $\epsilon_{\text{c}}$ 中，则相当于没有 CFG 了，而如果只加到 $\epsilon_{\text{c}}$ 中，那么 CFG 引导又太强。这里作者的策略是使用分辨率加权（CFG-RW），即 $w_i=64/h_i$ 其中 $h_i$ 是第 i 层的空间分辨率，如 $h_1=8,h_2=16,\dots$ 。