В апреле этого года мы выпустили модель Kandinsky 3.1, которая поддерживает много различных режимов, в том числе и возможность генерировать изображения в 4К с помощью диффузионной модели повышения разрешения KandiSuperRes. Подробнее про архитектуру и результаты можно почитать в этой статье. Модель позволяет генерировать более чёткие изображения в высоком разрешении, однако не устраняет артефакты, получившиеся на этапе генерации с помощью модели Kandinsky 3.1. Чтобы устранить эти недочёты, мы разработали модель KandiSuperRes Flash, которая улучшает изображение, делает его более эстетичным, и при этом повышает разрешение в два раза.

Описание пайплайна

Основная цель разработки обновленной версии модели KandiSuperRes Flash — это улучшение качества синтезируемых изображений по двум аспектам: повышение разрешения изображений и исправление артефактов, полученных на этапе генераций с помощью моделей семейства Kandinsky. Чтобы решить одновременно обе задачи, был разработан пайплайн, в который встроено две модели: super resolution и refiner.

Super resolution

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

• замороженная downsample‑часть U‑Net от модели KandiSuperRes с обучаемыми головами после каждого слоя понижения разрешения в качестве дискриминатора;

• Wasserstein Loss в качестве функции потерь.

В результате дистиллированная модель позволяет создавать более чёткие и детализированные изображения благодаря дискриминатору и обучению в режиме GAN. Ниже приведено сравнение модели KandiSuperRes и её дистиллированной версии.

Refiner

В качестве рефайнера мы используем дистиллированную модель Kandinsky 3.0 Flash, потому что в ходе экспериментов заметили, что данная модель довольно‑таки хорошо справляется с этой задачей. Kandinsky 3.0 Flash позволяет генерировать изображения из шума за 4 шага диффузии. А если использовать 1 или 2 шага диффузии и зашумлять соответственно до 229 или 479 шага, то модель как раз‑таки вносит изменения на изображение, тем самым исправляя артефакты. Чем больше шагов, тем сильнее изменяется изображение. В текущем пайплайне рефайнер работает за 1 шаг диффузии, чтобы исправлялись артефакты, но при этом сохранялись структура и содержание изображения. Ниже показаны примеры работы Kandinsky 3.0 Flash в качестве рефайнера за 1 шаг диффузии.

Таким образом, мы объединили дистиллированные модели KandiSuperRes и Kandinsky 3.0 Flash в общий пайплайн под названием KandiSuperRes Flash. На первом этапе вдвое повышается разрешение за 4 шага пиксельной диффузии, а на втором этапе выполняется рефайн изображения за 1 шаг латентной диффузии. Время инференса всего пайплайна занимает 6 секунд на H100 при повышение разрешения из 1К в 2К.

Результаты и примеры

Оценка результатов модели KandiSuperRes Flash проводилась двумя способами — с помощью автоматических метрик и side‑by‑side (SBS) сравнения. Для дистиллированной модели KandiSuperRes, выполняющей super resolution, мы посчитали метрики FID, PSNR, SSIM, L1 на датасете Set14 и сравнили с моделями Real‑ESRGAN и Stable Diffusion x2 Upscaler. В итоге наша модель превосходит своих конкурентов по качеству. Результаты приведены в таблице ниже.

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

Для модели KandiSuperRes Flash, выполняющей super resolution + refine, на том же наборе Set14 мы посчитали автоматическую метрику aesthetic score, которая показывает, насколько людям в среднем нравится изображение. Эта метрика рассчитывается с помощью MLP, который на вход принимает эмбеддинги CLIP. В таблице ниже мы сравнили модели Real‑ESRGAN, Stable Diffusion x2 Upscaler, KandiSuperRes distillation и KandiSuperRes Flash. В итоге дистиллированная версия KandiSuperRes превосходит Real‑ESRGAN и Stable Diffusion, а модель KandiSuperRes Flash, выполняющая ещё и рефайн изображения, показывает наилучшие результаты.

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

Чтобы понимать, насколько сильно модель KandiSuperRes Flash (super resolution + refiner) улучшает изображения, сгенерированные Kandinsky 3.1, мы провели side‑by‑side (SBS) сравнение. SBS проводился на фиксированной корзине запросов из 2100 промптов (100 промптов по каждой из 21 категорий). В исследовании участвовало 12 человек. Каждый разметчик выбирал лучшее изображение из двух по визуальному качеству (KandiSuperRes Flash vs Kandinsky 3.1). Сравнение происходило как суммарно по всем категориям, так и по каждой категории в отдельности. Подробнее про методологию SBS можно прочитать в статье Kandinsky 3.0.

Результаты SBS показаны на графиках ниже. Разметчики вдвое чаще выбирали изображения после повышения разрешения с помощью KandiSuperRes Flash, чем до повышения (28 % против 15 %). В 19 % случаев качество генераций обеих моделей удовлетворяло разметчиков, а в 37 % случаев ни одно из изображений их не устроило. Однако это можно объяснить тем, что на картинках, сгенерированных Kandinsky 3.1., были сильные артефакты (такие как неправильные формы рук у людей, нарушение строения тела у животных, несоблюдение физических законов, отсутствие реалистичности и пр.), которые не исправлялись моделью повышения разрешения.

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

Бесконечный super resolution

Поскольку KandiSuperRes Flash теперь не только повышает разрешение и чёткость изображения, но и дорисовывает детали, появилась возможность бесконечно увеличивать изображение (до x16 и даже сильнее).

Заключение

Мы представили новую версию модели повышения разрешения KandiSuperRes Flash, которая стала значительно лучше генерировать изображения. KandiSuperRes Flash теперь не только повышает чёткость, но и исправляет артефакты, дорисовывает детали, повышает эстетичность изображения. И одно из самых главных преимуществ — это возможность использования модели в режиме «бесконечного super resolution». Код и веса можно найти на Github и HuggingFace.

Коллектив авторов: Анастасия Мальцева, Владимир Архипкин, Николай Герасименко, Андрей Кузнецов и руководитель научной группы Sber AI Research Денис Димитров.

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

• DendiMath

• CompleteАI