ОминиКонтроль: Минимальный и универсальный контроль для модели Diffusion Transformer

В последние годы технологии генерации изображений с помощью искусственного интеллекта сделали огромный скачок вперёд, особенно благодаря использованию диффузионных моделей. Эти модели, такие как Diffusion Transformers (DiT), показали себя как лидеры в создании визуально реалистичных и разнообразных изображений, превосходя по качеству и многообразию даже такие известные методы, как Generative Adversarial Networks (GANs). Однако, несмотря на их успехи, одна из ключевых проблем остаётся нерешённой – это возможность точного и гибкого управления процессом генерации для удовлетворения разнообразных и сложных требований пользователей.

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

Существующие подходы

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

Пространственно-согласованные задачи (например, преобразование эскиза в изображение, заполнение недостающих областей изображения), где требуется прямая соответствие между условиями и выходным изображением. Такие методы, как ControlNet, вводят условные признаки в модель таким образом, чтобы сохранять пространственную структуру.
Непространственно-согласованные задачи (например, генерация изображений на основе определённых объектов, стилизация), где используются предобученные энкодеры, такие как CLIP, для извлечения глобальных признаков и их интеграции через механизмы кросс-внимания.

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

Специализация на конкретные типы задач, что требует использования разных методов для различных сценариев управления.
Значительное увеличение количества параметров модели из-за необходимости в дополнительных сетевых структурах.

Предлагаемый подход: ОминиКонтроль

Чтобы решить эти проблемы, мы представляем ОминиКонтроль, новый подход к интеграции управления на основе изображений в архитектуру Diffusion Transformers. Наш метод отличается следующими ключевыми особенностями:

Эффективность параметров: Мы используем существующий энкодер VAE модели для обработки условных изображений, что позволяет избежать необходимости в дополнительных сетевых структурах и сокращает количество параметров до 0.1% от общей модели.
Универсальность: ОминиКонтроль способен обрабатывать как пространственно-согласованные, так и непространственно-согласованные задачи в едином фреймворке, используя одинаковый процесс обработки токенов для всех типов условий.
Многоуровневое внимание: Наш метод позволяет токенам условий взаимодействовать с токенами шума через механизм многоуровневого внимания, что способствует эффективному обмену информацией и распространению сигналов управления.

Реализация

Мы реализовали наш метод на базе высокопроизводительной модели FLUX.1-dev, которая содержит 12 миллиардов параметров. Эксперименты показали, что наш подход превосходит как UNet-основанные реализации, так и их адаптации для модели FLUX.1 по ряду задач, включая генерацию изображений по краям, синтез с учётом глубины, редактирование определённых областей и генерацию изображений с сохранением идентичности объектов.

Датасет Subjects200K

Для обучения и оценки наших моделей мы разработали и выпустили датасет Subjects200K, содержащий более 200,000 изображений, которые сохраняют идентичность объектов при различных условиях. Этот датасет, наряду с эффективным пайплайном синтеза данных, предоставляет ценный ресурс для дальнейших исследований в области генерации изображений с сохранением идентичности объектов.

Заключение

ОминиКонтроль предлагает новый, параметрически эффективный и универсальный подход к управлению генерацией изображений с помощью Diffusion Transformers. Наши эксперименты подтвердили, что этот метод превосходит существующие решения по качеству и гибкости управления, открывая новые возможности для создания изображений с точным соответствием заданным условиям.