Diffusers 文件

ControlNet

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

ControlNet

ControlNet 模型是在其他預訓練模型之上訓練的介面卡。它透過使用額外的輸入影像來條件化模型,從而實現對影像生成的更高程度控制。輸入影像可以是 Canny 邊緣、深度圖、人體姿勢等等。

如果你正在 vRAM 有限的 GPU 上訓練,你應該嘗試在訓練命令中啟用 gradient_checkpointinggradient_accumulation_stepsmixed_precision 引數。你還可以透過使用 xFormers 的記憶體高效注意力來減少記憶體佔用。JAX/Flax 訓練也支援在 TPU 和 GPU 上進行高效訓練,但它不支援梯度檢查點或 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 是一個幫助你在多個 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 記憶體通常能處理的更大的批次大小進行訓練

最小信噪比加權

最小信噪比 (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 庫的瓶頸,該庫並未針對影像進行最佳化。為確保最大吞吐量,建議你探索其他資料集格式,如 WebDatasetTorchDataTensorFlow 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 的視訊記憶體。如果你在多個 GPU 上進行訓練,請將 --multi_gpu 引數新增到 accelerate launch 命令中。

16GB
12GB
8GB

在 16GB GPU 上,你可以使用 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!要了解如何使用新模型,以下指南可能會有所幫助

< > 在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.