ControlNet

ControlNet — это архитектура нейронной сети, которая добавляет пространственное управление к предварительно обученной диффузионной модели для генерации изображений по тексту, например Stable Diffusion. Помимо текстового промпта, модель принимает изображение-условие, например карту границ, карту глубины, скелет человеческой позы или карту сегментации. Сгенерированное изображение следует заданной структуре, а промпт определяет содержание и стиль. Например, промпт «рыцарь в серебряных доспехах» в сочетании со скелетом позы создаёт рыцаря именно в этой позе. Лвминь Чжан, Аньи Рао и Маниш Агарвала из Стэнфордского университета представили ControlNet в 2023 году в статье «Добавление условного управления к диффузионным моделям генерации изображений по тексту». Статья разделила премию Марра за лучшую работу на ICCV 2023. Веса базовой модели остаются замороженными, поэтому ControlNet добавляет управление, не ухудшая уже имеющиеся знания модели.

Как работает ControlNet#

Stable Diffusion генерирует изображение, многократно устраняя шум в латентном представлении с помощью U-Net. ControlNet блокирует U-Net и создаёт обучаемую копию её кодировщика: 12 кодирующих блоков и средний блок. Небольшая свёрточная сеть сначала преобразует изображение-условие размером 512×512 в карту признаков размером 64×64, соответствующую латентному представлению. Обучаемая копия обрабатывает латентное представление вместе с признаками условия, а её выходы добавляются к 12 пропускающим соединениям замороженной U-Net и её среднему блоку. Так замороженный декодер получает управляющий сигнал:

Копия соединяется с замороженной моделью только через нулевые свёртки: свёртки 1×1 с изначально нулевыми весами и смещениями. На первом шаге обучения они выдают нули, поэтому объединённая модель ведёт себя точно так же, как исходная, и в предварительно обученные слои не поступает случайный шум. При этом выходная нулевая свёртка получает ненулевой градиент весов, поскольку её вход не нулевой. Поэтому уже на первом шаге она отходит от нуля, и управляющий сигнал начинает усиливаться. В этом коде на PyTorch показан один блок:

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}")

Он выводит True и zero_in grad: 0.0, zero_out grad: 2040.6. Выход блока идентичен выходу замороженной модели, внешняя нулевая свёртка получает градиент, поскольку её вход ненулевой, а внутренняя начинает обучаться, когда внешняя отходит от нуля.

Типы условий#

В исходной статье обучались отдельные модели ControlNet для границ Canny, карт глубины, карт нормалей поверхности, прямых линий M-LSD, мягких границ HED, карт сегментации ADE20K, поз человека OpenPose и пользовательских эскизов. Каждое изображение-условие извлекается из эталонной фотографии с помощью готовой модели, например детектора границ Canny из OpenCV или сети для оценки глубины. Одновременно можно применять несколько моделей ControlNet, задавая силу влияния каждой. Модели оценки позы Ultralytics YOLO26 извлекают подобные условия позы и предсказывают для каждого человека 17 ключевых точек COCO. Эти точки можно нарисовать в виде скелета на пустом холсте:

import cv2
import numpy as np
from ultralytics import YOLO

# Загрузи модель позы YOLO26, которая предсказывает 17 ключевых точек для каждого человека
model = YOLO("yolo26n-pose.pt")
results = model("https://ultralytics.com/images/bus.jpg")
result = results[0]

# Нарисуй только скелеты на чёрном холсте размером с фотографию
skeletons = result.plot(img=np.zeros_like(result.orig_img), boxes=False)
cv2.imwrite("pose_condition.png", skeletons)

На сохранённом изображении конечности каждого человека показаны цветными линиями на чёрном фоне; цвета, текстуры и фон фотографии удалены.

ControlNet и родственные методы#

• LoRA: Дообучает небольшой набор добавленных весов, чтобы научить модель новому стилю, персонажу или объекту. LoRA меняет то, что рисует модель, а ControlNet фиксирует расположение объектов. Эти методы часто используют вместе.
• T2I-Adapter: Более лёгкая альтернатива примерно с 77 миллионами параметров. Она добавляет признаки условия к признакам кодировщика замороженной U-Net, а не запускает полную копию кодировщика. Адаптеры для разных условий можно объединять во время генерации.
• Генерация «изображение в изображение»: Начинает удалять шум из зашумлённой копии входной фотографии, поэтому результат сохраняет её цвета и примерную компоновку. ControlNet ограничивает только извлечённую структуру, например границы, а цвета и текстуры задаёт промпт.
• Восстановление изображений: Заново генерирует замаскированную часть изображения так, чтобы она соответствовала остальному. Этот метод определяет, какая область изменится, а ControlNet — структуру сгенерированного содержимого.

Примеры применения в реальных условиях#

• Архитектура и дизайн интерьера: Эскизы, чертежи в САПР или карты линий комнаты M-LSD превращаются в фотореалистичные изображения с разными материалами и стилями при сохранении той же геометрии.
• Иллюстрации персонажей и анимация: Скелеты поз фиксируют позу персонажа при изменении стиля, а позы каждого кадра, извлечённые из эталонного видео, задают согласованные последовательности анимации.
• Изображения продуктов и маркетинговые материалы: Карты глубины или границ объекта сохраняют его форму и расположение, а промпт меняет обстановку, освещение и фон.
• Синтетические данные: Генерация изображений по картам сегментации или позам создаёт новые обучающие изображения с заранее известными метками, поскольку компоновка была задана до генерации.

Обучение ControlNet#

Для обучения нужны пары изображений и карт условий, которые можно автоматически получить, запустив извлекатель условий на наборе изображений с подписями. Замороженным весам не нужны обновления градиентов, поэтому каждая итерация обучения требует лишь примерно на 23% больше памяти GPU и на 34% больше времени, чем оптимизация Stable Diffusion без ControlNet. В статье показано, что метод надёжно обучается на наборах данных объёмом от менее чем 50 000 до более миллиона изображений. Также при обучении наблюдается «резкая сходимость»: на протяжении тысяч шагов модель игнорирует условие, а затем внезапно начинает ему следовать — обычно в пределах 10 000 шагов оптимизации, а в примере из статьи это происходит на шаге 6133.