Hugging Face's logo

Diffusers 文档

ControlNet

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

ControlNet

ControlNet 模型是在另一个预训练模型之上训练的适配器。它允许通过使用额外的输入图像来调节模型，从而更好地控制图像生成。

如果您在 vRAM 有限的 GPU 上进行训练，您应该尝试在训练命令中启用 gradient_checkpointing、gradient_accumulation_steps 和 mixed_precision 参数。您还可以通过使用 xFormers 的内存高效注意力来减少内存占用。JAX/Flax 训练也支持在 TPUs 和 GPUs 上进行高效训练，但它不支持 gradient checkpointing 或 xFormers。如果您想使用 Flax 更快地进行训练，您应该拥有一个内存大于 30GB 的 GPU。

本指南将探讨 train_controlnet.py 训练脚本，以帮助您熟悉它，以及如何将其适配到您自己的用例。

在运行脚本之前，请确保您从源代码安装了库

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

然后导航到包含训练脚本的示例文件夹，并安装您正在使用的脚本所需的依赖项

PyTorch
Flax
cd examples/controlnet
pip install -r requirements.txt

🤗 Accelerate 是一个库，旨在帮助您在多个 GPUs/TPUs 上或使用混合精度进行训练。它将根据您的硬件和环境自动配置您的训练设置。请查看 🤗 Accelerate Quick tour 了解更多信息。

初始化 🤗 Accelerate 环境

accelerate config

要设置默认的 🤗 Accelerate 环境而无需选择任何配置

accelerate config default

或者，如果您的环境不支持交互式 shell，例如 notebook，您可以使用

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 内存通常可以处理的更大的批次大小进行训练

Min-SNR 权重

Min-SNR 权重策略可以通过重新平衡损失以实现更快的收敛来帮助训练。训练脚本支持预测 epsilon（噪声）或 v_prediction，但 Min-SNR 与这两种预测类型都兼容。此权重策略仅受 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() 函数中，您将找到用于加载 tokenizer、文本编码器、scheduler 和模型的代码。这也是 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,
)

如果您想了解更多关于训练循环如何工作的信息，请查看 理解 pipelines、模型和 schedulers 教程，其中分解了去噪过程的基本模式。

启动脚本

现在您已准备好启动训练脚本！🚀

本指南使用 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
12GB
8GB

在 16GB 显卡上，您可以使用 bitsandbytes 8 位优化器和梯度检查点来优化您的训练运行。安装 bitsandbytes

pip install bitsandbytes

然后，将以下参数添加到您的训练命令中

accelerate launch train_controlnet.py \
  --gradient_checkpointing \
  --use_8bit_adam \
PyTorch
Flax
export MODEL_DIR="stable-diffusion-v1-5/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 对各种任务进行推理。
< > 在 GitHub 上更新

←Wuerstchen T2I-Adapters→