KandiSuperRes Flash — обновлённая модель повышения разрешения изображений
В апреле этого года мы выпустили модель Kandinsky 3.1, которая поддерживает много различных режимов, в том числе и возможность генерировать изображения в 4К с помощью диффузионной модели KandiSuperRes. Подробнее про архитектуру и результаты можно почитать в этой статье. Модель позволяет генерировать более чёткие изображения в высоком разрешение, однако не устраняет артефакты, получившиеся на этапе генерации с помощью модели Kandinsky 3.1. Чтобы устранить эти недочёты, мы разработали модель KandiSuperRes Flash, которая улучшает изображение, делает его более эстетичным, и при этом повышает разрешение в два раза.
Пример повышения разрешения с помощью KandiSuperRes Flash
Описание конвейера
Обновлённый конвейер KandiSuperRes Flash содержит: дистиллированную версию KandiSuperRes и дистиллированную модель Kandinsky 3.0 Flash. Подробно про дистилляцию диффузионных моделей мы уже писали в своей предыдущей статье. Там мы рассказали, как обучали Kandinsky 3.0 Flash. По аналогии мы обучили и дистиллированную версию модели повышения разрешения. В качестве дискриминатора использовали downsample-часть U-Net от модели KandiSuperRes, а в качестве функции потерь — Wasserstein Loss. Заметили, что такой подход не только помогает ускорить генерирование изображений, но и позволяет создавать более чёткие и детализированные изображения благодаря дискриминатору и обучению в режиме GAN. Ниже приведено сравнение модели KandiSuperRes и её дистиллированной версии.
Сравнение модели KandiSuperRes и её дистиллированной версии
В ходе экспериментов мы обнаружили, что дистиллированную модель Kandinsky 3.0 Flash можно использовать в качестве refiner, генерируя за 1 или 2 шага, в зависимости от того, насколько сильно мы хотим внести изменения в изображение.
Результат работы Kandinsky 3.0 Flash в качестве рефайнера
Мы объединили эти две модели в общий конвейер под названием KandiSuperRes Flash. На первом этапе вдвое повышается разрешение за 4 шага пиксельной диффузии, а на втором этапе выполняется рефайн изображения за 1 шаг латентной диффузии (зашумление до 229 шага и обратный процесс). Время инференса всего конвейера занимает 6 секунд на H100 при повышение разрешения из 1К в 2К.
Результаты и примеры
Примеры генераций с помощью KandiSuperRes Flash
Чтобы понимать, насколько сильно модель KandiSuperRes Flash улучшает изображения, мы провели side-by-side (SBS) сравнение с моделью Kandinsky 3.1. SBS проводился на фиксированной корзине запросов из 2100 промптов (100 промптов по каждой из 21 категорий). Каждая генерация оценивалась по визуальному качеству (какое из двух изображений вам больше нравится). Подробнее про методологию SBS можно прочитать в статье Kandinsky 3.0. Результаты сравнения показаны на графиках ниже. Разметчики вдвое чаще выбирали изображения после повышения разрешения с помощью KandiSuperRes Flash, чем до повышения (28% против 15%). В 19% случаев разметчикам нравились оба изображения, а в 37% случаев ни одно из изображений не было выбрано. Однако это можно объяснить тем, что на этих изображениях присутствовали сильные артефакты, которые не исправлялись моделью повышения разрешения.
Сравнение моделей по визуальному качеству изображений
Сравнение визуального качества изображений по темам
Бесконечный super resolution
Поскольку KandiSuperRes Flash теперь не только повышает разрешение и чёткость изображения, но и детализирует его, дорисовывая некоторые детали, появилась возможность бесконечно увеличивать изображение (до x16 и даже сильнее).
Примеры бесконечного повышения разрешения с помощью KandiSuperRes Flash
Примеры бесконечного повышения разрешения с помощью KandiSuperRes Flash
Заключение
Мы представили новую версию модели повышения разрешения KandiSuperRes Flash, которая стала значительно лучше генерировать изображения. KandiSuperRes Flash теперь не только повышает чёткость, но и исправляет артефакты, дорисовывает детали, повышает эстетичность изображения. И одно из самых главных преимуществ — это возможность использования модели в режиме «бесконечного super resolution». Код и веса можно найти на Github и HuggingFace.
Коллектив авторов: Анастасия Мальцева, Владимир Архипкин, Николай Герасименко, Андрей Кузнецов и руководитель научной группы Sber AI Research Денис Димитров.
