Kandinsky 3.0 — новая модель генерации изображений по тексту

«Без чувства современности художник останется непризнанным»,
Михаил Пришвин, русский и советский писатель, прозаик и публицист

В прошлом году на АI Journey мы представили модель Kandinsky 2.0 — первую диффузионную мультиязычную модель генерации изображений по тексту, которая может генерировать изображения на основе русскоязычного текста. За ней последовали новые версии — Kandinsky 2.1 и Kandinsky 2.2, которые значительно отличались по качеству и своим возможностям от версии 2.0, и стали для нашей команды серьёзными вехами на пути к достижению лучшего качества генерации.

Но бесконечность не предел, и всегда есть куда расти. Количество научных работ и инженерных решений, касающихся генерации той или иной модальности, увеличивается, возникают новые постановки задач генерации — сейчас это в том числе генерация видео, 3D и (даже) 4D по тексту. Область генеративного обучения захватывает всё большие просторы в информационном поле. Недавно наша команда представила подход Deforum-Kandinsky, который позволяет получить анимированные видеоролики на основе базовой модели генерации изображения по тексту, и подход создания zoom-in / zoom-out видео на основе модели Kandinsky Inpainting. Также параллельно с этим релизом мы представляем первую в России end-to-end модель генерации видео по тексту Kandinsky Video.

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

Итак, спустя год после релиза нашей первой диффузионной модели мы представляем новую версию модели генерации изображений по тексту — Kandinsky 3.0! Это результат длительной работы нашей команды, которую мы вели параллельно с разработками версий Kandinsky 2.1 и 2.2. Мы провели много экспериментов по выбору архитектуры и проделали большую работу с данными, чтобы сделать понимание текста и качество генераций лучше, а саму архитектуру — проще и лаконичнее. Также мы сделали нашу модель более «отечественной»: теперь она значительно лучше ориентируется в российском и советском культурном поле.

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

Архитектура модели и обучение

Kandinsky 3.0 представляет собой диффузионную модель для генерации изображений по тексту (как и все модели линейки Kandinsky 2.X). Цель обучения любой диффузионной модели заключается в том, чтобы при обратном процессе диффузии научиться восстанавливать реальное изображение, которое было зашумлено во время процесса прямой диффузии (подробнее можно прочитать, например, тут). При обучении Kandinsky 3.0 мы отошли от концепции двухстадийной генерации, которая использовалась в предыдущих версиях. Напомню, что её суть заключалась в том, что мы обучаем две модели:

• Diffusion Mapping (Image Prior) — генерирует латентный картиночный вектор, полученный кодированием картинки визуальной частью модели CLIP, принимающей на вход латентный текстовый вектор, закодированный текстовой частью модели CLIP.

• Decoder (U-Net) — генерирует картинку (если быть более точным, то ) по латентному картиночному вектору CLIP.

Подробнее про устройство архитектуры версии Kandinsky 2.2 можно прочитать, например, вот тут.

В Kandinsky 3.0 генерация изображения происходит непосредственно из закодированных текстовых токенов. Такой подход позволяет упростить обучение, так как теперь необходимо обучать лишь одну часть модели (а именно Decoder). Также этот подход сильно улучшает понимание текста, ведь теперь мы можем использовать не текстовый энкодер CLIP, который обучался на довольно примитивных текстах (по факту описаниях изображений из сети Интернет), сильно отличающихся от естественного языка, а взять мощную языковую модель, обученную на большом корпусе качественных текстовых данных.

Помимо обновления подхода к кодированию текста мы провели очень большое исследование архитектуры U-Net, отвечающей за удаление шума с картинки. Главная дилемма состояла в том, какой тип слоёв будет содержать основную часть параметров сети: трансформерные или свёрточные. При обучении на больших объёмах данных трансформеры показывают лучшие результаты при работе с изображениями, однако U-Net-архитектуры практически всех диффузионных моделей преимущественно свёрточные. Чтобы разрешить эту дилемму, мы проанализировали различные архитектуры и отметили для себя следующие модели:

• ResNet-18, ResNet-50 — всем известная архитектура, но с одной примечательной особенностью. Свёрточные блоки в маленькой версии отличаются от свёрточных блоков в большой наличием bottleneck, отвечающего за сжатие числа каналов перед обработкой тензора свёрткой 3x3. Это позволяет уменьшить количество параметров, а следовательно увеличить глубину сети, что на практике даёт лучшие результаты при обучении.

• CoAtNet — архитектура, совмещающая в себе как свёрточные блоки, так и блоки внимания. Основная её идея состоит в том, что на начальном этапе картинку стоит обрабатывать локальными свёртками, а уже сжатое её представление — слоями трансформера, обеспечивающими глобальное взаимодействие элементов изображения.

• MaxViT — архитектура, уже практически полностью основанная трансформаторных блоках, но адаптированная для работы с картинками путём уменьшения квадратичной сложности self-attention.

Идея использовать модели классификации была навеяна тем, что много хороших архитектурных решений берутся из моделей, которые показывают высокое качество на бенчмарке ImаgeNet. Однако в ходе наших экспериментов выяснилось, что перенос качества работает неоднозначно. Лучшая на задаче классификации, из представленных выше, архитектура MaxViT показывает уже не столь хорошие результаты в задаче генерации (после «превращения» её в U-Net). Исследовав таким образом все вышеперечисленные архитектуры, мы остановились на блоке ResNet-50 как основном блоке U-Net. При этом мы дополнили его ещё одним свёрточным слоем с ядром 3x3, позаимствовав эту идею из статьи BigGAN.

В итоге архитектура Kandinsky 3.0 получилась из трёх основных частей:

• FLAN-UL2 — языковая Encoder-Decoder-модель. Для кодирования текста мы брали только Encoder, составляющий половину параметров всей архитектуры. Помимо предобучения на корпусе текстов, эта версия была также дообучена в стиле SFT на большом корпусе языковых задач. Наши эксперименты показали, что такое дообучение существенно улучшает генерацию картинок. Во время обучения картиночной части модели языковая модель была полностью заморожена.

• U-Net с архитектурой, проиллюстрированной ниже, состоящий преимущественно из блоков BigGAN-deep. Это позволило сделать архитектуру в два раза глубже в сравнении с иными диффузиями, основанными на обычных блоках BigGAN, при сохранении того же количества параметров.

В качестве автоэнкодера мы использовали Sber-MoVQGAN, который хорошо зарекомендовал себя в предыдущих версиях. Подробнее про него прочитать можно в отдельной статье. При этом была взята самая большая версия модели с 270 млн параметров, которая показывала лучшие результаты в задаче восстановления изображения.

Сравнение модели Kandinsky 3.0 с моделями Kandinsky 2.2 и SDXL по количеству параметров приведено в таблице ниже:

Данные

В ходе обучения (как обычно) использовался большой датасет пар «текст-картинка», собранных в интернете. Эти данные проходили многочисленные фильтры: эстетичности картинки, соответствия картинки тексту, наличия дубликатов, фильтры разрешения и соотношения сторон. По сравнению с Kandinsky 2.2 мы расширили используемые наборы данных, обогатив их новыми данными, добавили изображения, содержащие российские сущности, а также изображения, описания к которым были сгенерированы при помощи современных мультимодальных моделей.

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

• На первой стадии модель училась на изображениях с разрешением 256 × 256 — это 1,1 млрд пар текст-картинка (размер батча 20, 600 тыс шагов, 100 Tesla A100).

• Затем для обучения использовались изображения в разрешении 384 × 384 — это 768 млн пар текст-картинка (размер батча 10, 500 тыс шагов, 100 Tesla A100).

• После брались изображения в разрешении 512 × 512 — это 450 млн пар текст-картинка (размер батча 10, 400 тыс шагов, 100 Tesla A100).

• Далее аналогичным образом — в обучение попадали изображения в разрешении 768 × 768 — это 224 млн пар текст-картинка (размер батча 4, 250 тыс шагов, 416 A100).

• На финальной стадии обучения модель училась на изображениях в смешанном разрешении: 768 ≤ ширина, высота ≤ 1024 — таких пар набралось 280 млн (размер батча 1, 350 тыс шагов, 416 A100).

Результаты и примеры генераций

Ниже приведу несколько избранных примеров генераций (по текстам в разных доменах) модели Kandinsky 3.0.

Результаты сравнения Kandinsky 3.0 с другими известными text-to-image моделями

Для сравнения моделей мы собрали свой сбалансированный набор из 2100 промптов по 21 категории и сравнили различные версии модели Kandinsky 3.0, чтобы выбрать лучший. Для этого мы провели три этапа side-by-side (SBS) сравнений, в которых принимали участие 28 разметчиков.

Затем, когда наилучшая версия модели Kandinsky 3.0 была выбрала, мы провели side-by-side сравнение с моделью Kandinsky 2.2. В исследовании участвовало 12 человек, которые проголосовали суммарно 24 800 раз. Для этого был разработан бот, который каждый демонстрировал одну из 2100 пар изображений. Каждый человек выбирал лучшее изображения по двум критериям:

1. соответствие картинки текстовому описанию,

2. визуальное качество картинки.

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

Примеры изображений с side-by-side сравнения:

Одно из основных преимуществ модели Kandinsky 3.0 состоит в улучшенном понимании текстового запроса (заметная разница с Kandinsky 2.2). Кроме того, кратно улучшилось качество понимания запросов на темы отечественного культурного кода, например:

Ниже также представлены примеры генераций популярных моделей в сравнении с Kandinsky 3.0:

Модель Kandinsky Inpainting

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

Для обучения inpainting-части модели за основу взяли подход из GLIDE, который уже был ранее реализован в моделях семейства Kandinsky 2.X, а также в семействе моделей Stable Diffusion: входной слой U-Net модифицируется так, чтобы на вход дополнительно можно было подать image latent и mask. Таким образом, U-Net принимает на вход целых 9 каналов: 4 для исходного latent, 4 для latent-изображения и дополнительный канал для маски. С точки зрения модификаций, в целом, это всё — далее обучение ничем не отличается от обучения стандартной диффузионной модели.

Важной особенностью задачи является то, как генерируются маски и какой текст используется для обучения. Пользователи могут нарисовать либо маску кисточкой, либо новое изображение через outpаinting. Чтобы учесть то, как работает пользователь, при обучении мы создали специальные маски, которые имитируют его поведение: маски произвольной формы, маски, нарисованные кисточкой, маски различных реальных объектов, а также разные формы масок, использованные исследователями Adobe при разработке Generative Fill. С коллегами, которые разрабатывали последнюю модель, мы даже консультировались и смогли переняли их лучшие практики :)

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

Примеры Inpainting

Примеры Outpainting

Kandinsky Animation

Мы также адаптировали модель Kandinsky 3.0 для использования в пайплайне Deforum — таким образом, с помощью Kandinsky 3.0 теперь можно генерировать не только изображения, но и анимации. Подробнее с принципом работы режима анимаций можно познакомиться в одной из наших недавних статей.

Основной сложностью интеграции новой модели стало различие в режимах добавления шума в процессе диффузии: у Kandinsky 2.2 шум добавляется по линейному закону (картинка сверху), в то время как у Kandinsky 3.0 – по косинусному (картинка снизу). Эта особенность потребовала некоторого количества экспериментов для успешной адаптации.

Примеры анимаций

Отметим, что разрешение изображений, а также количество кадров в секунду снижено в связи с техническими требованиями Хабра.

Вывод и планы

Мы представили нашу новую архитектуру генерации изображений по тексту — Kandinsky 3.0. Значительно улучшили понимание текста и отечественного культурного кода в сравнении с нашими предыдущими моделями и обязательно будем продолжать двигаться в этом направлении. Что касается научной стороны, в наши планы входит создание ещё одной новой модели генерации, которая должна будет сказать своё новое слово на международной AI-арене.

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

Авторы и их вклад

Модель Kandinsky 3.0 разработана командой Sber AI при партнёрской поддержке учёных из Института искусственного интеллекта AIRI на объединённых датасетах Sber AI и компании SberDevices.

Коллектив авторов: Владимир Архипкин, Андрей Филатов, Анастасия Мальцева, Арсений Шахматов, Саид Азизов, Вячеслав Васильев, Игорь Павлов, Михаил Шойтов, Юлия Агафонова, Сергей Нестерук, Анастасия Лысенко, Илья Рябов, Ангелина Куц, Татьяна Паскова, София Кирилова, Михаил Мартынов, Сергей Марков, Андрей Кузнецов и Денис Димитров.

Каналы авторов:

• CompleteАI

• DendiMath

• Градиентное погружение

• AwesomeDL

• Канал Сергея Маркова

Контакты для коммуникации

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

Полезные ссылки

• fusionbrain.ai (доступна генерация по тексту и режимы inpainting/outpainting; а также генерация анимации и видео, о последней новинке можно почитать в статье)

• Telegram-bot (доступна генерация по тексту)

• rudalle.ru (доступна генерация по тексту)

• Github (доступен код модели генерации изображений и инпейнтинга)

• HuggingFace (доступны веса модели генерации изображений и инпейтинга)