51-50 两万字长文解读ControlNet论文和代码，以及自定义模型训练和图片精确控制生成实验

今天咱们来看 ICCV2023 最佳论文Adding Conditional Control to Text-to-Image Diffusion Models，又称为ControlNet。提到图像生成Finetuning工程方法，有Textual inversion、DreamBooth、LoRA、T2I-Adapter以及ControlNet，其中最著名的当属ControlNet。它也是Stable Diffusion核心插件，业内把ControlNet 称为精确控制 AI 图像生成的破冰方案。目前文生图经典的做法是集成Stable Diffusion、LoRA、ControlNet一起使用。

本文第一部分，论文精读，重点内容做了备注、解释。
本文第二部分，模型训练。自定义“数据集”进行模型训练，譬如训练智驾场景文生图。
本文第三部分，代码讲解。重点讲了LDM的encoder即ControlNet的实现代码、LDM参数被复制代码、零卷积和权重初始化为零的巧妙设计。
本文第四部分，Stable Diffusion+LoRA+ControlNet精确控制图片生成实验。

自载模型可以不受限制地按照您的意图生成图片，是否蛮期待？

第一部分论文精读

Abstract

我们提出了 ControlNet，这是一种神经网络架构，用于将空间条件控制添加到大型预训练的文本到图像扩散模型中。ControlNet锁定可生产的大型扩散模型，并将其深度和稳健的编码层重新用作学习一组不同条件控制的强大骨干，这些编码层是用数十亿张图像预训练的。神经架构与“零卷积”（零初始化卷积层）连接，后者从零逐步增长参数，并确保没有有害的噪声会影响微调。我们在单一/多个条件，有/没有提示情形下测试各种条件控制Stable Diffusion，例如边缘、深度、分割和人体姿势等。ControlNets 的训练对于小型 (<50k) 和大型 (>1m) 数据集是稳健的。大量的结果表明，ControlNet可以促进控制图像扩散模型的更广泛应用。

1. Introduction

我们中的许多人都经历过视觉灵感的闪现，我们希望用独特的图像捕捉这些灵感。随着文本到图像扩散模型的出现，我们现在可以通过输入文本提示来创建视觉上令人惊叹的图像。然而，文本到图像模型在控制图像的空间组成方面是有限的；仅仅通过文本提示来精确表达复杂的布局、姿势、形状和形式是很困难的。生成一个与我们的心理意象准确匹配的图像通常需要无数次的试错循环，如编辑提示，检查生成的图像，然后重新编辑提示。

我们能否通过让用户提供额外的图像来启用更细粒度的空间控制，直接指定他们想要的图像组合？在计算机视觉和机器学习中，这些额外的图像（例如边缘图、人体姿势骨架、分割图、深度特征图、法线等）通常被视为图像生成过程的条件。

### 备注 1，目前ControlNet实现的功能蛮多，如下图，集成到Stable Diffusion的ControlNet功能示意图。

图像到图像转换模型学习从条件图像到目标图像的映射。研究界还采取了一些措施来控制文本-图像模型，包括空间掩码、图像编辑指令、通过微调进行个性化等。虽然一些问题（例如，生成图像变体、修复）可以通过无训练的技术来解决，如约束去噪扩散过程或编辑注意力层，但更广泛的问题，如深度到图像、姿势到图像等，需要端到端学习和数据驱动的解决方案。

以端到端的方式学习大型文本到图像扩散模型的条件控制具有挑战性。特定条件的训练数据数量可能明显小于通用文本到图像训练数据。例如，各种特定问题的最大数据集通常约为 100K，比用于训练Stable Diffusion 的 LAION-5B 数据集小 50,000 倍。

在有限的数据下对大型预训练模型进行直接微调或持续训练可能会导致过拟合和灾难性遗忘。

研究人员表明，可以通过限制可训练参数的数量或排名来缓解这种遗忘。

### 备注 2：

Vision Transformer Adapter for Dense Predictions
International Conference on Learning Representations (ICLR)
LoRA: Low-Rank Adaptation of Large Language Models
Side-Tuning: Network Adaptation via Additive Side Networks

对于我们的问题，设计更深或定制的神经架构，处理具有复杂形状、不同高级语义、开放条件下的图像是必要的。

本文介绍了 ControlNet，这是一种端到端的神经网络架构，它学习大型预训练文本到图像扩散模型的条件控制（在我们的实现中是可扩散的）。ControlNet 通过锁定其参数来保留大型模型的质量和功能，并复制了一份可训练的编码层。该架构将大型预训练模型视为学习不同条件控制的强大主干。可训练副本和原始锁定模型通过零卷积层连接，权重被初始化为零，以便它们在训练期间逐渐增长。这种架构确保在训练开始时不会将有害的噪声添加到大型扩散模型的深度特征中，并保护可训练副本中的大规模预训练主干不被此类噪声破坏。

##### 备注 3：

ControlNet复制了一份Stable Diffusion可训练的编码层Encoder。训练时Stable Diffusion的编码层、中间层、解码层都是锁定的。
通过巧妙的设计，即可训练Encoder副本和原始锁定模型Stable Diffusion通过零卷积层连接，权重被初始化为零，这样既可以保留从大型数据集中学到的大量先验知识，又可以定制训练，以精确生成特定任务需要的图片。
零卷积可以确保噪声不会添加到初始训练副本中；权重初始化为零，但可以训练。
参考代码发现，ControlNet核心代码有完整的LDM代码，复制了一份LDM的encoder、middle，放在cldm。

实验表明，ControlNet 可以通过各种条件输入来控制Stable Diffusion，包括 Canny edges, Hough lines, user scribbles, human key points, segmentation maps, shape normals, depths等等（如图 1）。

本文贡献：

提出了ControlNet，这是一种神经网络架构，可以通过有效的微调将空间局部化的输入条件添加到预训练文本到图像扩散模型中。
提出了预训练 ControlNets 来控制Stable Diffusion，以Canny edges, Hough lines, user scribbles, human key points, segmentation maps, shape normals, depths, and cartoon line drawings等等为条件。（数了一下，当前版本有20种控制方法）
我们与几种替代架构进行了比较，通过消融实验验证该方法，并针对不同任务的几个先前基线进行了用户研究。

2. Related Work

2.1. Finetuning Neural Networks

微调神经网络的一种方法是通过额外的训练数据直接继续训练它。但是这种方法会导致过度拟合、模式崩溃和灾难性遗忘。广泛的研究集中在制定避免此类问题的微调策略上。

HyperNetwork

是一种起源于自然语言处理社区的方法，目的是训练一个小的循环神经网络来影响更大神经网络的权重。它已被应用于生成对抗网络GAN的图像生成。HyperNetwork style trainning 为Stable Diffusion实现了超网络，以改变其输出图像的艺术风格。

Adapter

在自然语言处理中广泛应用，通过嵌入新模块层来定制预训练transformer模型以适应其他任务。在计算机视觉中，适配器用于增量学习和域适应。该技术通常与CLIP一起使用，用于将预训练的骨干模型转移到不同的任务。最近，适配器在vision transformers和VIT-Adapter中取得了成功。在与我们同时进行的工作中，T2I-Adapter使Stable Diffusion适应外部条件。

Additive Learning

通过冻结原始模型权重并使用学习的权重掩码、修剪或hard attention添加少量新参数来规避遗忘。Side-Tuning使用侧分支模型通过线性混合冻结模型和添加网络的输出来学习额外的功能，并使用预定义的混合权重时间表。

Low-Rank Adaptation (LoRA)

通通过使用低秩矩阵学习参数的偏移量，防止灾难性遗忘，这是基于许多过度参数化模型位于低维子空间的观察结果。

Zero-Initialized Layers

ControlNet使用零初始化层来连接网络块。对神经网络的研究已经广泛讨论了网络权重的初始化和操作。例如，权值的高斯初始化可能比用零初始化的风险更小。最近，IDDPM讨论了如何在扩散模型中缩放卷积层的初始权重以改进训练，并且它们的“零模块”实现是将权重缩放到零的极端情况。Stability's model cards也提到了在神经层中使用零权重。ProGAN、StyleGAN和Noise2Noise中也讨论了操纵初始卷积权值。

### 备注 4：作者总结了文生图微调方法，如HyperNetwork、Adapter、LoRA经典方法等。咱们也可以通过以下模型了解到一些实用工程方法。

Textual inversion，简易的实现控制。
DreamBooth，完整的跑通了微调全流程，个人认为属于经典论文。
LoRA，灵巧的实现，应用较广。
ControlNet，彻底的实现精确控制，支持的条件最多。
T2I-Adapter，腾讯出品，出现比ControlNet晚一点，效果也差一点，光芒被遮挡。

2.2. Image Diffusion

Image Diffusion Models

SohlDickstein等人首先介绍Image Diffusion Models，最近已应用于图像生成。潜在扩散模型LDM在潜像空间中执行扩散步骤，降低了计算成本。文本到图像扩散模型通过CLIP等预训练语言模型将文本输入编码为潜在向量，实现了最先进的图像生成结果。GLIDE是一个文本引导的扩散模型，支持图像生成和编辑。Disco Diffusion使用剪辑引导处理文本提示。Stable Diffusion是潜在扩散的大规模实现。Imagen直接使用金字塔结构扩散像素，而不使用潜在图像。商业产品包括DALL-E2和Midjourney。

Controlling Image Diffusion Models

控制图像扩散模型有助于个性化、自定义或特定任务的图像生成。图像扩散过程直接提供了对颜色变化和修复的一些控制。文本引导控制方法侧重于调整提示、操纵 CLIP 特征和修改交叉注意力。MakeAScene 将分割掩码编码为令牌以控制图像生成。SpaText将分割掩码映射到本地化令牌嵌入中。GLIGEN在扩散模型的注意力层中学习新参数以进行grounded生成。Textual Inversion和DreamBooth可以通过使用一小组用户提供的示例图像微调图像扩散模型来个性化生成图像中的内容。基于提示的图像编辑提供了用提示操作图像的实用工具。Sketch-Guided Text-to-lmage Diffusion Models出了一种将扩散过程与草图拟合的优化方法。

并行工作MultiDiffusion、Masksketch、Composer、T2I-Adapter研究了控制扩散模型的各种方法。

2.3. Image-to-Image Translation

条件GAN和transformer可以学习不同图像域之间的映射，例如Taming Transformer是一种vision transformer方法；Palette是从头开始训练的条件扩散模型；PITI是一种基于预训练的图像到图像转换条件扩散模型。操纵预训练的 GAN 可以处理特定的图像到图像任务，例如 StyleGAN 可以通过额外的编码器来控制。

3. Method

图 2：神经块将一个特征图 x 作为输入，并输出另一个特征映射 y，如（A）所示。为了将ControlNet添加到这样的块中，我们锁定原始块并创建可训练的副本，并使用零卷积层将它们连接在一起，即权重和偏差都初始化为零的1×1卷积。这里，c 是我们希望添加到网络中的条件向量，如（b）所示。

### 备注 5

ControlNet的功效相当于一种辅助神经网络，在原Stable Diffusion模型中添加了一个辅助的可训练模块，从而引入额外条件c来控制图像生成过程。
原SD模型权重被复制了一份，给最右边ControlNet模型训练时使用trainable copy，同时原SD模型权重被锁定Locked，不参与训练。
ControlNet模型结果和原SD模型结果 $y$ 相加add获得最终输出结果 $y_c$ 。

3.1. ControlNet

ControlNet将附加条件注入到神经网络的块中(图 2)。在这里，我们使用术语网络块来指代一组通常组合在一起形成神经网络的单个单元的神经层，例如resnet block, conv-bn-relu block, multi-head attention block, transformer block等。假设F(·;Θ)就是这样一个经过训练的神经块，参数为Θ，将输入的特征映射x转换为另一个特征映射y为

在我们的设置中，x 和 y 通常是 2D 特征图，即 x ∈Rh×w×c，其中 {h, w, c} 分别是映射中通道的高度、宽度和数量（图 2a）。为了将ControlNet添加到这样一个预先训练的神经块中，我们锁定原始块的参数Θ，同时将块克隆到参数为Θc的可训练副本(图2b)。可训练副本以外部条件向量 c 作为输入。当这种结构应用于Stable Diffusion等大型模型时，锁定参数保留了经过数十亿张图像训练的生产就绪模型性能，而可训练副本重用这种大规模预训练模型，以建立一个深度、鲁棒和强大的主干来处理不同的输入条件。可训练副本连接到具有零卷积层的锁定模型，表示为 Z(·; ·)。具体来说，Z(·;·) 是一个 1×1 的卷积层，权重和偏差都初始化为 0。

为了构建 ControlNet，我们分别使用两个参数为 Θz1 和 Θz2 的零卷积实例。然后计算完整的ControlNet，其中 $y_c$ 是ControlNet块的输出。

在第一个训练步骤中，由于零卷积层的权重和偏差参数都被初始化为零，因此等式 (2) 中的 Z(·;·) 项都评估为零。

因此， $y_c= y$ ，这样，当训练开始时，有害噪声不会影响可训练副本中神经网络层的隐藏状态。此外，由于 Z(c; Θz1) = 0，可训练副本也接收输入图像 x，可训练副本完全具备功能性，保留了大型预训练模型的能力，使其能够作为进一步学习的强大主干。零卷积通过在初始训练步骤中消除作为梯度的随机噪声来保护该主干。

### 备注 6

Stable Diffusion模型被Locked，保留原本能力；与此同时，使用额外数据对可训练副本trainable copy进行微调，学习想要添加条件如边缘、深度、分割和人体姿势控制等。这样ControlNet使用小批量数据集就能对控制条件进行训练，同时不会破坏原SD模型经过数十亿张图像训练后的能力。
不直接在原始权重上训练，主要是为了规避在小数据集上训练容易过拟合。
在训练步骤的第一步中，零卷积层权重w和偏差参数b都被初始化为零。首先确保了有害噪声不会在训练之初就影响可训练副本中神经网络层的隐藏状态。其次通过多次迭代，可以不断优化可训练参数权重。

3.2. ControlNet for Text-to-Image Diffusion

我们使用Stable Diffusion作为一个例子来展示ControlNet如何将条件控制添加到大型预训练扩散模型中。Stable Diffusion本质上是一个U-Net网络，具有编码器、中间块和跳接解码器。编码器和解码器都包含 12 个块。完整模型包含 25 个块。在 25 个块中，8 个块是下采样或上采样卷积层，而其他 17 个块是每个块包含 4 个 ResNet 层和 2 个 Vision Transformers (ViT) 的主要块。每个ViT都包含几个交叉注意和自注意。例如，在图 3a 中，“SD Encoder Block A”包含 4 个 ResNet 层和 2 个 ViT，而“×3”表示该块重复 3 次。文本提示使用CLIP文本编码器进行编码，扩散时间步长使用位置编码的时间编码器进行编码。

图 3：Stable Diffusion U-net架构在编码器块和中间块上连接了一个ControlNet。锁定的灰色块显示了Stable Diffusion V1.5结构。添加可训练的蓝色块和白色零卷积层来构建ControlNet。

ControlNet 结构应用于 U-net 的每个编码器层级（图 3b）。特别是，使用 ControlNet 创建了 12 个编码块和 1 个稳定扩散中间块的可训练副本。12个编码块分辨率为4个(64 × 64， 32 × 32， 16 × 16， 8 × 8)，每个编码块复制3次。输出被添加到 U-net 的 12 个跳接和 1 个中间块中。由于Stable Diffusion是一种典型的 U-net 结构，因此该 ControlNet 架构可能适用于其他模型。我们连接 ControlNet 的方式在计算上是有效的——因为锁定的副本参数被冻结，所以最初锁定的编码器不需要梯度计算进行微调。这种方法加快了训练速度并节省了 GPU 内存。在单个 NVIDIA A100 PCIE 40GB 上进行测试，与没有ControlNet的情况下优化Stable Diffusion相比，使用ControlNet优化Stable Diffusion在每次训练迭代中只需要增加约23%的GPU内存和34%的时间。

图像扩散模型学习逐步去噪图像并从训练域生成样本。去噪过程可以在像素空间或从训练数据编码的潜在空间中发生。Stable Diffusion使用潜在图像作为训练域，因为在这个空间中工作已被证明可以稳定训练过程。具体来说，Stable Diffusion 使用类似于 VQ-GAN 的预处理方法将 512 × 512 像素空间图像转换为更小的 64 × 64 潜在图像。为了将ControlNet添加到Stable Diffusion中，我们首先将每个输入条件图像(如边缘、姿势、深度等)从512 × 512的输入大小转换为与Stable Diffusion大小相匹配的64 × 64特征空间向量。特别是，我们使用四个卷积层的微小网络 $\varepsilon$ (·)，具有 4 × 4 内核和 2 × 2 步幅（分别由 ReLU 激活，分别使用 16、32、64、128 通道，用高斯权重初始化并与完整模型联合训练）将图像空间条件 $c_i$ 编码为特征空间条件向量 $c_f$

条件向量 $c_f$ 被传递到 ControlNet 中。

### 备注 7 ：

Stable Diffusion使用潜在图像LDM作为基础模型。SD参数被冻结，不需要进行梯度计算，相对来说比较节省GPU。
整体架构是UNet主干，在Unet骨干的部分层引入了Vision Transformer。ControlNet主要将SD UNet的Encoder、Middle部分进行复制训练，然后把结果输入到SD Decoder模块中，通过zero convolution模块、skip connection加入处理后的特征。
从WebUI可以看到整个模型输入包括原图Map Input、Prompt、Added Prompt、Negative Prompt、Random Input。更细节一点，ControlNet输入包括Time Embedding、Text Embedding，condition $c_f$ 。

3.3. Training

给定一个输入图像 $z_0$ ，图像扩散算法逐步向图像添加噪声并产生噪声图像 $z_t$ ，其中 t 表示添加噪声的次数。给定一组条件，包括时间步长 t、文本提示 $c_t$ 和特定任务条件 $c_f$ ，图像扩散算法学习一个网络 εθ 来预测添加到噪声图像 $z_t$ 中的噪声，其中

其中 L 是扩散模型的整体学习目标。该学习目标直接用于使用 ControlNet 微调扩散模型。在训练过程中，我们将 50% 的文本提示 $c_t$ 随机替换为空字符串。这种方法增加了 ControlNet 在输入条件图像（例如边缘、姿势、深度等）中直接识别语义的能力，作为提示的替代品。在训练过程中，由于零卷积不会向网络添加噪声，因此模型应该始终能够预测高质量的图像。

### 随机替换，训练时将 50% 的文本提示 $c_t$ 随机替换为空字符串。

我们观察到该模型不会逐渐学习控制条件，而是在输入条件图像之后突然成功；通常在不到 10K 优化步骤中。如图 4 所示，我们称之为“突然收敛现象”。

3.4. Improved Training

Small-Scale Training

在计算设备有限的情况下，我们发现部分断ControlNet与SD之间的连接可以加快收敛速度。默认情况下，我们将ControlNet连接到“SD中间块”和“SD解码器块1,2,3,4”，如图3所示。我们发现，断开到解码器1、2、3、4的链接，只连接中间块，可以将训练速度提高约1.6倍(在RTX 3070TI笔记本电脑GPU上测试)。当模型显示出结果与条件之间的合理关联时，断开的环节可以在继续训练中重新连接起来，以方便精确控制。

### 只连接中间块

Large-Scale Training

这里的大规模训练是指同时拥有强大的计算集群(至少8个Nvidia A100 80G或同等)和大型数据集(至少100万对训练图像)的情况。这通常适用于数据容易获得的任务，例如，Canny。在这种情况下，由于过度拟合的风险相对较低，我们可以首先训练ControlNets进行足够多的迭代(通常超过50k步)，然后解锁SD所有权重，并将整个模型作为一个整体进行联合训练。这将训练出一个针对更具体问题的模型。

3.5. Implementation

我们提出了几种基于不同图像条件的ControlNets实现，以各种方式控制大型扩散模型。

Classifier-free guidance resolution weighting

Stable Diffusion依赖于一种称为无分类器指导CFG的技术来生成高质量的图像，CFG 被表述为

$\epsilon_{prd} = \epsilon _{uc} +\beta_{cfg}(\epsilon_c - \epsilon_{uc})$

其中 εprd、εuc、εc、βcfg 分别是模型的最终输出、无条件输出、条件输出和用户指定的权重。当通过ControlNet添加条件图像时，它可以添加到εuc和εc中，或者只添加到εc中。在具有挑战性的情况下，例如，当没有给出提示时，将其添加到 εuc 和 εc 将完全消除 CFG 指导（图 5b）；仅使用 εc 将使指导非常强大（图 5c）。

我们的解决方案是首先将条件图像添加到 $\epsilon _c$ ，然后根据每个块的分辨率 $w_i = 64/h_i$ ，将权重 $w_i$ 乘以Stable Diffusion和ControlNet之间的每个连接，其中 $h_i$ 是第 i 个块的大小，例如，h1=8，h2=16,...,h13＝64。

通过减少 CFG 引导强度，我们可以得到图 5d 所示的结果，我们称之为 CFG 分辨率加权。

Composing multiple ControlNets

为了将多个条件图像(如Canny边缘和姿态)应用于Stable Diffusion的单个实例，我们可以直接将相应ControlNets的输出添加到Stable Diffusion模型中(图 6)。这种组合不需要额外加权或线性插值。

4. Experiments

4.2. Ablative Study

我们通过以下方法来研究 ControlNets 的替代结构

用用高斯权重初始化的标准卷积层替换零卷积
用单个卷积层替换每个块的可训练副本，我们称之为 ControlNet-lite。

我们提出了 4 个提示设置来测试现实世界用户可能的行为：

没有提示；
不足以完全覆盖条件图像中的对象，例如，本文的默认提示“高质量、详细和专业图像”；
改变条件图像语义的冲突提示；
描述必要内容语义的完美提示，例如“好房子”。

图 8a 在所有 4 个设置中显示 ControlNet 成功。轻量级ControlNet-lite(图8c)不足以解释条件图像，并且在不充分和没有提示条件下失败。当替换零卷积时，ControlNet 的性能下降到与 ControlNet-lite 大致相同，这表明可训练副本的预训练主干在微调期间被破坏（图 8b）。

4.3. Quantitative Evaluation

4.4. Comparison to Previous Methods

4.5. Discussion

Influence of training dataset sizes

Figure 22: The influence of different training dataset sizes.

Capability to interpret contents

Figure 11: Interpreting contents. If the input is ambiguous and the user does not mention object contents in prompts, the results look like the model tries to interpret input shapes.

Transferring to community models

Figure 12: Transfer pretrained ControlNets to community models without training the neural networks again.

### 备注 8：具体使用Stable Diffusion+ControlNet时，如第一栏，在Stable Diffusion主界面随机生成一只Default turtle。然后把这只Default turtle输入主界面下面一点的ControlNet栏，勾选Scribble（需要同时勾选预处理器、提前下载Scribble模型），点击“生成”就会生成Scribble turtle。然后再把这只Scribble turtle输入ControlNet栏（这次不勾选预处理器，模型仍然选择Scribble），然后在主界面提示prompt栏输入咒语“a masterpiece of cartoon-style turtle illustration”，执行两次，就会生成最后用户期望的两只乌龟。

5. Conclusion

ControlNet 是一种神经网络架构，它学习大型预训练文本到图像扩散模型的条件控制。它重用源模型的大规模预训练层来构建一个深度和强大的编码器来学习特定的条件。原始模型和可训练副本通过“零卷积”层连接，这些层在训练期间消除有害噪声。大量实验表明，ControlNet 在有或没有提示的情况下，可以有效地控制具有单一或多个条件的Stable Diffusion。不同条件数据集的结果表明ControlNet结构真正的可能适用于更广泛的条件，并促进相关应用发展。

第二部分自定义数据集模型训练

1）确保Anaconda、显卡、Python、Pytorch准备就绪。

2）准备ControlNet仓库文件。

在 E:\Aigraphx\AIGC 目录下执行

git clone https://github.com/lllyasviel/ControlNet.git

3）创建conda虚拟环境。

E:\Aigraphx\AIGC

conda env create -f environment.yaml
conda activate control

4）Stable Diffusion 模型准备。

保守起见，先下载SD1.5版本到 E:\Aigraphx\AIGC\models。由于特殊原因，https://huggingface.co/

网站经常上不去，大家选用国内某个镜像：

https://aliendao.cn/models/runwayml/stable-diffusion-v1-5 ###v1.5 512*512 版本
https://gitee.com/modelee/stable-diffusion-2-1-base/tree/main/ ### V2.1 版本，512和768版本

请注意，ControlNet 复制了SD内所有权重，因此没有必要从头开始训练，微调整个模型。

在 E:\Aigraphx\AIGC 下执行

python tool_add_control.py ./models/v1-5-pruned.ckpt ./models/control_sd15_ini.ckpt

此时会报错，提示没有CLIP tokenizer。

5）准备OpenAI CLIP模型

E:\Aigraphx\AIGC\openai\clip-vit-large-patch14目录下执行

https://aliendao.cn/models/openai/clip-vit-large-patch14

再次执行

python tool_add_control.py ./models/v1-5-pruned.ckpt ./models/control_sd15_ini.ckpt

6）公开数据集数据准备

下载公开数据集fill50k，准备就绪的文档目录结构如下。

### 9：原ControlNet官方文档这样描述数据集的。准备好的数据集像是一个包含 50000 个类似数组的对象。每个项目都是一个包含三个条目“jpg”、“txt”和“hint”的字典。其中“jpg”为目标图像target image，“hint”为控制图像 control image，“txt”为提示promt。不要问我们为什么使用这三个名称 - 这与一个名为 LDM 的库的黑暗历史有关。<->

train.jsonl文件内容摘要如下：

{"text": "pale golden rod circle with old lace background", "image": "images/0.png", "conditioning_image": "conditioning_images/0.png"}
{"text": "light coral circle with white background", "image": "images/1.png", "conditioning_image": "conditioning_images/1.png"}

6）准备自定义"数据集"。

参考fill50k数据集，自定义训练需要source：conditioning images，target：images，和Json三类文件。参考公开数据集，咱们先跑通一张source、一张target的训练情形。

拷贝Fill50K数据集target目录下第一张图到E:\Aigraphx\AIGC\Training\target

拷贝Fill50K数据集conditioning images目录下第一张图到E:\Aigraphx\AIGC\Training\source

更名1.png为响应的0.png，编辑promt.josn文件，内容如下。

{"prompt": "pale golden rod circle with old lace background"，"source": "source/0.png", "target": "target/0.png"}

### 10：可以使用ControlNet-WebUI批量生成prompt，然后再手工修正。

7）准备训练脚本文件

编辑内存管理脚本config.py

save_memory = True

编辑数据集脚本tutorial_dataset.py

class MyDataset(Dataset):def __init__(self):self.data = []with open('./training/prompt.json', 'rt') as f:for line in f:self.data.append(json.loads(line))def __len__(self):return len(self.data)def __getitem__(self, idx):item = self.data[idx]source_filename = item['source']target_filename = item['target']prompt = item['prompt']source = cv2.imread('./training/' + source_filename)target = cv2.imread('./training/' + target_filename)

编辑训练脚本tutorial_train.py

# Configs
resume_path = './models/control_sd15_ini.ckpt'
batch_size = 4
logger_freq = 300
learning_rate = 1e-5
sd_locked = True
only_mid_control = True# First use cpu to load models. Pytorch Lightning will automatically move it to GPUs.
model = create_model('./models/cldm_v15.yaml').cpu()
model.load_state_dict(load_state_dict(resume_path, location='cpu'),strict=False)
model.learning_rate = learning_rate
model.sd_locked = sd_locked
model.only_mid_control = only_mid_control

### 11：结合3.4节的Small-Scale Training 模式，修改成只用中间块only_mid_control = True。

因为训练时报GPU错误，修改

model.load_state_dict(load_state_dict(resume_path, location='cpu')) 为：

model.load_state_dict(load_state_dict(resume_path, location='cpu'),strict=False)

8）训练Training

9）精确控制图片生成测试

请参考第四部分的完整功能图形界面测试。

第三部分代码讲解

1.ControlNet结构示意图

结合代码来看，ControlNet核心代码cldm.py和ldm/modules/diffusionmodules/openaimodel.py的UnetModel的input_block, middle_block一样，重点加入了zero_conv操作。

如下图所示参数的定义，几乎一模一样。

2.Controlnet逻辑架构

cldm.py中定义了三个类

class ControlLDM(LatentDiffusion)
class ControlledUnetModel(UNetModel)
class ControlNet(nn.Module)

ControlLDM是总调度。ControlledUnetModel是模型骨干，如果没有Condition Control信息，就跑常规Stable Diffusion流程；否则跑ControNet流程，并把结果添加到原SD中间块。ControlNet负责处理input block，Hint block、middle block各个block，并zero_conv输出。

### 11

类继承： ControlLDM--->LatentDiffusion--->DDPM。
调用：DDPM.training_step--->ControlLDM.get_input--->LatentDiffusion.get_input--->LatentDiffusion.forward--->LatentDiffusion.p_losses。

from ldm.modules.attention import SpatialTransformer
from ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock
from ldm.models.diffusion.ddpm import LatentDiffusion
from ldm.models.diffusion.ddim import DDIMSampler

### 12：拓扑图来自以下博文。

ControlNet_stable diffusion controlnet-CSDN博客

3.ControlNet代码解读

ControlNet模型包括ControlNet、ControlledUnetModel、ControlLDM三部分内容。请自行补充SD中潜在空间、自注意、交叉自注意、时间标签等内容。

ControlNet实现了timestep、hint Map、和PromptEmbedding特征融合，同时进行下采样，增加特征图的通道数。

如果ControlLDM要引入条件控制部分Condition Control加入训练，则启用这部分模块，内容来自cldm.py

class ControlNet(nn.Module):
# 13 input 块self.input_blocks = nn.ModuleList([TimestepEmbedSequential(conv_nd(dims, in_channels, model_channels, 3, padding=1))])
# 14 零卷积 块self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels)])# 15 Hint 块。处理control image输入，原文作者解释为何要取Hint这个名字，说是LDM"黑暗"历史。self.input_hint_block = TimestepEmbedSequential(conv_nd(dims, hint_channels, 16, 3, padding=1),nn.SiLU(),conv_nd(dims, 16, 16, 3, padding=1),nn.SiLU(),conv_nd(dims, 16, 32, 3, padding=1, stride=2),nn.SiLU(),conv_nd(dims, 32, 32, 3, padding=1),nn.SiLU(),conv_nd(dims, 32, 96, 3, padding=1, stride=2),nn.SiLU(),conv_nd(dims, 96, 96, 3, padding=1),nn.SiLU(),conv_nd(dims, 96, 256, 3, padding=1, stride=2),nn.SiLU(),
# 16 Hint块中，先经历几层卷积、激活之后，再次零卷积操作。zero_module(conv_nd(dims, 256, model_channels, 3, padding=1)))self._feature_size = model_channelsinput_block_chans = [model_channels]ch = model_channelsds = 1for level, mult in enumerate(channel_mult):for nr in range(self.num_res_blocks[level]):layers = [ResBlock(ch,time_embed_dim,dropout,out_channels=mult * model_channels,dims=dims,use_checkpoint=use_checkpoint,use_scale_shift_norm=use_scale_shift_norm,)]ch = mult * model_channelsif ds in attention_resolutions:if num_head_channels == -1:dim_head = ch // num_headselse:num_heads = ch // num_head_channelsdim_head = num_head_channelsif legacy:# num_heads = 1dim_head = ch // num_heads if use_spatial_transformer else num_head_channelsif exists(disable_self_attentions):disabled_sa = disable_self_attentions[level]else:disabled_sa = Falseif not exists(num_attention_blocks) or nr < num_attention_blocks[level]:layers.append(AttentionBlock(ch,use_checkpoint=use_checkpoint,num_heads=num_heads,num_head_channels=dim_head,use_new_attention_order=use_new_attention_order,) if not use_spatial_transformer else SpatialTransformer(ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,use_checkpoint=use_checkpoint))self.input_blocks.append(TimestepEmbedSequential(*layers))self.zero_convs.append(self.make_zero_conv(ch))self._feature_size += chinput_block_chans.append(ch)if level != len(channel_mult) - 1:out_ch = chself.input_blocks.append(TimestepEmbedSequential(ResBlock(ch,time_embed_dim,dropout,out_channels=out_ch,dims=dims,use_checkpoint=use_checkpoint,use_scale_shift_norm=use_scale_shift_norm,down=True,)if resblock_updownelse Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)))ch = out_chinput_block_chans.append(ch)self.zero_convs.append(self.make_zero_conv(ch))ds *= 2self._feature_size += chif num_head_channels == -1:dim_head = ch // num_headselse:num_heads = ch // num_head_channelsdim_head = num_head_channelsif legacy:# num_heads = 1dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
# 17 中间块middle_block,包含ResBlock、AttentionBlock、ResBlock块，最后零卷积。self.middle_block = TimestepEmbedSequential(ResBlock(ch,time_embed_dim,dropout,dims=dims,use_checkpoint=use_checkpoint,use_scale_shift_norm=use_scale_shift_norm,),AttentionBlock(ch,use_checkpoint=use_checkpoint,num_heads=num_heads,num_head_channels=dim_head,use_new_attention_order=use_new_attention_order,) if not use_spatial_transformer else SpatialTransformer(  # always uses a self-attnch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,use_checkpoint=use_checkpoint),ResBlock(ch,time_embed_dim,dropout,dims=dims,use_checkpoint=use_checkpoint,use_scale_shift_norm=use_scale_shift_norm,),)self.middle_block_out = self.make_zero_conv(ch)self._feature_size += chdef make_zero_conv(self, channels):return TimestepEmbedSequential(zero_module(conv_nd(self.dims, channels, channels, 1, padding=0)))
# 18 整体guided_hint输入包括：加噪音的潜在z、Control image Input、Prompt、Negative Prompt、Random seed Input、timesteps。def forward(self, x, hint, timesteps, context, **kwargs):t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)emb = self.time_embed(t_emb)guided_hint = self.input_hint_block(hint, emb, context)outs = []h = x.type(self.dtype)for module, zero_conv in zip(self.input_blocks, self.zero_convs):if guided_hint is not None:h = module(h, emb, context)h += guided_hintguided_hint = Noneelse:h = module(h, emb, context)outs.append(zero_conv(h, emb, context))h = self.middle_block(h, emb, context)outs.append(self.middle_block_out(h, emb, context))return outs

4. ControlledUnetModel代码解读

ControlLDM骨干网，内容来自cldm.py

class ControlledUnetModel(UNetModel):def forward(self, x, timesteps=None, context=None, control=None, only_mid_control=False, **kwargs):# 19 hs 用于保存下采样的每一层的输出，这样在上采样的时候方便进行skip-connection        hs = []  with torch.no_grad():  t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)emb = self.time_embed(t_emb)  h = x.type(self.dtype)  for module in self.input_blocks:h = module(h, emb, context)hs.append(h)h = self.middle_block(h, emb, context)# 20 上面是control=None时，Locked Encoder(DownSample) + Middle 计算出来结果
# 21 下面则是有Control情形时，完成条件control及编码的相加。然后和h进行cat，之后再进行Up sample操作if control is not None:h += control.pop()for i, module in enumerate(self.output_blocks):if only_mid_control or control is None:h = torch.cat([h, hs.pop()], dim=1)else:h = torch.cat([h, hs.pop() + control.pop()], dim=1)h = module(h, emb, context)h = h.type(x.dtype)return self.out(h)

5. ControlLDM代码解读

ControlNet模型的总调度ControlLDM，内容来自cldm.py

class ControlLDM(LatentDiffusion):
# 22 训练时选择是否只用中间块，启用小规模训练模式def __init__(self, control_stage_config, control_key, only_mid_control, *args, **kwargs):super().__init__(*args, **kwargs)self.control_model = instantiate_from_config(control_stage_config)self.control_key = control_keyself.only_mid_control = only_mid_controlself.control_scales = [1.0] * 13
# 23 训练时锁定SD参数训练，AdamW只优化ControlNet参数def configure_optimizers(self):lr = self.learning_rateparams = list(self.control_model.parameters())if not self.sd_locked:params += list(self.model.diffusion_model.output_blocks.parameters())params += list(self.model.diffusion_model.out.parameters())opt = torch.optim.AdamW(params, lr=lr)return opt
# 24 训练时加入LoRA微调，适当时把操作放到CPU。LatentDiffusion模型VAE部分是self.first_stage_model，CLIP部分是self.cond_stage_model。def low_vram_shift(self, is_diffusing):if is_diffusing:self.model = self.model.cuda()self.control_model = self.control_model.cuda()self.first_stage_model = self.first_stage_model.cpu()self.cond_stage_model = self.cond_stage_model.cpu()else:self.model = self.model.cpu()self.control_model = self.control_model.cpu()self.first_stage_model = self.first_stage_model.cuda()self.cond_stage_model = self.cond_stage_model.cuda()

第四部分 ControlNet精确控制测试

Stable Diffusion和ControlNet WebUI 有手动+自动+整合包三种方式，本文假设大家都有一定基础，直接用秋叶整合包安装。

Stable Diffusion秋叶整合包是中国大神秋叶（bilibili@秋葉aaaki）基于Stable Diffusion WebUI内核开发的整合包，内置了与电脑本身系统隔离的Python环境和Git命令集，包含了第三部分需要下载和安装的依赖项、GitHub依赖包、预训练模型以及LoRA、ControlNet插件。