ControlNet-图像控制文生图模型

iResearch666

发布于 2023-09-13 14:17:18

7450

文章被收录于专栏：AI算法能力提高班AI算法能力提高班

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

paper https://arxiv.org/abs/2302.05543
code https://github.com/lllyasviel/ControlNet

ControlNet pose examples

Abstract

除了text-prompt文本提示作为输入，还增加了其它控制条件作为输入（文本+图像作为输入）
在小数据集上（<50k）鲁棒性较好
是一个特定任务条件，且是一个end2end，对diffusion model进行控制（Stable diffusion是一种）

Contributions

ControlNet 将diffusion model 克隆成2个模型，一个是trainable copy，一个是locked copy
- trainable copy 在特定任务上以端到端方式学习条件控制
- locked copy 保留网络来自数亿图像的生成能力，所以具有较好的鲁棒性
- 2个copy的联系采用 zero convolution ，其权重从0开始增加到优化参数，因为不需要向深层特征增加noise，所以训练速度和微调模型一样快
Conditions: Canny edges, Hough lines, user scribbles, human key points, segmentation maps, shape normals, depths, etc

Related Work

Text-to-Image Diffusion

Diffusion models can be applied to text-to-image generating tasks to achieve state-of-the-art image generating results. This is often achieved by encoding text inputs into latent vectors using pretrained language models like CLIP

Glide is a text-guided diffusion models supporting both image generating and editing.
Disco Diffusion is a clip-guided implementation of to process text prompts.
Stable Diffusion is a large scale implementation of latent diffusion to achieve text-to-image generation.
Imagen is a text-to-image structure that does not use latent images and directly diffuse pixels using a pyramid structure.

Image-to-Image Translation

在应用上，ControlNet 和 image-to-image translation 存在重叠，但是二者动机不同
- image-to-image translation 是学习不同domain图像之间的映射关系
- ControlNet 用特定任务条件对diffusion model 进行控制

Methodology

Overview

image-20230828105432263

克隆预训练模型：ControlNet 首先创建了预训练图像扩散模型的两个副本，其中一个是 “锁定” 的，不能被修改，而另一个是 “可训练” 的，可以在特定任务上进行微调。 ControlNet 使用了一种称为 “权重共享” 的技术，该技术可以将预训练模型的权重复制到两个不同的神经网络中。这样，在微调可训练副本时，锁定副本仍然保留着从预训练中学到的通用知识，并且可以提供更好的初始状态。需要注意的是，在克隆预训练模型之前，需要先选择一个适合特定任务的预训练模型，并对其进行必要的调整和优化。这样才能确保克隆出来的模型能够更好地适应特定任务，并取得更好的效果。
定义输入条件：ControlNet 然后定义了一组输入条件，可以用来控制模型的输出。这些条件可能包括颜色方案、对象类别或其他特定任务参数。该技术可以将输入条件与预训练模型进行连接，并将其作为额外的输入信息传递给神经网络。这样，在微调可训练副本时，神经网络可以根据这些输入条件来调整输出结果，并更好地适应特定任务。需要注意的是，在定义输入条件之前，需要先确定哪些参数对于特定任务是最重要的，并选择合适的方式来表示这些参数。例如，在图像分类任务中，可能需要考虑对象类别、颜色和纹理等因素，并选择合适的方式来表示它们。同时，还需要确保这些输入条件能够被有效地传递给神经网络，并且不会对模型性能产生负面影响。
训练可训练副本：ControlNet 然后使用反向传播和其他标准训练技术，在特定数据集上对可训练副本进行训练。很自然的就是一个标准的神经网络训练流程。训练可训练副本之前，需要先选择一个适合特定任务的预训练模型，并将其克隆为可训练副本。同时还需要确定哪些输入条件对于特定任务是最重要的，并将其与预训练模型进行连接。
合并输出：最后，ControlNet 将两个模型副本的输出组合起来，产生一个最终结果，既包含从预训练中学到的通用知识，也包含从微调中学到的特定知识。

数据流转主要流程如下：

a.timesteps经过embedding转换为特征向量送入Stable Diffusion和ControlNet；

b.随机噪声被送入Stable Diffusion；

c.图像的Map经过HintBlock，与随机噪声相加，送入ControlNet;

d.Prompt的Embedding送入Stable Diffusion和ControlNet；

e.Stable Diffusion的所有参数被冻结不参与训练

f.ControlNet的结构与Stable Diffusion一致，只是每层后面增加一个零卷积；

g.Stable Diffusion和ControlNet中的ResBlock将上一层的输出和timesteps作为输入；

h.Stable Diffusion和ControlNet中的SpatialTransformer将上一层的输出和Prompt Embedding 作为输入。

ControlNet

为什么要将diffusion model分成2个copy，而不是在原始模型上训练权重？

因为一方面要保留采用数亿图像训练好模型的生成质量，另一方面在特定任务上训练，其数据量少容易overfitting

2个copy通过zero convolution连接零卷积是指卷积的权重和偏差初始化为0的1X1普通卷积，在迭代过程中参数更新在第一次迭代时候，零卷积为0，对2个copy网络贡献为0，所以不影响网络性能，网络输入仍为x，输出仍为y 。x + Zc x Zin -> y + Zy x Zout 满足的条件：特征图像非零

image-20230828142618638

ControlNet in Image Diffusion Model

以Stable Diffusion为例，它是由encoder，a middle block，a skip-connected decoder三部分组成。其中U-net encoder和decoder都是12个blocks，总共25个blocks。有8个block是下采样或者上采样，其它17block包含4个resnet层和2个ViT层。每个ViT层包括self or cross注意力机制。Text encoder采用OpenAI CLIP，每个time steps采用positional encoding。
SD是将512X512分辨率的图像转化到64X64分辨率的 latent images，所以ControlNet需要将图像空间编码到64X64的隐图像空间Cf，设计一个卷积网络即可。 We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides (activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full model)
训练时候内存增加了23%，时间增加了34%，计算效率比较高
trainable copy中，只是encoder和middle block参数可调，其它decoder参数固定