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.
Подробнее про устройство архитектуры версии Kandinsky 2.2 можно прочитать, например, вот тут.
В Kandinsky 3.0 генерация изображения происходит непосредственно из закодированных текстовых токенов. Такой подход позволяет упростить обучение, так как теперь необходимо обучать лишь одну часть модели (а именно Decoder). Также этот подход сильно улучшает понимание текста, ведь теперь мы можем использовать не текстовый энкодер CLIP, который обучался на довольно примитивных текстах (по факту описаниях изображений из сети Интернет), сильно отличающихся от естественного языка, а взять мощную языковую модель, обученную на большом корпусе качественных текстовых данных.
Архитектура модели Kandinsky 3.0
Помимо обновления подхода к кодированию текста мы провели очень большое исследование архитектуры U-Net, отвечающей за удаление шума с картинки. Главная дилемма состояла в том, какой тип слоёв будет содержать основную часть параметров сети: трансформерные или свёрточные. При обучении на больших объёмах данных трансформеры показывают лучшие результаты при работе с изображениями, однако U-Net-архитектуры практически всех диффузионных моделей преимущественно свёрточные. Чтобы разрешить эту дилемму, мы проанализировали различные архитектуры и отметили для себя следующие модели:
ResNet-18, ResNet-50 — всем известная архитектура, но с одной примечательной особенностью. Свёрточные блоки в маленькой версии отличаются от свёрточных блоков в большой наличием bottleneck, отвечающего за сжатие числа каналов перед обработкой тензора свёрткой 3×3. Это позволяет уменьшить количество параметров, а следовательно увеличить глубину сети, что на практике даёт лучшие результаты при обучении.
Характерные блоки ResNet-18 (слева) и ResNet-50 (справа)
CoAtNet — архитектура, совмещающая в себе как свёрточные блоки, так и блоки внимания. Основная её идея состоит в том, что на начальном этапе картинку стоит обрабатывать локальными свёртками, а уже сжатое её представление — слоями трансформера, обеспечивающими глобальное взаимодействие элементов изображения.
Архитектура CoAtNet
MaxViT — архитектура, уже практически полностью основанная трансформаторных блоках, но адаптированная для работы с картинками путём уменьшения квадратичной сложности self-attention.
Архитектура MaxViT
Идея использовать модели классификации была навеяна тем, что много хороших архитектурных решений берутся из моделей, которые показывают высокое качество на бенчмарке ImаgeNet. Однако в ходе наших экспериментов выяснилось, что перенос качества работает неоднозначно. Лучшая на задаче классификации, из представленных выше, архитектура MaxViT показывает уже не столь хорошие результаты в задаче генерации (после «превращения» её в U-Net). Исследовав таким образом все вышеперечисленные архитектуры, мы остановились на блоке ResNet-50 как основном блоке U-Net. При этом мы дополнили его ещё одним свёрточным слоем с ядром 3×3, позаимствовав эту идею из статьи BigGAN.
В итоге архитектура Kandinsky 3.0 получилась из трёх основных частей:
FLAN-UL2 — языковая Encoder-Decoder-модель. Для кодирования текста мы брали только Encoder, составляющий половину параметров всей архитектуры. Помимо предобучения на корпусе текстов, эта версия была также дообучена в стиле SFT на большом корпусе языковых задач. Наши эксперименты показали, что такое дообучение существенно улучшает генерацию картинок. Во время обучения картиночной части модели языковая модель была полностью заморожена.
U-Net с архитектурой, проиллюстрированной ниже, состоящий преимущественно из блоков BigGAN-deep. Это позволило сделать архитектуру в два раза глубже в сравнении с иными диффузиями, основанными на обычных блоках BigGAN, при сохранении того же количества параметров.
Архитектура U-Net (Kandinsky 3.0)
В качестве автоэнкодера мы использовали Sber-MoVQGAN, который хорошо зарекомендовал себя в предыдущих версиях. Подробнее про него прочитать можно в отдельной статье.
Сравнение с конкурентами по размеру:
Kandinsky 2.2 | SD-XL | Kandinsky 3.0 | |
Тип модели | Latent Diffusion | Latent Diffusion | Latent Diffusion |
Количество параметров | 4,6 млрд | 3,33 млрд | 11,9 млрд |
Текстовый энкодер | 0,62 млрд | 0,8 млрд | 8,6 млрд |
Diffusion Mapping | 1,0 млрд | - | - |
U-Net | 1,2 млрд | 2,5 млрд | 3,0 млрд |
MoVQGAN | 0,08 млрд | 0,08 млрд | 0,26 млрд |
Объём датасета | 1,5 млрд пар | 1,5 млрд пар | 1,5 млрд пар |
Данные
В ходе обучения использовалось много пар «текст-картинка», собранных в интернете. Эти данные проходили многочисленные фильтры: эстетичности картинки, соответствия картинки тексту, наличия дубликатов, разрешения и соотношения сторон. По сравнению с Kandinsky 2.2 мы расширили используемые наборы, обогатив их новыми данными, добавили русские сущности, а также изображения, описания к которым были сгенерированы при помощи современных мультимодальных моделей.
Процесс обучения был разделён на несколько стадий, что позволило использовать больше обучающих данных, а также генерировать изображения различного размера.
256 × 256: 1,1 млрд пар текст-картинка, размер батча 20, 600 тыс шагов, 100 A100.
384 × 384: 768 млн пар текст-картинка, размер батча 10, 500 тыс шагов, 100 A100.
512 × 512: 450 млн пар текст-картинка, размер батча 10, 400 тыс шагов, 100 A100.
768 × 768: 224 млн пар текст-картинка, размер батча 4, 250 тыс шагов, 416 A100.
Смешанное разрешение: 768 ≤ ширина × высота ≤ 1024, 280 млн пар текст-картинка, размер батча 1, 350 тыс шагов, 416 A100.
Примеры генераций
Результаты сравнения и примеры генераций
Для сравнения моделей мы собрали свой сбалансированный набор из 2100 промптов по 21 категории и сравнили различные веса Kandinsky 3.0, чтобы выбрать лучший. Для этого провели три этапа side-by-side, в которых принимали участие 28 разметчиков.
Затем, когда наилучшая версия модели Kandinsky 3.0 была выбрала, мы провели side-by-side сравнение с моделью Kandinsky 2.2. В исследовании участвовало 12 человек, которые проголосовали суммарно 24 800 раз. Для этого разработали бот, который демонстрировал одну из 2100 пар изображений. Каждый человек выбирал лучшее изображения по двум критериям:
соответствие тексту,
визуальное качество картинки.
Сравнение происходило как по качеству визуализации и пониманию текста суммарно по всем категориям, так и по каждой категории в отдельности:
Примеры изображений с side-by-side сравнения:
Одно из основных преимуществ модели Kandinsky 3.0 состоит в улучшенном понимании текстового запроса (разница с Kandinsky 2.2 в 31%). Кроме того, кратно улучшилось качество понимания запросов на темы отечественного культурного кода, например:
Ниже представлены примеры генераций популярных моделей в сравнении с Kandinsky 3.0:
A beautiful girl
A highly detailed digital painting of a portal in a mystic forest with many beautiful trees. A person is standing in front of the portal.
A man with a beard
A 4K dslr photo of a hedgehog sitting in a small boat in the middle of a pond. It is wearing a Hawaiian shirt and a straw hat. It is reading a book. There are a few leaves in the background.
Barbie and Ken are shopping
Extravagant mouthwatering burger, loaded with all the fixings. Highlight layers and texture
A bear in Russian national hat with a balalaika
Inpainting + Outpainting
Наша команда провела отдельную работу по разработке inpainting/outpainting-модели, с помощью которой можно было бы редактировать изображения: менять в них необходимые объекты и целые области внутри изображения (inpainting) или расширять их новыми деталями до больших панорам (outpainting). Задача inpainting намного сложнее, чем стандартная генерация, потому что модели нужно научиться генерировать полноценное изображение не только с учётом текста, но и с учётом известной части этого изображения.
Для обучения inpainting-части модели за основу взяли подход из GLIDE, который уже был ранее реализован в моделях семейства Kandinsky, а также в семействе моделей Stable Diffusion: входной слой U-Net модифицируется так, чтобы на вход дополнительно можно было подать image latent и mask. Таким образом, U-Net принимает на вход целых 9 каналов: 4 для исходного latent, 4 для latent-изображения и дополнительный канал для маски. С точки зрения модификаций, в целом, это всё — далее обучение ничем не отличается от обучения стандартной диффузионной модели.
Важной особенностью задачи является то, как генерируются маски и какой текст используется для обучения. Пользователи могут нарисовать либо маску кисточкой, либо новое изображение через outpаinting. Чтобы учесть то, как работает пользователь, при обучении мы создали специальные маски, которые имитируют его поведение: маски произвольной формы, нарисованные кисточкой, маски объектов, а также заполнение изображения по аналогии с Generative Fill от Adobe. С коллегами, которые разрабатывали последнюю модель, мы даже консультировались и переняли их лучшие практики :)
Примеры масок для обучения модели Inpainting
Как результат, модель хорошо справляется с заменой изображения, а также с его дополнением (см. примеры ниже).
Примеры Inpainting
Примеры Outpainting
Animation
Кроме этого, мы также адаптировали модель Kandinsky 3.0 для использования в пайплайне Deforum — таким образом, с помощью Kandinsky 3.0 теперь можно генерировать не только изображения, но и анимации. Подробнее с принципом работы этого режима можно познакомиться в одной из наших недавних статей.
Разные принципы зашумления изображения в процессе диффузии
Основной сложностью интеграции новой модели стало различие в режимах добавления шума в процессе диффузии: у Kandinsky 2.2 шум добавляется по линейному закону (картинка сверху), в то время как у Kandinsky 3.0 — по косинусному (картинка снизу). Эта особенность потребовала некоторого количества экспериментов для успешной адаптации.
Примеры анимаций
Отметим, что разрешение изображений, а также количество кадров в секунду снижено в связи с техническими требованиями Хабра.
Красивая женщина, темные волосы, веснушки, цветочная корона из крупных пионов и роз, красивый градиентный розовый фон, верхнее освещение, профессиональная фотография, студийная фотография, 4k, режим: «live»
Предельная детализация, 8k, ультра качество, шедевр, глубина резкости, плавное освещение, иллюстрация, очень милый реалистичный чебурашка в курточке с мандарином в руках, сложные острые детали, b объемное освещение, невероятно подробные масштабы, невероятно подробные глаза, большие уши, невероятно подробный вид крупным планом, радужный свет, детализированная четкая шерсть, снег, зимний город, фейерверк, огни, бенгальские огни, гирлянды, радость, смех, улыбка, доброта, счастье, режим: «zoom out»
Красивый зимний лес, artstation, режим: «shift right»
Вывод и планы
Мы представили нашу новую архитектуру генерации изображений по тексту — Kandinsky 3.0. Значительно улучшили понимание текста и российской культуры в сравнении с нашими предыдущими моделями и обязательно будем продолжать двигаться в этом направлении. Что касается научной стороны, в наши планы входит создание ещё одной новой модели генерации, которая должна будет сказать своё новое слово на международной AI-арене.
Область искусственного интеллекта и генеративного обучения открывает широкие просторы для дальнейшего развития. При этом мы уверены, что это развитие в любом случае будет очень интересным и потребует смены взглядов на многие окружающие нас вещи. Нам, всему человечеству, ещё только предстоит познать всю силу генеративного обучения. Оставайтесь с нами на нашей волне, чтобы не пропустить то, как будет меняться этот мир, в том числе нашими усилиями!
Авторы и их вклад
Модель Kandinsky 3.0 разработана командой Sber AI при партнёрской поддержке учёных из Института искусственного интеллекта AIRI на объединённом наборе данных Sber AI и компании SberDevices.
Коллектив авторов: Владимир Архипкин, Вячеслав Васильев, Андрей Филатов, Анастасия Мальцева, Саид Азизов, Арсений Шахматов, Игорь Павлов, Михаил Шойтов, Юлия Агафонова, Сергей Нестерук, Анастасия Лысенко, Илья Рябов, Ангелина Куц, София Кирилова, Сергей Марков, Андрей Кузнецов и Денис Димитров.
Контакты для коммуникации
По всем возникающим вопросам и предложениям по развитию модели, добавлению новых возможностей и сотрудничеству в части использования наших разработок можно писать Денису или Андрею.
Полезные ссылки
fusionbrain.ai (доступна генерация по тексту и режимы inpainting/outpainting, а также генерация анимации и видео, об этой новинке подробнее будет в этой статье)
Telegram-bot (доступна генерация по тексту)
rudalle.ru (доступна генерация по тексту)
Telegram-bot для видео генерации (доступна генерация анимация и видео, об этой новинке речь будет в этой статье)
HuggingFace
Github
Deforum (project page по генерации анимаций)
Каналы авторов:
