Diffusers 文件
ControlNet
並獲得增強的文件體驗
開始使用
ControlNet
ControlNet 模型是在其他預訓練模型之上訓練的介面卡。它透過使用額外的輸入影像來條件化模型,從而實現對影像生成的更高程度控制。輸入影像可以是 Canny 邊緣、深度圖、人體姿勢等等。
如果你正在 vRAM 有限的 GPU 上訓練,你應該嘗試在訓練命令中啟用 gradient_checkpointing
、gradient_accumulation_steps
和 mixed_precision
引數。你還可以透過使用 xFormers 的記憶體高效注意力來減少記憶體佔用。JAX/Flax 訓練也支援在 TPU 和 GPU 上進行高效訓練,但它不支援梯度檢查點或 xFormers。如果你想用 Flax 更快地訓練,你應該有一個 >30GB 記憶體的 GPU。
本指南將探討 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 記憶體通常能處理的更大的批次大小進行訓練
最小信噪比加權
最小信噪比 (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()
函式中,你會找到用於載入分詞器、文字編碼器、排程器和模型的程式碼。這也是 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)
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 GPU 上,你可以使用 bitsandbytes 8 位最佳化器和梯度檢查點來最佳化你的訓練執行。安裝 bitsandbytes
pip install bitsandbytes
然後,將以下引數新增到您的訓練命令中
accelerate launch train_controlnet.py \ --gradient_checkpointing \ --use_8bit_adam \
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 進行各種任務的推理。