使用扩散器来训练您的ControlNet

使用扩散器训练ControlNet

介绍

ControlNet 是一种神经网络结构,通过添加额外的条件,可以对扩散模型进行细粒度的控制。这项技术首次在论文《为文本到图像扩散模型添加条件控制》中亮相,并迅速在开源扩散社区中获得了广泛应用,作者发布了8种不同的条件来控制 Stable Diffusion v1-5,包括姿势估计、深度图、Canny 边缘、草图等。

在本博客文章中,我们将详细介绍训练 Uncanny Faces 模型的每个步骤——这是一个基于3D合成人脸的面部姿势模型(实际上,uncanny faces 是一个意外的结果,请继续关注以了解它是如何产生的)。

开始训练 Stable Diffusion 的 ControlNet

训练自己的 ControlNet 需要三个步骤:

  1. 规划条件:ControlNet 可以根据多个任务对 Stable Diffusion 进行灵活的控制。预训练模型展示了多种条件,并且社区还构建了其他条件,例如基于像素化的调色板。

  2. 构建数据集:一旦确定了条件,就需要构建数据集。为此,您可以从头开始构建数据集,或者使用现有数据集的子集。您的数据集需要三列来训练模型:一个真实的 image、一个 conditioning_image 和一个 prompt

  3. 训练模型:一旦准备好数据集,就可以开始训练模型了。由于有了扩散器的训练脚本,这是最简单的部分。您需要一块至少有8GB VRAM的GPU。

1. 规划条件

规划条件时,有两个问题很有用:

  1. 我想使用什么类型的条件?
  2. 是否已经存在可以将“常规”图像转换为我的条件的模型?

对于我们的示例,我们考虑使用面部标记条件。我们的理由是:1. 通用的标记条件 ControlNet 效果很好。2. 面部标记是一种广泛使用的技术,有多个模型可以在常规图片上计算面部标记。3. 让 Stable Diffusion 跟随特定的面部标记或模仿您自己的面部表情可能很有趣。

2. 构建数据集

好的!所以我们决定使用面部标记 Stable Diffusion 条件。为了准备数据集,我们需要:

  • 真实的 image:在这种情况下,是人脸图像
  • conditioning_image:在这种情况下,是显示面部标记的图像
  • caption:描述使用的图像的标题

对于这个项目,我们决定使用 Microsoft 的 FaceSynthetics 数据集:它是一个包含10万个合成人脸的数据集。其他包含真实人脸的面部研究数据集,如 Celeb-A HQFFHQ,但我们决定在这个项目中使用合成人脸。

FaceSynthetics 数据集听起来是一个很好的起点:它包含了真实的人脸图像,以 iBUG 68 个面部标记格式注释的面部标记图像,以及人脸的分割图像。

很完美,对吗?不幸的是,实际上并不是。还记得“规划条件”步骤中的第二个问题吗——我们应该有可以将面部转换为该数据集注释的标记格式的已知模型。结果发现,并没有已知的模型可以将人脸转换为该数据集的标记格式。

因此,我们决定走另一条路:

  • 使用FaceSynthetics数据集中的人脸image作为真实值
  • 使用可以将任何人脸图像转换为iBUG 68个面部标记格式的已知模型(在我们的案例中,我们使用了SOTA模型SPIGA)
  • 使用自定义代码将面部标记转换为漂亮的插图掩码,作为conditioning_image
  • 将其保存为Hugging Face数据集

在这里,您可以找到用于将FaceSynthetics数据集中的真实值图像转换为插图掩码并保存为Hugging Face数据集的代码。

现在,在数据集中拥有真实值imageconditioning_image之后,我们还缺少一步:即每个图像的标题。强烈建议执行此步骤,但您可以尝试使用空提示进行实验,并向我们报告结果。由于FaceSynthetics数据集没有标题,我们对其进行了BLIP字幕处理。您可以在此处查看用于给所有图像添加字幕的代码。

有了这些,我们得到了最终的数据集!Face Synthetics SPIGA与字幕包含了FaceSynthetics数据集中100K个图像的真实值图像、分割和字幕。我们已经准备好训练模型了!

3. 训练模型

准备好我们的数据集后,现在是时候训练模型了!尽管这本应该是整个过程中最困难的部分,但借助diffusers训练脚本,它变得最容易。我们在LambdaLabs上租用了一台单个A100,每小时1.10美元。

我们的训练经验

我们将模型训练了3个epoch(这意味着将100K个图像的批次显示给模型3次),批量大小为4(每步向模型显示4个图像)。结果证明这是过度拟合的,它忘记了与真实脸部有些不同的概念,所以例如在提示中使用”shrek”或”a cat”不会生成shrek或猫,而是生成一个人,并且还开始忽略风格。

只需1个epoch(即模型”看到”100K个图像后),它就已经收敛到跟随姿势而不是过度拟合。所以它有效果,但是…由于我们使用了合成人脸数据集,模型最终学习到的是看起来不真实的3D人脸,而不是真实的人脸。考虑到我们使用的是合成人脸数据集而不是真实数据集,这是有道理的,它可以用于娱乐/模因目的。这是uncannyfaces_25K模型。

在这个互动表格中,您可以使用下面的拨轮来查看模型经过多少训练步骤以及它对训练过程的影响。在大约15K步骤时,它已经开始学习姿势。并且在25K步骤时达到成熟。这里

我们如何进行训练

我们所要做的就是安装依赖项:

pip install git+https://github.com/huggingface/diffusers.git transformers accelerate xformers==0.0.16 wandb
huggingface-cli login
wandb login 

然后运行train_controlnet.py代码

!accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base" \
 --output_dir="model_out" \
 --dataset_name=multimodalart/facesyntheticsspigacaptioned \
 --conditioning_image_column=spiga_seg \
 --image_column=image \
 --caption_column=image_caption \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./face_landmarks1.jpeg" "./face_landmarks2.jpeg" "./face_landmarks3.jpeg" \
 --validation_prompt "High-quality close-up dslr photo of man wearing a hat with trees in the background" "Girl smiling, professional dslr photograph, dark background, studio lights, high quality" "Portrait of a clown face, oil on canvas, bittersweet expression" \
 --train_batch_size=4 \
 --num_train_epochs=3 \
 --tracker_project_name="controlnet" \
 --enable_xformers_memory_efficient_attention \
 --checkpointing_steps=5000 \
 --validation_steps=5000 \
 --report_to wandb \
 --push_to_hub

让我们分解一些设置,并且让我们介绍一些优化技巧,以便在训练时将 VRAM 降低到 8GB。

  • pretrained_model_name_or_path:您想使用的稳定扩散基础模型(我们选择了 v2-1,因为它可以更好地渲染面部)
  • output_dir:您希望保存模型的目录
  • dataset_name:用于训练的数据集。在我们的案例中,是带有标题的面部合成 SPIGA 数据集
  • conditioning_image_column:数据集中包含条件图像的列的名称(在我们的案例中是 spiga_seg
  • image_column:数据集中包含真实图像的列的名称(在我们的案例中是 image
  • caption_column:数据集中包含图像标题的列的名称(在我们的案例中是 image_caption
  • resolution:条件和真实图像的分辨率(在我们的案例中是 512x512
  • learning_rate:学习率。我们发现对于这些示例,1e-5 效果很好,但您可以尝试不同范围内的不同值,例如 1e-42e-6 之间。
  • validation_image:这是让您在训练过程中偷看的机会!验证图像将在每个 validation_steps 的数量上运行,以便您可以看到训练进展如何。在此处插入本地路径,以包含任意数量的条件图像。
  • validation_prompt:与验证图像一起运行的提示。可以是任何可以测试模型训练效果的内容。
  • train_batch_size:用于适应 GPU 的训练批次大小。由于我们有 A100,我们可以承担 4,但如果您的 GPU 的 VRAM 较低,我们建议将此值降低到 1
  • num_train_epochs:每个 epoch 对应训练集中的图像被模型“观察”多少次。我们尝试了 3 个 epoch,但结果表明最佳结果只需要略多于 1 个 epoch,3 个 epoch 时我们的模型出现过拟合。
  • checkpointing_steps:每隔 x 步保存一个中间检查点(在我们的案例中是 5000 )。每 5000 步,保存一个中间检查点。
  • validation_steps:每隔 x 步运行一次 validaton_promptvalidation_image
  • report_to:报告训练结果的位置。在这里,我们使用了 Weights and Biases,给我们提供了这个漂亮的报告。但是,将 train_batch_size4 减少到 1 可能不足以让训练适应小型 GPU,因此还有一些额外的参数可供每个 GPU VRAM 大小添加:
  • push_to_hub:将最终训练好的模型推送到 Hugging Face Hub 的参数。

在 16GB VRAM GPU 上适应

pip install bitsandbytes

--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--use_8bit_adam

批次大小为 1 结合 4 个梯度累积步骤相当于在示例中使用的原始批次大小 4。此外,我们启用了梯度检查点和 8 位 Adam 以节省额外的内存。

在 12GB VRAM GPU 上适应

--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--use_8bit_adam
--set_grads_to_none

适配8GB VRAM GPU

请按照我们的指南这里

4. 结论!

训练ControlNet的这次经历非常有趣。我们成功训练了一个能够跟随真实面部姿势的模型 – 但它学会了制作离奇的3D面部而不是真实的3D面部,因为这是它所训练的数据集,具有其自身的魅力和风格。

试试我们的Hugging Face空间:

关于我们的下一步计划 – 为了创建真实看起来的面部,同时不使用真实的面部数据集,一个想法是将整个FaceSynthetics数据集通过稳定扩散图像转换为真实看起来的面部,然后再训练另一个ControlNet。

敬请关注,我们将很快举办ControlNet培训活动!请关注Hugging Face的Twitter或加入我们的Discord以保持最新信息。