ControlNet

ControlNet là kiến trúc mạng nơ-ron bổ sung khả năng điều khiển không gian cho một model khuếch tán văn bản thành ảnh tiền huấn luyện như Stable Diffusion. Ngoài prompt văn bản, model nhận một ảnh điều kiện, chẳng hạn bản đồ cạnh, bản đồ độ sâu, bộ khung pose người hoặc bản đồ phân đoạn; ảnh được tạo sẽ tuân theo cấu trúc đó, còn prompt quyết định nội dung và phong cách. Kết hợp prompt như "hiệp sĩ mặc giáp bạc" với bộ khung pose sẽ tạo ra hình ảnh hiệp sĩ đúng tư thế đó. Lvmin Zhang, Anyi Rao và Maneesh Agrawala tại Đại học Stanford giới thiệu ControlNet vào năm 2023 trong bài báo Bổ sung điều khiển có điều kiện cho model khuếch tán văn bản thành ảnh, công trình đồng nhận giải Marr cho bài báo xuất sắc nhất tại ICCV 2023. Trọng số của model nền được giữ cố định, vì vậy ControlNet bổ sung khả năng điều khiển mà không làm suy giảm những gì model đã biết.

Cách ControlNet hoạt động#

Stable Diffusion tạo ảnh bằng cách liên tục khử nhiễu latent với một U-Net. ControlNet khóa U-Net này và tạo một bản sao có thể huấn luyện của encoder gồm 12 block mã hóa cùng block ở giữa. Trước tiên, một mạng tích chập nhỏ chuyển ảnh điều kiện 512×512 thành bản đồ đặc trưng 64×64 tương ứng với latent. Bản sao có thể huấn luyện xử lý latent cùng các đặc trưng điều kiện; đầu ra của nó được cộng vào 12 kết nối skip và block giữa của U-Net đã khóa, nhờ đó decoder đã khóa nhận được tín hiệu điều khiển:

Bản sao chỉ kết nối với model đã khóa thông qua các phép tích chập zero: phép tích chập 1×1 có trọng số và bias khởi tạo bằng 0. Ở bước huấn luyện đầu tiên, chúng tạo ra đầu ra bằng 0, nên model kết hợp hoạt động chính xác như bản gốc và không có nhiễu ngẫu nhiên truyền đến các lớp tiền huấn luyện. Phép tích chập zero ở đầu ra vẫn nhận gradient trọng số khác 0 vì đầu vào của nó khác 0; do đó, lần cập nhật đầu tiên đưa nó ra khỏi giá trị 0 và tín hiệu điều khiển bắt đầu tăng lên. Đoạn mã PyTorch sau minh họa một block:

import copy

import torch
from torch import nn

# A frozen block from the pretrained model and a trainable copy of it
torch.manual_seed(0)
locked = nn.Conv2d(64, 64, 3, padding=1).requires_grad_(False)
trainable = copy.deepcopy(locked).requires_grad_(True)

# Zero convolutions: 1x1 convolutions whose weights and biases start at zero
zero_in, zero_out = nn.Conv2d(64, 64, 1), nn.Conv2d(64, 64, 1)
for conv in (zero_in, zero_out):
    nn.init.zeros_(conv.weight)
    nn.init.zeros_(conv.bias)

x = torch.randn(1, 64, 32, 32)  # features entering the block
condition = torch.randn(1, 64, 32, 32)  # encoded edge, depth, or pose map
y = locked(x) + zero_out(trainable(x + zero_in(condition)))

# At step 0 the output equals the frozen block's, and only the outer zero conv gets a gradient
print(torch.equal(y, locked(x)))
y.sum().backward()
print(f"zero_in grad: {zero_in.weight.grad.norm():.1f}, zero_out grad: {zero_out.weight.grad.norm():.1f}")

Lệnh này in True và zero_in grad: 0.0, zero_out grad: 2040.6. Đầu ra của block giống hệt model đã đóng băng; phép tích chập zero bên ngoài nhận gradient vì đầu vào khác 0, còn phép bên trong bắt đầu học sau khi phép bên ngoài đã lệch khỏi giá trị 0.

Các loại điều kiện#

Bài báo gốc huấn luyện các ControlNet riêng biệt cho cạnh Canny, bản đồ độ sâu, bản đồ pháp tuyến bề mặt, đường thẳng M-LSD, cạnh mềm HED, bản đồ phân đoạn ADE20K, pose người OpenPose và phác thảo của người dùng. Mỗi ảnh điều kiện được trích xuất từ ảnh tham chiếu bằng model có sẵn, chẳng hạn bộ phát hiện cạnh Canny của OpenCV hoặc một mạng ước lượng độ sâu; có thể áp dụng nhiều ControlNet cùng lúc, mỗi model có cường độ riêng. Các model ước lượng pose của Ultralytics YOLO26 trích xuất loại điều kiện pose này, dự đoán 17 keypoint COCO cho mỗi người; các keypoint có thể được vẽ thành bộ khung trên canvas trống:

import cv2
import numpy as np
from ultralytics import YOLO

# Tải model pose YOLO26 để dự đoán 17 keypoint cho mỗi người
model = YOLO("yolo26n-pose.pt")
results = model("https://ultralytics.com/images/bus.jpg")
result = results[0]

# Chỉ vẽ bộ khung trên canvas màu đen có kích thước bằng ảnh
skeletons = result.plot(img=np.zeros_like(result.orig_img), boxes=False)
cv2.imwrite("pose_condition.png", skeletons)

Ảnh đã lưu hiển thị tay chân của mỗi người dưới dạng các đường màu trên nền đen; màu sắc, kết cấu và phông nền của ảnh đã bị loại bỏ.

ControlNet và các phương pháp liên quan#

• LoRA: Tinh chỉnh một tập hợp nhỏ trọng số được bổ sung để dạy model phong cách, nhân vật hoặc vật thể mới. LoRA thay đổi nội dung model vẽ, còn ControlNet cố định vị trí các thành phần; hai phương pháp thường được kết hợp.
• T2I-Adapter: Một giải pháp thay thế nhẹ hơn với khoảng 77 triệu tham số, bổ sung đặc trưng điều kiện vào đặc trưng encoder của U-Net đã đóng băng thay vì chạy bản sao đầy đủ của encoder. Có thể kết hợp các adapter cho những điều kiện khác nhau khi tạo ảnh.
• Tạo ảnh từ ảnh: Bắt đầu khử nhiễu từ bản sao có nhiễu của ảnh đầu vào, nên đầu ra kế thừa màu sắc và bố cục tổng thể của ảnh đó. ControlNet chỉ ràng buộc cấu trúc được trích xuất, chẳng hạn cạnh, còn màu sắc và kết cấu do prompt quyết định.
• Tô lại ảnh: Tạo lại một phần ảnh được mask để khớp với phần còn lại. Phương pháp này kiểm soát vùng thay đổi, còn ControlNet kiểm soát cấu trúc của nội dung được tạo.

Ứng dụng thực tế#

• Kiến trúc và thiết kế nội thất: Phác thảo, bản vẽ đường nét CAD hoặc bản đồ đường M-LSD của căn phòng có thể được chuyển thành hình ảnh kết xuất chân thực với vật liệu và phong cách khác nhau nhưng giữ nguyên hình học.
• Minh họa nhân vật và hoạt hình: Bộ khung pose giữ nguyên tư thế nhân vật qua các biến thể phong cách; pose trích xuất theo từng khung hình từ video tham chiếu tạo ra các chuỗi hoạt hình nhất quán.
• Hình ảnh sản phẩm và tiếp thị: Bản đồ độ sâu hoặc cạnh của ảnh sản phẩm giữ nguyên hình dạng và vị trí sản phẩm, trong khi prompt thay đổi bối cảnh, ánh sáng và phông nền.
• Dữ liệu tổng hợp: Tạo ảnh từ bản đồ phân đoạn hoặc pose sẽ tạo ra ảnh huấn luyện mới có nhãn được biết trước, vì bố cục đã được xác định trước khi tạo ảnh.

Huấn luyện ControlNet#

Huấn luyện cần các cặp ảnh và bản đồ điều kiện, được tạo tự động bằng cách chạy bộ trích xuất điều kiện trên một tập dữ liệu ảnh có chú thích. Trọng số bị khóa không cần cập nhật gradient, vì vậy mỗi vòng lặp huấn luyện chỉ tốn thêm khoảng 23% bộ nhớ GPU và 34% thời gian so với tối ưu hóa Stable Diffusion không có ControlNet; bài báo ghi nhận quá trình huấn luyện ổn định trên các tập dữ liệu từ dưới 50.000 đến hơn một triệu ảnh. Quá trình huấn luyện còn cho thấy hiện tượng "hội tụ đột ngột": model bỏ qua điều kiện trong hàng nghìn bước rồi bất ngờ bắt đầu làm theo điều kiện, thường trong vòng 10.000 bước tối ưu hóa; ở ví dụ trong bài báo, hiện tượng này xảy ra tại bước 6.133.