Добавляем контроль к диффузии

Введение

В прошлом посте мы рассмотрели разные методы файнтюна модели Stable Diffusion, которые помогают добавлять новые стили или объекты в процесс генерации. Они могут быть полезны для получения более качественных результатов и расширения знаний модели в какой-то области.

Да, результаты генерации модели устраивают нас. Но что делать, если хочется управлять генерацией? Как добиться того, чтобы сами генерации были похожи на референс (например, соответствовали скетчу ребенка или позе человека на фотографии), или по-другому, чтобы сгенерированный контент отвечал дополнительному условию?

Умение генерировать изображение по дополнительному условию (кроме текста) — полезное умение. Оно позволяет ограничивать генеративную способность модели и держать ее в нужных нам рамках, получать вариации и быстрее выбирать наилучшую из них.

В этом посте мы обсудим способы, с помощью которых можно добавлять дополнительные условия в генерацию диффузионными моделями, а также определим достоинства и недостатки каждого из них.

Виды дополнительных условий

Сначала рассмотрим полезные для генерации дополнительные условия. Все виды перечислить будет сложно, поэтому давайте пройдемся по основным группам.

Скетчи

Счетчи — наброски рисунков в виде черно-белого изображения: границы объектов, прямые линии, произвольные рисунки и др.

Частично их можно создавать самостоятельно от руки. Но также есть автоматические алгоритмы (например, Canny Edge, HED Boundary).

Поза

Контроль генерации позой персоны особенно важен при генерации людей или персонажей.

Поза может быть автоматически найдена предобученными нейросетевыми моделями типа OpenPose (в том числе с модификациями для нахождения точек рук и лица).

Карта глубины

Карта глубины — черно-белое изображение с градациями серого, которое кодирует удаленность содержимого сцены от камеры.

Она применяется, чтобы придать 2D-изображению глубину (создать 3D-изображение).

Рисунок 5. Пример визуализации 2D-изображения в 3D с помощью предсказанной карты глубины

Генерация изображений по карте глубины позволяет частично сохранять композицию сцены и положение объектов, в каком-то смысле осуществлять полную перерисовку имеющегося изображения в конкретной стилистике.

Дополнительно рекомендуем ознакомиться с инструкцией для генерации изображения с помощью карты глубины от civitai.com.

Карта нормалей

Карта нормалей схожа с картой глубины, поскольку также предоставляет информацию о свойствах поверхности в 3D (для каждого пикселя задается вектор из 3-х значений, который формирует направление точки, перпендикулярное поверхности в точке).

Карта нормалей может быть получена как отдельной моделью, так и через преобразования из карты глубины. С помощью ее свойств мы также частично сохраняем композиционные составляющие сцены при генерации.

Bounding box для объектов

Одно из удобных условий — bounding box’ы с указанием позиции объектов на изобрежении.

Оно позволяет достаточно свободно описывать композиционные составляющие сцены.

Сегментация

Конечно, нельзя пройти мимо сегментации — одного из самых важных условий генерации.

Идея применения сегментации в генеративных моделях не новая: этот подход активно использовался в моделях pix2pix/HD, GauGAN 1/2.

Рисунок 9. Пример использования сегментации как условия в генеративной модели GauGAN

Изображения

Предыдущие виды условий так или иначе подразумевают работу со входной картинкой: ее предобработку через создание скетча, исполнение предобученных моделей или создание таких условий с помощью других инструментов.

Однако часто описать то, что мы хотим получить, довольно трудно. Например, мы хотим сгенерировать снимок, схожий с другим по содержанию или стилистике, но отличающийся от него в деталях.

Используя генеративную модель, мы также можем решать задачу колоризации через черно-белое изображение в качестве условия.

Другое интересное применение — генерация с лицом в качестве условия. Пожалуй, это одно из самых сложных условий: на текущий момент неплохо передаются общие черты лица, но добиться полной схожести достаточно трудно, особенно если менять выражение лица генерации.

Класс изображений в качестве условий активно используется на практике.

Важно отметить: для более широкой практической пользы нам бы хотелось иметь возможность комбинировать несколько условий в рамках одной генерации (например, позу и лицо).

Методы добавления условий для диффузионных моделей

Итак, мы разобрались с условиями. Давайте теперь поговорим о механизмах их внедрения. Условно они разделяются на несколько групп в зависимости от принципа работы:

• дообучение модели;
• модификация текстового промпта;
• “добавки” между блоками Unet-части модели;
• модификация механизма Cross-Attention.

Рассмотрим каждый принцип по порядку 😉

Дообучение модели

Самый очевидный способ добавления условия — дообучение имеющегося Unet, чтобы тот принимал его в виде одного или нескольких входных каналов.

Для этого достаточно переопределить входной сверточный слой conv_in в Unet (слой, который принимает латенты с 4-мя каналами) и увеличить количество входных каналов. Веса нового слоя можно проинициализировать либо случайно, либо так, чтобы слой в самом начале не использовал новое условие.

При дообучении модели веса данного слоя подстраиваются, а оставшиеся веса используют подаваемое условие. Однако процесс может потребовать много ресурсов. Но у нас есть возможность их снизить с помощью LoRA-подходов [см. наш пост].

Яркие представители такого подхода — Stable Diffusion Depth и Stable Diffusion Inpainting. Эти модели получены из txt2image моделей с использованием файнтюна.

Например, Depth модель принимает на вход латент, как и раньше, +1 канал, в котором закодирована карта глубины (суммарно 5 каналов), карта глубины при этом интерполируется до разрешения латента. Inpaint модель же принимает на вход целых 9 каналов (4 канала от текущего зашумленного латента, 1 канал с интерполированной до размеров латента маской inpaint, 4 канала с замаскированным латентом от чистой картинки).

Хотя данный подход достаточно мощный, он требует значительно больше вычислительных ресурсов по сравнению с другими методами, поэтому не является популярным. Еще он хорошо работает при использовании легковесных условий (1 канал — скетч, сегментация, карта глубины, карта нормалей). Однако при применении условий типа изображения (3 канала изображения или 4 канала латента) трудно получить парные данные (здесь выступают другие способы).

Модификация текстового промпта

Один из способов добавления дополнительных условий — их пробрасывание внутрь процесса кодирования промпта. Самое простое такое применение — добавление токенов, которые кодируют позицию (bounding box) интересующих объектов внутрь промпта.

Reco

Авторы Reco предлагают дообучить новые эмбеддинги, кодирующие нормализованную позицию объектов на изображении через параметры bounding box. Относительные координаты квантизируются, и при кодировании используются полученные значения, они вставляются перед объектом в промпт.

Для обучения авторы используют 2 возможных подхода: применяют датасет COCO, а также строят свой небольшой размеченный датасет, подвыборку LAION.

“Добавки” между блоками Unet-части модели

Один из самых популярных способов проброса дополнительного условия в Unet — дополнительные добавки на выходах блоков внутри Unet (на выходах down-блоков внутри encoder’а / на skip-connection между encoder’ом и decoder’ом).

Яркие представители метода:

• ControlNet (и его модификации — ControlNet-XS, Control-LoRA, ControlNet++) и UniControl
• T2I-Adapter

Давайте рассмотрим каждый из этих подходов подробнее.

ControlNet

ControlNet [статья] — достаточно простой и при этом очень мощный подход, который позволяет добавлять пространственный контроль к предобученной диффузионной модели.

Он базируется на следующих принципах:

• переиспользование диффузионного Unet-энкодера;
• применение подхода zero-convolution при инициализации весов, связанных со смешиванием условия и текущих признаков.

Для кодирования входного условия используется легкая обучаемая сверточная сеть с простой архитектурой (свертки с ядрами размером 4×4 и stride 2×2). Она переводит условие из изображения в размерность, равную размерности латента \( c_f \).

Полученное закодированное условие обрабатывается через zero-convolution (свертка 1×1 с нулевыми изначальными весами), и результат прибавляется к текущему закодированному латенту. Инициализация нулями такой свертки помогает добиться отсутствия изначального влияния входного сигнала: при тренировке веса свертки подстраиваются нужным образом (ведь после свертки отсутствует нелинейность).

Такая техника позволяет не изменять поведение модели, следовательно, добиваться лучших результатов при меньших затратах.

Полученное общее условие \( z_t + \mathbb{Z_{conv}}(c_f,\theta_{z_1}) \) затем подается в отдельный Unet-энкодер, который является полной тренируемой копией диффузионного энкодера. Каждый его блок использует закодированный text condition, как и обычный энкодер от Stable Diffusion. Выходы всех down и mid-блоков сохраняются, каждый выход отдельно обрабатывается через zero-convolution и прибавляется к фичам в декодер-блоке.

Другими словами, ControlNet — дополнительный энкодер с архитектурой Stable Diffusion, выходы которого являются добавками для декодера Stable Diffusion.

Тренировка модели не отличается от тренировки обычной диффузионной модели: мы предсказываем добавленный на изображение шум. Однако во время предсказания используется внешнее условие, которое может быть предварительно посчитано для изображения. При тренировке, как и при тренировке диффузии, текстовое условие периодически обнуляется, чтобы модель научилась работать без него.

Благодаря zero-convolution модель способна сразу предсказывать высококачественные изображения (поскольку изначального влияния ControlNet нет). Авторы отмечают: диффузионная модель с ControlNet не выучивает, как использовать условие постепенно, а неожиданно выдает результаты, схожие с поданным условием. Этот феномен называют “внезапной сходимостью”.

Также мы получаем возможность смешивать разные условия между собой за счет перевзвешивания суммы добавок от разных ControlNet. Такой подход называется MultiControlNet. Он позволяет добиваться еще более интересных результатов.

Стоит отметить: ControlNet вызывается по умолчанию на каждой итерации вместе с диффузионным Unet, что сильно замедляет процесс генерации. На практике мы видим, что можно ограничивать шаги, на которых он вызывается (на более поздних шагах генерации влияние ControlNet минимально).

Еще один важный вопрос: является ли архитектура ControlNet (с encoder’ом от диффузионного Unet) оптимальной? Или аналогичных результатов можно добиться за счет более легковесных подходов?

В статье авторы рассматривают zero-convolution и сравнивают результаты с ControlNet-Lite, использующим легковесные свертки вместо тяжеловесного Unet-энкодера.

Полученные метрики отмечают важность zero-convolution и переиспользования энкодера от диффузионного Unet.

Также рекомендуем ознакомиться с дискуссией на GitHub: автор подсвечивает важность диффузионного энкодера, когда генерация производится без текстового промпта.

Вдохновившись ControlNet, сообщество разработало ряд смежных методов.

Control-LoRA

Например, Control-LoRA предлагает соединить идеи ControlNet и легковесность LoRA-подхода. Здесь вместо обучения отдельного энкодера на условие обучается LoRA поверх диффузионного Unet. Несмотря на отсутствие статьи, подход успешно переиспользовался Stability AI.

ControlNet-XS

ControlNet-XS предлагает глубже проанализировать ControlNet с точки зрения архитектуры. Одна из ключевых причин, по которой мы не можем заменить диффузионный энкодер на что-то более легковесное, — необходимость энкодера ControlNet обладать генеративной силой, чтобы успешно добавлять информацию в декодер.

Здесь можно привести такую аналогию: генеративная модель в данной ситуации — автомобиль, который движется куда-то, а ControlNet — система навигации. И системе навигации в классическом ControlNet’е необходимо предугадывать, как и куда будет двигаться машина в следующий промежуток времени.

Но решить проблему можно — для этого энкодер ControlNet’а и диффузионный энкодер должны обмениваться информацией в процессе работы!

Авторы рассматривают три возможные архитектуры (см. Type A, Type B и Type C на рисунке ниже) и экспериментально доказывают — Type B работает лучше с точки зрения качества и количества параметров.

Энкодером в данном случае выступает облегченная версия диффузионного Unet, тренируемая с нуля (размер энкодера при этом может варьироваться от 1.7 до 55 миллионов параметров). Также при обучении используется механизм zero-convolution для пробрасывания информации между энкодером и декодером.

Авторы представляют преимущество этого подхода как в качестве генерации, так и в скорости его работы в сравнении с ControlNet (получаем ускорение почти в 2 раза).

ControlNet++

ControlNet++ — недавняя работа, дальнейшее развитие базового ControlNet. Авторы утверждают: ControlNet не всегда выдает результаты, которые идеально соответствуют нашему условию. Поэтому они предлагают добавить идеи Cycle-Consistency loss’а в тренировку ControlNet’а, что позволяет ControlNet’у лучше соответствовать условию.

UniControl

UniControl — интересное развитие идей MultiControlNet’а. Эмпирические наблюдения показывают: соединение большого количества контролей через ControlNet приводит к более негативным результатам. Частично это объясняется тем, что направления сдвигов в шуме не обязательно коррелируются, а значит, могут мешать друг другу.

Вторая проблема — применение большого количества ControlNet’ов требует многократного прогона Unet-энкодеров (по одному на каждый ControlNet), что приводит к замедлению генерации, плюс требует хранения весомого количества параметров.

Для борьбы с этими проблемами авторы предлагают построить один ControlNet, решающий сразу несколько задач.

Отдельно для каждого типа условий строится своя небольшая сверточная модель, результаты которых объединяются в стиле Mixture-Of-Experts.

Важен и дополнительный HyperNet, куда подается текстовое описание выполняемой задачи (например, “depth map to image”, “segmentation map to image” и др.). CLIP text embedding обрабатывается через дополнительно обучаемые слои, и данный сигнал модулирует одно / несколько закодированных условий в стилистике StyleGAN2.

Авторы строят UniControl на датасете с 9 задачами:

1. Hed edge to image;
2. Canne edge to image;
3. Segmentation map to image;
4. Depth map to image;
5. Normal surface to image;
6. Human pose to image;
7. Sketch to image;
8. Bbox to image;
9. Image outpainting.

Подобная архитектура после обучения позволяет производить Zero-Shot генерацию на новые задачи, которых не было в обучающей выборке.

Для этого авторы предлагают присваивать веса в MOE Adapter вручную, либо по их косинусной близости с эмбеддингами задач, на которых обучалась модель. В статье приведены 3 Zero-Shot задачи (их не было при обучении исходной модели): image deblurring, колоризация и inpaint.

T2I-Adapter

T2I-Adapter — более легковесная альтернатива ControlNet, в основе которой лежит схожий принцип добавок в фичи Unet. Статья опубликована почти на неделю позже статьи по ControlNet. Возможно, это повлияло на меньшую популярность данного подхода.

Вместо переиспользования диффузионного энкодера T2I-Adapter предлагает обучать новый легковесный энкодер (со структурой Pixel-Unshuffle + Resnet блоки). Выход каждого блока — “добавка” для down-блока диффузионного Unet-энкодера.

Важно отметить: в отличие от ControlNet’а T2I-Adapter использует “добавки” для энкодера, а не для декодера. А еще авторы не применяют механизм zero-convolution, а пробрасывают сигнал простым сложением.

Несмотря на кажущуюся простоту, подход вполне сопоставим по качеству с ControlNet. У него есть преимущество: добавки в down-блоки диффузионного Unet-энкодера считаются один раз, а не на каждом шаге диффузии. Это позволяет сэкономить много времени в отличие от ControlNet.

Модификация механизма Cross-Attention

С точки зрения результирующего качества работы самый удачный подход — модификация механизма Cross-Attention в модели таким образом, чтобы он обуславливался еще на один внешний сигнал с условием.

Часто к Self-Attention (для visual token’ов) и Cross-Attention (для соединения visual token’ов и внешнего условия, обычно — закодированного текстового промпта) добавляют Cross-Attention блок, который обрабатывает текущие visual token’ы и смешивает их с еще одним условием. Это позволяет доучивать только добавленные блоки, а все оставшиеся веса модели остаются фиксированными.

GLIGEN

Авторы GLIGEN в первую очередь фокусируются на добавлении bounding box объектов в качестве условия. При этом они показывают, как можно использовать другие условия (например, позу, карту глубины, HED map, семантическую карту).

Для кодирования названия объектов, соответствующих bounding box, используется CLIP text encoder. Его выходы вместе с закодированными fourier координатами bounding box смешиваются через MLP и в дальнейшем используются в новых Self-Attention блоках.

С целью кодирования других типов информации авторы используют следующие подходы:

• для изображения в качестве условия — CLIP image encoder как кодировщик информации;
• для ключевых точек — 2 координаты ключевых точек вместо 4-х в bbox координатах;
• для других пространственно согласованных условий (например, семантическая карта, depth map) — легковесную сверточную модель. Она закодирует условие в размерность, на которой будет оперировать модель. Авторы статьи отмечают: конкатенация условия в первый сверточный слой (как это было описано в секции с дообучением модели) сильно ускоряет сходимость.

Полученное закодированное условие предлагается сконкатенировать с текущими visual token’ами и обработать новым блоком с gated self-attention.

\( v = v + \beta \cdot \tanh\left(\gamma\right) \cdot TS \left( SelfAttn\left( [v, h^e] \right) \right) \)

где \( TS(\cdot) \) — операция token-selection, которая выбирает только визуальные токены, \( \gamma \) — обучаемый скаляр, инициализируемый нулем (чтобы изначально блок ничего не делал), \( \beta \) отвечает за силу эффекта и выставляется в 1 по умолчанию.

Также авторы показывают: применение новых блоков на всех этапах диффузионного процесса приводит к ухудшению результатов генерации. Поэтому они предлагают выставлять $\beta$ на более поздних шагах в 0.

IP-Adapter

Идея IP-Adapter в целом схожа с GLIGEN, но авторы здесь больше фокусируются на использовании изображения в качестве условия.

Для кодирования image prompt авторы предлагают использовать CLIP image encoder — его выходы обрабатываются через несколько линейных слоев и подаются в новые Cross-Attention блоки в качестве условия.

Задача новых Cross-Attention блоков — предсказать “добавку” в visual token’ы, которая будет учитывать внешний сигнал от закодированного изображения. В качестве visual token’ов используются те же токены, что и в Cross-Attention блоках. Они предсказывают текстовым эмбеддингом в качестве условия. Авторы называют этот подход Decoupled Cross-Attention и утверждают, что он работает лучше обычного метода.

При обучении IP-Adapter’а, как и при обучении обычной диффузионной модели, с какой-то вероятностью image prompt обнуляется, чтобы блоки научились работать без него.

Также авторы рассматривают 2 подхода к использованию эмбеддингов от CLIP encoder’а (до pooling слоя — fine-grained embedding, после — global image embedding) и показывают, что использование фичей до pooling слоя приводит к более согласованным с промптом результатам.

Также авторы показывают: получаемый IP-Adapter остается совместимым с остальной экосистемой Stable-Diffusion (например, с ControlNet) и позволяет комбинировать его с другими подходами.

Полученный адаптер способен работать и в других режимах диффузии: img2img и inpaint.

Таким образом, IP-Adapter — достаточно гибкий подход, совместимый с остальной экосистемой Stable Diffusion и способный работать в разных режимах без потери качества.

Заключение

Сообщество разработало множество методов расширения предобученных диффузионных моделей для лучшего следования различным дополнительным условиям.

Большая часть таких подходов не требует полного переобучения модели, что позволяет использовать их как plug-and-play решения между чекпоинтами.

А это, в свою очередь, помогает людям тратить меньше времени на генерации, экономить финансы многим компаниям и добиваться более качественных результатов генерации.