Прошёл ровно год с момента релиза модели Kandinsky 2.1 — именно эта модель принесла известность нашей исследовательской группе Sber AI Research и дала толчок развитию всей линейки моделей Kandinsky. В честь этой даты мы выпускаем новую версию модели Kandinsky 3.1, о которой я расскажу подробнее в этой статье.

Введение

В 2023 году мы все наблюдали очень большой прогресс в области исследования языковых, мультимодальных и генеративных моделей. Этот прогресс коснулся и привычной нам модальности изображений (Kandinsky 2.1, 2.2, 3.0, Stable Diffusion XL, IF, Шедеврум, MJv6 и др.), и модальности текстов (ChatGPT, GPT-4, LLaMA, Falcon, GigaChat и др.), и аудио (VALL-E, MusicLM и др.), и 3D (Magic3D и др.), и даже модальности видео (Kandinsky Video, Gen-2, CogVideo и др.). В 2024 всё движется вперёд ещё более впечатляющими темпами: выходят на абсолютно новый уровень модели генерации изображений (SD3, SD3-Turbo), видео (Sora), музыки и целых песен (Suno) и т. д. При этом все основные игроки стараются равномерно развиваться и повышать качество синтеза. Текстовые чат‑боты научились взаимодействовать с внешними системами посредством плагинов, синтез изображений вышел на уровень фотореалистичных генераций, длина генерируемых видео постепенно увеличивается с сохранением сюжетной связности между кадрами (вполть до 1 минуты у Sora). И такой прогресс обусловлен уже не только наращиванием вычислительных мощностей, но и большим числом неординарных архитектурных решений, которые позволяют добиваться лучшего качества, а также сложными технологиями инженерии данных, обеспечивающими создание огромных и в то же время очень качественных датасетов для обучения моделей.

В ноябре прошлого года на конференции AI Journey наша команда представила text-to-image модель нового поколения Kandinsky 3.0. В целом, пересмотрев архитектуру и взяв более мощный текстовый энкодер (по сравнению с семейством моделей Kandinsky 2.X), нам удалось добиться значительного роста в качестве изображений с точки зрения реалистичности и детализации, улучшить понимание текста и побить качество модели SDXL на side-by-side-сравнении с точки зрения человеческих предпочтений. Это является наиболее показательной и справедливой метрикой качества в задачах генерации любой модальности (хотя у этой методологии тоже есть свои минусы). Подробнее о модели Kandinsky 3.0 можно прочитать в этой статье. Одновременно с этим на базе этой модели мы выпустили первую российскую модель генерации видео по тексту Kandinsky Video, о которой можно больше узнать здесь.

Сегодня мы представляем модель Kandinsky 3.1 — идейное продолжение модели Kandinsky 3.0, которую мы улучшили и обогатили набором различных полезных функций и режимов, предоставляющих пользователям больше возможностей полноценно использовать всю силу нашей новой модели. В Kandinsky 3.1 входит сразу несколько улучшений, о каждом из которых я расскажу подробнее в статье:

• дистилляция по числу шагов диффузии (Kandinsky 3.0 Flash);

• улучшение промпта с помощью языковой модели Neural-Chat-v3-1;

• IP-Adapter, который позволяет дополнительно (помимо текста) обуславливаться на изображение (за счёт этого удалось вернуть режимы смешивания изображений, изображения и текста, которые были в версиях Kandinsky 2.X, но работали там за счёт наличия в этих моделях специального блока image prior);

• ControlNet — механика, реализующая способы дополнительного обуславливания (контроля за генерацией) на основе canny edges, depth maps и т. д.;

• Inpainting — модификация основной модели для дорисования по тексту недостающих частей изображения;

• SuperRes — специальная диффузионная модель, повышающая разрешение изображения (так, в Kandinsky 3.1 теперь можно генерировать 4K изображения).

Кроме этого, мы обучили маленькую версию модели Kandinsky 3.0 Small (1B), с которой более просто и удобно экспериментировать.

Ускорение модели

Серьёзной проблемой модели Kandinsky 3.0, как и всех диффузионных моделей, была скорость генерации. Для получения одного изображения необходимо было пройти 50 шагов в обратном процессе диффузии, то есть 50 раз пропустить данные через U-Net с размером батча 2 для classifier free guidance. Для решения этой проблемы мы использовали подход Adversarial Diffusion Distillation, впервые описанный в статье от Stability AI, но с рядом существенных модификаций:

1. В случае использования предобученных пиксельных моделей в качестве дискриминатора возникла бы необходимость декодировать сгенерированную картинку при помощи MoVQ Decoder и пробрасывать через него градиенты, что привело бы к огромным затратам памяти. Данные затраты не позволили бы обучать модель в разрешении 1024 × 1024. Поэтому в качестве дискриминатора мы использовали замороженную downsample-часть U-Net от Kandinsky 3.0 с обучаемыми головами после каждого слоя понижения разрешения. Это связано с желанием сохранить возможность генерировать картинки в высоком разрешении.

2. Мы добавили Cross Attention на текстовые эмбеддинги от FLAN-UL2 в головы дискриминатора вместо добавления текстового CLIP эмбеддинга. Это позволило улучшить понимание текста дистиллированной моделью.

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

4. Мы убрали регуляризацию в виде Distillation Loss, так как по нашим экспериментам она не оказывала существенного влияния на качество модели.

Мы обнаружили, что довольно быстро генератор становится сильнее дискриминатора, что приводит к нестабильности обучения. Чтобы решить эту проблему, мы значительно увеличили learning rate у дискриминатора. У дискриминатора он равнялся 1e-3, в то время как у генератора — 1e-5. Для предотвращения расходимости мы также использовали gradient penalty, как и в оригинальной работе.

Обучение происходило на «эстетичном» (отобранном руками) датасете размером 100K пар «текст-изображение», который является подсетом датасета для претрейна Kandinsky 3.0.

В результате этого подхода получилось ускорить Kandinsky 3.0 почти в 20 раз, сделав возможной генерацию изображения всего за 4 прохода через U-Net. Также на скорость повлиял тот факт, что теперь нет необходимости использовать classifier free guidance. Kandinsky 3.0 из диффузионной модели по факту превратился в GAN (Kandinsky 3.0 Flash), обученный с хорошей начальной инициализацией весов после претрейна.

Однако для серьёзного ускорения пришлось пожертвовать качеством понимания текста, что показывают результаты side-by-side (SBS) сравнения. SBS проводится на фиксированной корзине запросов из 2100 промптов (100 промптов по каждой из 21 категорий). Каждая генерация оценивается по визуальному качеству (какое из двух изображений вам больше нравится) и по соответствию тексту (какое из двух изображений лучше соответствует запросу). Про методологию side-by-side (SBS) сравнения можно подробнее прочитать в статье Kandinsky 3.0.

Приведём примеры изображений, сгенерированных моделью Kandinsky 3.0 Flash.

Бьютификация (улучшение) запроса

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

Для решения этой проблемы в Kandinsky 3.1 была встроена опция бьютификации запроса — способ улучшения запроса (добавление в него деталей) пользователя с помощью большой языковой модели (LLM). Бьютификация работает очень просто: на вход языковой модели подаётся инструкция с просьбой улучшить запрос, а далее ответ модели подается на вход Kandinsky для генерации.

В качестве LLM мы использовали neural-chat-7b-v3-1 от Intel (дообученная модель Mistral-7B) со следующим системным промптом:

### System:\nYou are a prompt engineer. Your mission is to expand prompts written by user. You should provide the best prompt for text to image generation in English. \n### User:\n{prompt}\n### Assistant:\n

Тут {prompt} — это запрос, который написал пользователь.

Примеры генераций на один и тот же запрос, но с бьютификацией и без неё представлены ниже (с помощью модели Kandinsky 3.0).

Также мы провели side-by-side сравнение качества генераций, сделанных Kandinsky с использованием функции бьютификации запроса и без него. Мы провели тестирование и для Kandinsky 3.0, и для Kandinsky 3.1, чтобы оценить, как сильно влияет языковая модель на генерируемые изображения.

Генерации и сравнение с другими моделями

Сравнение Kandinsky 3.1 (Flash + улучшение промпта) с предыдущими версиями:

Сравнение Kandinsky 3.1 (Flash + улучшение промпта) с другими моделями:

Также интересно проследить эволюцию моделей Kandinsky, начиная с Kandinsky 2.1:

Редактирование изображений и Text-Image Guidance

В отличие от модели Kandinsky 3.0, в версию Kandinsky 3.1 мы внедрили функционал генерации изображения не только с помощью текстового запроса, но и/или с помощью визуальной подсказки в виде подаваемого на вход изображения. Это позволяет редактировать уже имеющееся изображение, изменять его стиль и добавлять к нему новые объекты. Для этого мы использовали IP-Adapter — подход, продемонстрировавший хорошие результаты в сравнении с традиционным дообучением.

Для реализации IP-Adapter-а на основе нашей базовой модели генерации и её имплементации в библиотеке diffusers мы использовали адаптеры внимания. В качестве энкодера изображений был взят ViT-L/14 (Visual Transformer, дообученный в пайплайне CLIP). C его помощью мы получаем эмбеддинги размера batch_size x 768, которые потом, посредством линейного слоя, преобразуются в тензоры размером batch_size × 4 × 4096. Добавив пару новых слоев для key и value изображений в механизме cross attention, мы складываем выход обычного текстового cross attention с выходом cross attention для изображений. Обучение производилось на датасете COYO-700M с размером батча 288 в течение 800K итераций.

Поддерживаемые варианты инференса:

1. Вариация изображения. Чтобы сделать вариацию изображения, мы просто считаем эмбеддинги изображения с помощью CLIP-ViT-L/14 и подаём их в модель.

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

3. Смешивание изображения и текста. Мы считаем эмбеддинги изображения и подаём их в модель вместе с текстом, поскольку мы сохранили стандартный cross attention на текст.

Редактирование в режиме ControlNet

Мы обнаружили, что подход на основе IP-адаптера не сохраняет форму объектов на изображении, поэтому мы решили обучить ControlNet в дополнение к нашей модели генерации для консистентного изменения внешнего вида изображения, сохраняющего больше информации в сравнении с исходным. В качестве модели для получения границ на изображении, подающихся на вход ControlNet, мы использовали HED detector. Обучение длилось 5000 итераций на датасете COYO 700m на 8 GPU Tesla A100 с размером батча 512.

Kandinsky 3.1 Inpainting

В Kandinsky 3.1 Inpainting мы сфокусировались на улучшении качества модели генерации объектов. В Kandinsky 3.0 мы учили модель восстанавливать изображение по его исходному описанию — из-за этого модель очень хорошо восстанавливает исходное изображение, но когда дело доходит до создания другого объекта на месте старого, то модель может не справиться. Один из способов исправления этой проблемы — дообучение на масках из датасетов для задачи object detection или segmentation (например, Paint by Example или SmartBrush). При использовании bounding box масок модель учится генерировать изображения чётко по текстовому запросу, а не по описанию целой картинки. Таким образом, качество модели при использовании её людьми растёт —именно в таком формате модель Inpainting и используется на инференсе. Чтобы модель не разучилась делать и «классический» inpainting (по полным описаниям изображения), мы сбалансировали наши обучающие сеты — 50% масок приходят с bounding boxes, а оставшиеся 50% выбираются случайным образом, как было раньше в Kandinsky 3.0 Inpainting.

Поскольку в качестве текстовых запросов мы используем только имена классов, то модель может разучиться генерировать изображения по длинным запросам. Поэтому мы решили доразметить наш датасет с помощью LLaVA 1.5. Для этого, после выбора bounding box, который мы использовали в качестве маски, мы подавали crop-изображения в LLaVA, чтобы получить текстовое описание этого кусочка изображения. Далее это текстовое описание использовалось как текстовый запрос.

Также мы провели сравнение с другими моделями, чтобы «оцифровать» качество нашего нового метода inpainting-а. Для этого мы взяли датасет COCO, случайным образом выбрали из него 1000 изображений и по одному объекту с каждого изображения, который мы перегенерировали. Далее на полученных картинках мы прогнали модель, обученную на датасете детекции YOLO-X, и посчитали её метрики качества детекции. Если детектор, обученный на реальных изображениях, сможет задетектировать сгенерированный объект, то можно сделать вывод, что объект сгенерирован достаточно естественно. Ниже приведены метрики:

Из таблицы следует, что наша модель отлично справляется с дорисовыванием больших объектов (да и в целом её качество при добавлении любых объектов очень хорошее).

KandiSuperRes

В новой версии Kandinsky 3.1 появилась возможность получать генерации изображений в разрешении 4K. Для этого была обучена диффузионная модель повышения разрешения KandiSuperRes.

За основу была взята архитектура Kandinsky 3.0, но с некоторыми модификациями. 

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

• Вместо обычного U-Net был реализован EfficientUnet подобно подходу, описанному в статье Imagen. В сравнении с U-Net Kandinsky 3.0, EfficientUnet потребляет меньше памяти и также имеет лучшую сходимость. Отличие в том, что вместо 3 Residual блоков на каждом понижении разрешения в EfficientUnet используется большее количество блоков на низких разрешениях и меньшее количество блоков на высоких. Также изменяется порядок выполнения свёртки и downsampling/upsampling-операций относительно исходного U-Net. Кроме этого, мы убрали обусловливание на текстовый промпт, так как он не вносит вклад в генерации в высоком разрешении 4К. В итоге EfficientUnet KandiSuperRes содержит 413M параметров.

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

Обучение осуществлялось в 2 этапа: сначала модель училась на датасете LAION на 32 A100 в течение 1,57М шагов с размером батча 2 на разрешение 256 → 1024. Затем модель училась на «эстетичных» сетах высокого разрешения, используемых при обучении Kandinsky 3.0, в течение 1,5М шагов. На втором этапе обучения было добавлено JPEG-сжатие подобной той схеме, которая описана в статье Real-ESRGAN.

Модель KandiSuperRes позволяет работать с изображениями различного разрешения, однако основной целью являются генерации в высоком разрешении 4K. Так как модель KandiSuperRes обучалась задаче увеличивать разрешение с 256 до 1024, а на больших разрешениях обучать не было возможности из-за переполнения памяти A100, для генерации в 4К использовался алгоритм MultiDiffusion, позволяющий создавать панорамы. Суть алгоритма заключается в том, что изначально изображение делится на перекрывающиеся патчи, и затем на каждом шаге диффузии удаляется шум, а значения пикселей/латентов перекрывающихся областей усредняются. И таким образом, пройдя все шаги диффузии, мы получаем бесшовное изображение любого разрешения. В итоге на инференсе модель KandiSuperRes работает за 5 шагов, используя DPMSolverMultistepScheduler. Время инференса для генерации изображения в 4К занимает 13 секунд, а в 1К — 0,5 секунды.

В таблице ниже приведено сравнение KandiSuperRes с моделями Real-ESRGAN и Stable Diffusion x4 Upscaler по метрикам FID, SSIM, PSNR и L1 на датасетах Wikidata 5К, RealSR(V3) и Set14. Wikidata 5К содержит 5000 изображений, собранных из Википедии, в разрешении 1К. RealSR(V3) содержит 100 тестовых изображений в разрешениях 1K и 2K. Set14 содержит 14 изображений в низких разрешениях с JPEG-артефактами. В итоге модель KandiSuperRes показала наилучшие результаты. 

Сравнение моделей Real-ESRGAN, Stable Diffusion и KandiSuperRes:

На рисунке 1 приведены примеры генераций моделей KandiSuperRes, Stable Diffusion и Real-ESRGAN в разрешение 1024. На рисунке 2 приведены примеры генераций KandiSuperRes в разрешение 4К.

Код и веса KandiSuperRes выложены на Github и HuggingFace.

Маленькие версии модели (1B)

Чтобы сделать нашу модель более доступной для запуска и дообучения при малых вычислительных ресурсах, мы решили обучить её маленькую версию (Kandinsky 3.0 Small). Для неё мы взяли U-Net, содержащий 1B параметров, и текстовый энкодер от FLAN-T5 XL. Обучение происходило в течение 1М итераций на датасете COYO-700M с размером батча 2048. Внизу представлены примеры генераций с помощью этой модели.

Вывод

В этой статье мы представили новую версию нашей диффузионной модели генерации изображений по тексту Kandinsky 3.1, которую мы сделали более эффективной с помощью современных методов дистилляции и дополнили новыми режимами работы, которых не было в версии Kandinsky 3.0, — возможностью редактирования изображений и переносом стиля. Дополнительно мы улучшили качество работы inpainting-а — метода повышения разрешения для перехода в пространство пикселей из пространства скрытых представлений. Кроме того, помимо доступа к основной большой модели, мы предоставили пользователям возможность работы с версиями нашей модели с меньшим числом параметров, что гораздо удобнее для локального запуска и дообучения при малых ресурсах для своих целей и задач. Можно сказать, что наше основное достижение — это то, что мы сделали творческую генерацию гораздо более удобной для пользователя, который теперь может полноценно использовать весь потенциал нашей модели.

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

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

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

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

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

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

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

• DendiMath

• CompleteАI

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

• AwesomeDL

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

• Igrek.log

• Коктейль Кокто

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

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

Пока что Kandinsky 3.1 доступен для ограниченного круга пользователей. При этом широкий доступ к модели и к её полному функционалу будет предоставлен в самое ближайшее время.

Полезные ссылки для пользования предыдущими версиями модели:

• fusionbrain.ai (доступна генерация по тексту и inpainting)

• Telegram-bot (доступна генерация по тексту, а также есть доступ к моделям линейки Kandinsky 2.X и их полному функционалу)

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

• Project Page (description of Kandinsky 3.1 in English)

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

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