stable diffusion(五)——ControlNet原理与实践


这个系列会分享下stable diffusion中比较常用的几种训练方式,分别是Dreambooth、textual inversion、LORA、Hypernetworks和ControlNet。在https://civitai.com/选择模型时也能看到它们的身影。本文该系列的第五篇

前言

论文链接:https://arxiv.53yu.com/abs/2302.05543

代码链接:https://github.com/lllyasviel/ControlNet

ControlNet可以从提供的参考图中获取布局或者姿势等信息, 并引导diffusion model生成和参考图类似的图片。这种功能对于生成模型来说至关重要, ControlNet开源到stable diffusion中时, 很多自媒体文章的标题都是《XX行业要结束了》, 有点杀死比赛的意思。训练和测试过stable diffusion的朋友应该知道, prompt至关重要, 而调整prompt确实是一件费时费力的事情, 想要控制好一些细节非常不容易。往往都是生成大量图片, 能得到几个满意的。而有了ControlNet则可以精确的控制生成图片的细节(当然, 也依赖于你已经有一张参考图了)

扩散模型中会需要使用text作为条件引导模型生成, 而ControlNet在引导的时候会增加一个条件实现更加可控的text-to-image生成.

diffusion model进展非常迅速, 针对diffusion model的一些痛点, 作者提出了几个问题和训练扩散模型时的限制, 并着手用ControlNet解决:

  • 基于prompt的控制是否能满足我们的需求?
  • 在图像处理中, 有许多有明确定义的长期任务, 是否可以用大模型来使这些任务达到新的高度?
  • 应该用什么样的框架处理用户想要的各种各样条件?
  • 在特定场景, 新提出的框架怎么能继承大模型(diffusion model)的能力.

三点限制:

  • 特定场景的数据量少, 在较少的数据量上很容易出现过拟合
  • 大的计算集群很奢侈, (stable diffusion base模型需要15万个A100显卡小时)
  • end2end是必要的, 一些特定场景需要将原始的输入转为更高语义的表达, 手工处理的方式不太可行

Method

对于上面的4个问题和3点限制, 作者提出了ControlNet, ControlNet具有以下的特点:

  • end2end
  • control large im age diffusion model(例如stable diffusion)
  • learn task-specific input conditions

为了能继承大模型的能力, ControlNet将原来大模型的参数分为两部分, 分别是可训练冻结的。对比下添加ControlNet结构前后的网络图很容易明白。

之前stable diffusion的模型结构如下图所示:

图(a)是之前stable diffusion的输出, 图(b)和图(a)的区别在于添加了ControlNet的结构, 具体而言是将neural network block复制了一份, 作为trainable copy,并且neural network block的网络参数会被冻结住。而且trainable copy前后会有zero convolution, zero convolution其实是1*1的卷积。最后会将neural network blocktrainable copy的特征相加。

neural network block的作用可以用下面的公式表示

$x$是一个2D的feature, 其shape为{h,w,c}分别代表height, width, channel.$\theta$代表neural network block的参数, y表示经过neural network block后的特征。

添加了control后, 网络会从公式1变为下面的公式2

$Z$代表zero convolution, $\theta_{z1}$,$\theta_{z2}$,$\theta_{c}$分别代表第一个, 第二个zero convolution的参数和trainable copy的参数。

在初始化训练的时候, zero convolution的weight和bias会被初始化为0, 因此ControlNet具有以下的特点:

实际上是为了让初始化的时候添加了ControlNet的模型和原始的模型效果是一样的, 后续再慢慢finetune

可以简化成下面的表达式

$y_c$是加上ControlNet的输出, 而$y$是原始的输出, 模型初始化时两者是相等的

文章中特地解释了zero convolution梯度更新的问题, 在模型开始迭代前, 有W=0, B=0.则

可以看到权重W的梯度是I, 而偏置B的梯度是1.说明zero convolution是可以正常更新的

ControlNet的详细结构如下图所示:

灰色部分是原来stable diffusion的结构, 蓝色部分是从U-Net的encode对应部分copy, 经过zero convolution后和U-Net的decode相加。

了解了结构后, 下面是训练的细节。$\epsilon-\epsilon_{\theta}$和原始的stable diffusion是一样的, 是每次生成噪声和实际噪声的差值。区别在于加入了ControlNet新加的控制条件$c_f$。其他的参数都和stable diffusion一样, $z_t$代表噪声图片, $t$代表迭代次数,$c_t$代表text prompt。

在实际训练的时候, 会随机让50%的text prompt变为空, 促使模型更多的关注$c_f$,让模型学习用$c_f$控制图像生成。

虽然ControlNet新加了些模型结构, 但由于大部分参数都被冻结, 因此训练时实际上只用原来stable diffusion模型的23%显存和34%的训练时间, 可以在单张40GB的A100卡上做训练

计算资源的限制: 如果数据量比较少, 可以只连接U-Net中的SD Middle Block, 而连接decoder的zero convolution都不用, 可以减少参数量。可以将ControlNet的训练速度提高1.6倍(在RTX 3070TI的消费级显卡上也可以做训练)。

有大量的训练资源: 如果有大量的训练资源(8张A100), 为了获取更好的效果可以将原始U-Net模型中的被冻结参数的decoder部分也参与模型训练。

消融实验

这部分没有定量的分析, 更多的是通过一些case表明ControlNet能有效的控制生成的图像, 因为毕竟新加了个条件, 没有现成的数据集用来评测。

有无ControlNet结构对比

下面的图最左边是ControlNet使用的canny图;中间是没有controNet, 只通过text prompt控制生成的房子及周边环境;右边是使用了ControlNet的图, 可以看到不实用ControlNet, 房子的结构可能会发生一些改变, 而加上了ControlNet后, 房子的结构始终固定, 改变只有房子的外观和周边的环境。

突然收敛现象

作者发现ControlNet在训练的时候会出现突然收敛, 这种收敛一般发生在5000-10000步, 而且收敛之后如果继续迭代模型的变化并不大。例如下面的图, 提供了苹果的canny图和text prompt: apple, 但模型最开始学习到的是苹果电脑相关的, 到6100步图片中仍然没有学习到生成提供的canny图类似的图片, 但仅仅过了33步后到6133步模型突然就学会了!而且该现象和batch size关系不大, 也就是说如果增大batch size 模型差不多也会在6133步收敛, 好处是大batch size收敛出的模型效果更好。

canny-edge-based模型和训练数据大小的关系

在50k数据之前增加数据效果会有明显的提升, 但50k之后狮子生成的已经很好了, 差异在于增大数据后背景会更加好点。有点像突然收敛之后继续迭代的情况。

和其他用图像作为条件引导生成方式的对比

depth

在stable diffusion V2.0中有个很重要的更新——Depth-to-Image Diffusion, 可以推断输入图像的深度信息, 然后利用文本和深度信息生成新图像。

这种方式和depth-base ControlNet很像, 对比如下图所示:

从生成效果上两者看起来都非常好(硬要比的话, ControlNet效果稍好些)。从训练成本上, stable diffusion V2.0需要大量A100的训练集群在1200万的训练数据训练2000GPU-hours, 而Depth-based ControlNet在stable diffusion V1.5的基础上预训练, 使用单张RTX 3090TI在20万数据上训练1周。这么对比确实ControlNet太香了

segmentation

下面的图对比的是PITI(Pretraining-Image-to-Image)和segmentation-base ControlNet

Scribbles

下面的图对比的是Sketch-guided diffusion和scribbles-base ControlNet

canney

下面的图对比的是Taming Transformer和canney-base ControlNet

控制条件对比

不同控制条件下生成效果的对比, 作者给了一个房间布局的图, 然后分别用cannry, depth, HED, Nromal, Line和Scribbles的方式做对比。其实这几个模型都有更适合的应用场景, 关以房间布局的例子而言, Line和HED是更合适的, 这个从Line和HED的图像特征就能看出。

其他

更多的应用,下面的图展示的是把原图的一部分mask掉, 然后随便画几笔草图, 生成的图就是没mask掉的部分结合草图的完成效果。(太牛了)

当图像细节比较简单时, ControlNet可以实现精确的控制

如果需要兼顾和提示图片类似与多样性两种需求, 可以在生成图片中灵活的设置ControlNet参与图片生成的步数。例如下面的图, 和原始图相似, 但也保证了多样性

局限性

ControlNet可能会曲解canney和depth等图像特征的含义, 例如下面的图是一个水杯的canney图, ControlNet生成的是一个农田, 加上a glass of water后, 生成的图片稍微好点, 但仍然和实际的有差异。

总结

ControlNet的原理到这里就介绍完了, 做的改进很简洁, 只需要较小的数据集就可以用作训练(50k, 相比diffusion model亿级别的确实小不少)。用的计算资源较少, 并且做一些参数调整后甚至可以在消费级的显卡上做训练。总共有10个不同图像特征的模型, 可以适应多种生成需求。ControlNet在stable diffusion中已经开源了, 并且在webui上很容易添加和尝试,这部分有非常多的blog和视频有讲解, 就不再赘述了。

-------------本文结束感谢您的阅读-------------