ControlNet
ControlNet 模型是在另一个预训练模型的基础上训练的适配器。通过用额外的输入图像对模型进行条件化,它可以更好地控制图像生成。输入图像可以是边缘检测图、深度图、人体姿态图等等。
如果您是在 vRAM 有限的 GPU 上进行训练,则应尝试在训练命令中启用 gradient_checkpointing
、gradient_accumulation_steps
和 mixed_precision
参数。您还可以使用内存高效的注意力机制(例如 xFormers)来减少内存占用。JAX/Flax 训练也支持在 TPU 和 GPU 上进行高效训练,但它不支持梯度检查点或 xFormers。如果您想使用 Flax 更快地进行训练,则您的 GPU 内存应大于 30GB。
本指南将探讨 train_controlnet.py 训练脚本,帮助您熟悉它,以及如何将其用于自己的用例。
在运行脚本之前,请确保您从源代码安装了库
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
然后导航到包含训练脚本的示例文件夹,并安装您正在使用的脚本所需的依赖项
cd examples/controlnet
pip install -r requirements.txt
🤗 Accelerate 是一个用于帮助您在多个 GPU/TPU 上或使用混合精度进行训练的库。它会根据您的硬件和环境自动配置您的训练设置。请查看 🤗 Accelerate 快速入门,了解更多信息。
初始化 🤗 Accelerate 环境
accelerate config
要设置默认的 🤗 Accelerate 环境,无需选择任何配置
accelerate config default
或者,如果您的环境不支持交互式 shell(例如笔记本),则可以使用
from accelerate.utils import write_basic_config
write_basic_config()
最后,如果您想使用自己的数据集训练模型,请查看 创建用于训练的数据集 指南,了解如何创建适用于训练脚本的数据集。
以下部分重点介绍训练脚本中理解如何修改脚本的重要部分,但没有详细介绍脚本的各个方面。如果您有兴趣了解更多信息,请随时浏览 脚本,如果您有任何问题或疑虑,请告知我们。
脚本参数
训练脚本提供许多参数,帮助您自定义训练运行。所有参数及其描述都可以在 parse_args()
函数中找到。此函数为每个参数提供默认值,例如训练批次大小和学习率,但您也可以在训练命令中设置自己的值(如果您愿意)。
例如,要使用 fp16 格式的混合精度来加速训练,请将 --mixed_precision
参数添加到训练命令中
accelerate launch train_controlnet.py \
--mixed_precision="fp16"
许多基本且重要的参数在 文本到图像 训练指南中进行了描述,因此本指南仅重点介绍与 ControlNet 相关的参数
--max_train_samples
:训练样本数量;这可以降低以加快训练速度,但如果您想流式传输非常大的数据集,则需要在训练命令中包含此参数和--streaming
参数--gradient_accumulation_steps
:在反向传播之前累积的更新步骤数量;这使您可以使用比 GPU 内存通常可以处理的更大的批次大小进行训练
最小信噪比加权
最小信噪比 加权 策略可以通过重新平衡损失来帮助训练,从而实现更快的收敛。训练脚本支持预测 epsilon
(噪声)或 v_prediction
,但最小信噪比与两种预测类型兼容。此加权策略仅受 PyTorch 支持,在 Flax 训练脚本中不可用。
添加 --snr_gamma
参数并将其设置为推荐值 5.0
accelerate launch train_controlnet.py \ --snr_gamma=5.0
训练脚本
与脚本参数一样,训练脚本的通用流程在 文本到图像 训练指南中进行了介绍。相反,本指南将介绍 ControlNet 脚本的相关部分。
训练脚本有一个 make_train_dataset
函数,用于使用图像变换和标题标记化对数据集进行预处理。您会看到,除了通常的标题标记化和图像变换外,脚本还包括针对条件化图像的变换。
如果您在 TPU 上流式传输数据集,性能可能会受到 🤗 Datasets 库的限制,该库并未针对图像进行优化。为了确保最大吞吐量,建议您探索其他数据集格式,例如 WebDataset、TorchData 和 TensorFlow Datasets。
conditioning_image_transforms = transforms.Compose( [ transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR), transforms.CenterCrop(args.resolution), transforms.ToTensor(), ] )
在 main()
函数中,您会找到加载标记器、文本编码器、调度器和模型的代码。这也是 ControlNet 模型从现有权重加载或从 UNet 随机初始化的地方
if args.controlnet_model_name_or_path:
logger.info("Loading existing controlnet weights")
controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path)
else:
logger.info("Initializing controlnet weights from unet")
controlnet = ControlNetModel.from_unet(unet)
优化器 设置 为更新 ControlNet 参数
params_to_optimize = controlnet.parameters() optimizer = optimizer_class( params_to_optimize, lr=args.learning_rate, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, eps=args.adam_epsilon, )
最后,在 训练循环 中,条件化文本嵌入和图像将传递到 ControlNet 模型的下层和中间层
encoder_hidden_states = text_encoder(batch["input_ids"])[0]
controlnet_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)
down_block_res_samples, mid_block_res_sample = controlnet(
noisy_latents,
timesteps,
encoder_hidden_states=encoder_hidden_states,
controlnet_cond=controlnet_image,
return_dict=False,
)
如果你想了解更多关于训练循环的工作原理,请查看 理解管道、模型和调度器 教程,它解释了去噪过程的基本模式。
启动脚本
现在你可以启动训练脚本了! 🚀
本指南使用 fusing/fill50k 数据集,但请记住,如果你愿意,你可以创建和使用自己的数据集(参见 创建用于训练的数据集 指南)。
将环境变量 MODEL_NAME
设置为 Hub 上的模型 ID 或本地模型的路径,并将 OUTPUT_DIR
设置为要保存模型的位置。
下载以下图像以用作训练条件:
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png
在启动脚本之前还有一件事!根据你的 GPU,你可能需要启用某些优化来训练 ControlNet。本脚本中的默认配置需要大约 38GB 的 vRAM。如果你在多个 GPU 上训练,请将 --multi_gpu
参数添加到 accelerate launch
命令中。
在 16GB 的 GPU 上,你可以使用 bitsandbytes 8 位优化器和梯度检查点来优化你的训练运行。安装 bitsandbytes
pip install bitsandbytes
然后,将以下参数添加到你的训练命令中
accelerate launch train_controlnet.py \ --gradient_checkpointing \ --use_8bit_adam \
export MODEL_DIR="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/save/model"
accelerate launch train_controlnet.py \
--pretrained_model_name_or_path=$MODEL_DIR \
--output_dir=$OUTPUT_DIR \
--dataset_name=fusing/fill50k \
--resolution=512 \
--learning_rate=1e-5 \
--validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
--validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--push_to_hub
训练完成后,你可以使用新训练的模型进行推理!
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch
controlnet = ControlNetModel.from_pretrained("path/to/controlnet", torch_dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
"path/to/base/model", controlnet=controlnet, torch_dtype=torch.float16
).to("cuda")
control_image = load_image("./conditioning_image_1.png")
prompt = "pale golden rod circle with old lace background"
generator = torch.manual_seed(0)
image = pipeline(prompt, num_inference_steps=20, generator=generator, image=control_image).images[0]
image.save("./output.png")
Stable Diffusion XL
Stable Diffusion XL (SDXL) 是一款功能强大的文本到图像模型,可以生成高分辨率图像,并在其架构中添加了第二个文本编码器。使用 train_controlnet_sdxl.py
脚本训练 SDXL 模型的 ControlNet 适配器。
SDXL 训练脚本在 SDXL 训练 指南中进行了更详细的讨论。
下一步
恭喜你训练了自己的 ControlNet!要了解更多关于如何使用新模型的信息,以下指南可能会有所帮助
- 学习如何 使用 ControlNet 对各种任务进行推理。