使用 diffusers 训练你的 ControlNet 🧨

简介

ControlNet 是一种神经网络结构，它通过添加额外的条件来实现对扩散模型的细粒度控制。这项技术首次出现在论文 《为文本到图像扩散模型添加条件控制》(Adding Conditional Control to Text-to-Image Diffusion Models) 中，并迅速在开源扩散社区中流行起来，作者发布了 8 种不同的条件来控制 Stable Diffusion v1-5，包括姿态估计、深度图、Canny 边缘、草图，以及更多。

在这篇博文中，我们将详细介绍我们是如何训练 诡异人脸模型 (Uncanny Faces model) 的每一步——这是一个基于 3D 合成人脸的面部姿势模型 (诡异人脸实际上是一个意想不到的结果，请继续关注它是如何产生的)。

开始为 Stable Diffusion 训练你的 ControlNet

训练你自己的 ControlNet 需要 3 个步骤

1. 规划你的条件: ControlNet 的灵活性足以让 Stable Diffusion 适应许多任务。预训练模型展示了广泛的条件，社区也构建了其他条件，例如基于像素化调色板的条件。

2. 构建你的数据集: 一旦确定了条件，就该构建数据集了。为此，你可以从头开始构建数据集，也可以使用现有数据集的子集。你需要数据集中的三列来训练模型：一个基准真相 image、一个 conditioning_image 和一个 prompt。

3. 训练模型: 数据集准备好后，就可以训练模型了。感谢 diffusers 训练脚本，这是最简单的部分。你需要一个至少有 8GB VRAM 的 GPU。

1. 规划你的条件

为了规划你的条件，思考以下两个问题会很有帮助

1. 我想使用什么样的条件？
2. 是否已经有现成的模型可以将“常规”图像转换为我的条件？

在我们的例子中，我们考虑使用面部关键点作为条件。我们的理由是：1. 基于通用关键点条件的 ControlNet 效果很好。2. 面部关键点是一种足够普及的技术，并且有多种模型可以在常规图片上计算面部关键点。3. 让 Stable Diffusion 遵循特定的面部关键点或模仿你自己的面部表情可能会很有趣。

2. 构建你的数据集

好的！我们决定使用面部关键点作为 Stable Diffusion 的条件。所以，为了准备数据集，我们需要

• 基准真相 image：在这里是人脸图像
• conditioning_image：在这里是可视化了面部关键点的图像
• caption：描述所用图像的标题

对于这个项目，我们决定使用微软的 FaceSynthetics 数据集：它是一个包含 10 万张合成人脸的数据集。其他包含真实人脸的面部研究数据集，如 Celeb-A HQ、FFHQ - 但我们决定在这个项目中使用合成人脸。

FaceSynthetics 数据集听起来是一个很好的起点：它包含了人脸的基准真相图像，以及以 iBUG 68 点面部关键点格式标注的关键点，还有人脸的分割图像。

完美。是吗？不幸的是，并非如此。还记得在“规划你的条件”步骤中的第二个问题吗——我们应该有模型能将常规图像转换为条件图像？结果发现，没有已知的模型可以将人脸转换成该数据集的标注关键点格式。

所以我们决定走另一条路

• 使用 FaceSynthetics 数据集中的人脸基准真相 image
• 使用一个已知的模型，该模型可以将任何一张人脸图像转换为 iBUG 的 68 点面部关键点格式 (在我们的例子中，我们使用了 SOTA 模型 SPIGA)
• 使用自定义代码将面部关键点转换为一个美观的插图蒙版，用作 conditioning_image
• 将其保存为 Hugging Face 数据集

你可以在这里找到用于将 FaceSynthetics 数据集中的基准真相图像转换为插图蒙版并将其保存为 Hugging Face 数据集的代码。

现在，数据集中有了基准真相 image 和 conditioning_image，我们还差一步：为每张图片添加标题。这一步是强烈推荐的，但你也可以尝试使用空提示词，并反馈你的结果。由于我们没有 FaceSynthetics 数据集的标题，我们通过 BLIP 图像描述模型 为其生成了标题。你可以在这里查看为所有图像生成标题的代码。

这样，我们得到了最终的数据集！带标题的 Face Synthetics SPIGA (Face Synthetics SPIGA with captions) 数据集包含了 FaceSynthetics 数据集中 10 万张图像的基准真相图像、分割图和标题。我们准备好训练模型了！

3. 训练模型

我们的数据集准备好了，是时候训练模型了！虽然这本应是整个过程中最难的部分，但有了 diffusers 训练脚本，它反而成了最简单的部分。我们使用了一台在 LambdaLabs 上租用的 A100，每小时 1.10 美元。

我们的训练经验

我们对模型进行了 3 个 epoch 的训练 (这意味着 10 万张图片的批次被模型看了 3 次)，批处理大小为 4 (每一步向模型展示 4 张图片)。结果证明这样做太多了并且导致了过拟合 (所以它忘记了那些与真实人脸稍有不同的概念，例如在提示词中输入“shrek”或“a cat”不会生成史莱克或猫，而是一个人，并且也开始忽略风格)。

仅用 1 个 epoch (即模型“看过”10 万张图片后)，它就已经收敛到能够遵循姿势并且没有过拟合。所以它成功了，但是...由于我们使用了 FaceSynthetics 数据集，模型最终学会了生成诡异的 3D 脸，而不是逼真的人脸。考虑到我们使用的是合成人脸数据集而不是真实人脸，这是合理的，并且可以用于有趣/梗图的目的。这是 uncannyfaces_25K 模型。

在这个交互式表格中，你可以通过下面的拨盘来查看模型经历了多少训练步骤，以及这对训练过程有何影响。在大约 1.5 万步时，它已经开始学习姿势。在大约 2.5 万步时，它已经成熟。

我们是如何进行训练的

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

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

让我们来分析一些设置，并介绍一些优化技巧，以便在低至 8GB VRAM 的显卡上进行训练。

• pretrained_model_name_or_path: 你想使用的 Stable Diffusion 基础模型 (我们在这里选择了 v2-1，因为它能更好地渲染人脸)
• output_dir: 你希望模型保存的目录
• dataset_name: 用于训练的数据集。在我们的例子中是 带标题的 Face Synthetics SPIGA
• conditioning_image_column: 数据集中包含条件图像的列名 (在我们的例子中是 spiga_seg)
• image_column: 数据集中包含基准真相图像的列名 (在我们的例子中是 image)
• caption_column: 数据集中包含该图像标题的列名 (在我们的例子中是 image_caption)
• resolution: 条件图像和基准真相图像的分辨率 (在我们的例子中是 512x512)
• learning_rate: 学习率。我们发现 1e-5 在这些例子中效果很好，但你可以尝试不同的值，例如在 1e-4 和 2e-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 步运行一次 validation_prompt 和 validation_image。
• report_to: 向哪里报告你的训练情况。这里我们使用了 Weights and Biases，它给了我们这份漂亮的报告。但是，将 train_batch_size 从 4 减少到 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 人脸，而不是逼真的人脸，因为这是它训练时使用的数据集，这也有它自己的魅力和风格。

试试我们的 Hugging Face Space

至于我们的下一步——为了创建看起来逼真的人脸，同时仍然不使用真实人脸数据集，一个想法是运行整个 FaceSynthetics 数据集通过 Stable Diffusion Image2Image，将 3D 脸转换为逼真的脸，然后训练另一个 ControlNet。

请继续关注，因为我们很快将举办一个 ControlNet 训练活动！在 Twitter 上关注 Hugging Face 或加入我们的 Discord 以获取最新消息。