ControlNet

Mô hình ControlNet là các adapter được huấn luyện trên một mô hình đã được huấn luyện trước khác. Nó cho phép kiểm soát nhiều hơn trong việc tạo ảnh bằng cách điều kiện hóa mô hình với một ảnh đầu vào bổ sung. Ảnh đầu vào có thể là canny edge, depth map, human pose, và nhiều loại khác.

Nếu bạn đang huấn luyện trên GPU có vRAM hạn chế, bạn nên thử bật các tham số gradient_checkpointing, gradient_accumulation_steps và mixed_precision trong lệnh huấn luyện. Bạn cũng có thể giảm lượng bộ nhớ sử dụng bằng cách dùng attention tiết kiệm bộ nhớ với xFormers.

Hướng dẫn này sẽ khám phá script huấn luyện train_controlnet.py để giúp bạn làm quen với nó, cũng như cách bạn có thể điều chỉnh nó cho trường hợp sử dụng của riêng mình.

Trước khi chạy script, hãy đảm bảo bạn cài đặt thư viện từ mã nguồn:

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

Sau đó, điều hướng đến thư mục ví dụ chứa script huấn luyện và cài đặt các phụ thuộc cần thiết cho script bạn đang sử dụng:

cd examples/controlnet
pip install -r requirements.txt

🤗 Accelerate là một thư viện giúp bạn huấn luyện trên nhiều GPU/TPU hoặc với độ chính xác hỗn hợp. Nó sẽ tự động cấu hình thiết lập huấn luyện của bạn dựa trên phần cứng và môi trường của bạn. Hãy xem 🤗 Accelerate Quick tour để tìm hiểu thêm.

Khởi tạo một môi trường 🤗 Accelerate:

accelerate config

Để thiết lập môi trường 🤗 Accelerate mặc định mà không cần chọn bất kỳ cấu hình nào:

accelerate config default

Hoặc nếu môi trường của bạn không hỗ trợ shell tương tác, chẳng hạn như notebook, bạn có thể sử dụng:

from accelerate.utils import write_basic_config

write_basic_config()

Cuối cùng, nếu bạn muốn huấn luyện một mô hình trên bộ dữ liệu của riêng mình, hãy xem hướng dẫn Tạo bộ dữ liệu để huấn luyện để tìm hiểu cách tạo một bộ dữ liệu tương thích với script huấn luyện.

Các phần tiếp theo làm nổi bật những phần quan trọng trong script huấn luyện để hiểu cách chỉnh sửa nó, nhưng không đi sâu vào mọi khía cạnh của script. Nếu bạn muốn tìm hiểu thêm, hãy tự do đọc qua script và cho chúng tôi biết nếu bạn có bất kỳ câu hỏi hoặc lo ngại nào.

Tham số của script

Tệp script huấn luyện cung cấp nhiều tham số để giúp bạn tùy chỉnh phiên huấn luyện của mình. Tất cả các tham số và mô tả của chúng được tìm thấy trong hàm parse_args(). Hàm này cung cấp các giá trị mặc định cho từng tham số, chẳng hạn như kích thước lô huấn luyện và tốc độ học, nhưng bạn cũng có thể tự đặt giá trị của riêng mình trong lệnh huấn luyện nếu muốn.

Ví dụ, để tăng tốc độ huấn luyện với độ chính xác hỗn hợp sử dụng định dạng fp16, hãy thêm tham số --mixed_precision vào lệnh huấn luyện:

accelerate launch train_controlnet.py \
  --mixed_precision="fp16"

Nhiều tham số cơ bản và quan trọng được mô tả trong hướng dẫn huấn luyện Text-to-image, vì vậy hướng dẫn này chỉ tập trung vào các tham số liên quan đến ControlNet:

--max_train_samples: số lượng mẫu huấn luyện; bạn có thể giảm giá trị này để huấn luyện nhanh hơn, nhưng nếu bạn muốn stream các bộ dữ liệu thực sự lớn, bạn sẽ cần bao gồm tham số này cùng với tham số --streaming trong lệnh huấn luyện của mình
--gradient_accumulation_steps: số bước cập nhật cần tích lũy trước khi thực hiện bước truyền ngược; điều này cho phép bạn huấn luyện với kích thước lô lớn hơn so với mức mà bộ nhớ GPU của bạn thường có thể xử lý

Trọng số Min-SNR

Chiến lược trọng số Min-SNR có thể hỗ trợ quá trình huấn luyện bằng cách cân bằng lại loss để đạt được sự hội tụ nhanh hơn. Script huấn luyện hỗ trợ dự đoán epsilon (noise) hoặc v_prediction, nhưng Min-SNR tương thích với cả hai loại dự đoán. Chiến lược trọng số này chỉ được hỗ trợ bởi PyTorch.

Thêm tham số --snr_gamma và đặt nó bằng giá trị khuyến nghị là 5.0:

accelerate launch train_controlnet.py \
  --snr_gamma=5.0

Script huấn luyện

Tương tự như các tham số của script, một hướng dẫn tổng quan về script huấn luyện được cung cấp trong hướng dẫn huấn luyện Text-to-image. Thay vào đó, hướng dẫn này sẽ xem xét các phần liên quan của script ControlNet.

Script huấn luyện có hàm make_train_dataset để tiền xử lý bộ dữ liệu với các biến đổi ảnh và token hóa caption. Bạn sẽ thấy rằng ngoài việc token hóa caption và biến đổi ảnh thông thường, script còn bao gồm các biến đổi cho ảnh điều kiện.

Nếu bạn đang stream một dataset trên TPU, hiệu năng có thể bị giới hạn bởi thư viện 🤗 Datasets vì nó chưa được tối ưu cho ảnh. Để đảm bảo thông lượng tối đa, bạn nên tìm hiểu các định dạng dataset khác như WebDataset, TorchData và TensorFlow Datasets.

conditioning_image_transforms = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution),
        transforms.ToTensor(),
    ]
)

Bên trong hàm main(), bạn sẽ tìm thấy mã để nạp tokenizer, text encoder, scheduler và các model. Đây cũng là nơi model ControlNet được nạp từ các trọng số có sẵn hoặc được khởi tạo ngẫu nhiên từ một 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)

optimizer được thiết lập để cập nhật các tham số của 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,
)

Cuối cùng, trong vòng lặp huấn luyện, các embedding văn bản điều kiện và hình ảnh được truyền vào các khối down và mid của mô hình 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,
)

Nếu bạn muốn tìm hiểu thêm về cách hoạt động của vòng lặp huấn luyện, hãy xem hướng dẫn Hiểu về pipeline, mô hình và trình lập lịch, nơi phân tích chi tiết mẫu cơ bản của quá trình khử nhiễu.

Chạy script

Bây giờ bạn đã sẵn sàng để chạy script huấn luyện! 🚀

Hướng dẫn này sử dụng bộ dữ liệu fusing/fill50k, nhưng hãy nhớ rằng bạn có thể tự tạo và sử dụng bộ dữ liệu của riêng mình nếu muốn (xem hướng dẫn Tạo bộ dữ liệu để huấn luyện).

Đặt biến môi trường MODEL_NAME thành một model id trên Hub hoặc một đường dẫn đến model cục bộ và OUTPUT_DIR thành nơi bạn muốn lưu model.

Tải xuống các hình ảnh sau để điều kiện cho quá trình huấn luyện của bạn:

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

Một điều nữa trước khi bạn chạy script! Tùy thuộc vào GPU bạn đang dùng, bạn có thể cần bật một số tối ưu hóa nhất định để huấn luyện ControlNet. Cấu hình mặc định trong script này yêu cầu khoảng 38GB vRAM. Nếu bạn đang huấn luyện trên nhiều hơn một GPU, hãy thêm tham số --multi_gpu vào lệnh accelerate launch.

16GB

12GB

8GB

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

Khi quá trình huấn luyện hoàn tất, bạn có thể sử dụng mô hình vừa được huấn luyện để suy luận!

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch

controlnet = ControlNetModel.from_pretrained("path/to/controlnet", dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "path/to/base/model", controlnet=controlnet, dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

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) là một mô hình text-to-image mạnh mẽ, tạo ra các hình ảnh độ phân giải cao, và bổ sung một text-encoder thứ hai vào kiến trúc của nó. Sử dụng script train_controlnet_sdxl.py để huấn luyện một adapter ControlNet cho mô hình SDXL.

Tài liệu hướng dẫn huấn luyện SDXL thảo luận chi tiết hơn về script huấn luyện SDXL.

Các bước tiếp theo

Chúc mừng bạn đã huấn luyện thành công ControlNet của riêng mình! Để tìm hiểu thêm về cách sử dụng mô hình mới của bạn, các hướng dẫn sau có thể hữu ích:

Tìm hiểu cách sử dụng ControlNet để suy luận trên nhiều tác vụ khác nhau.

Update on GitHub

Diffusers

ControlNet

Tham số của script

Trọng số Min-SNR

Script huấn luyện

Chạy script

Stable Diffusion XL

Các bước tiếp theo