Эксперименты со Schedulers в Stable Diffusion

В этой статье разберём, что есть scheduler в диффузионных моделях и как можно подменять их, пользуясь возможностями библиотеки diffusers.

Stable Diffusion

Для начала пару слов и картинок о рассматриваемой модели. Stable Diffusion успел нашуметь в 2022 году, и, кажется, многие новости будут связаны с этим в будущем. Модель эта примечательна в первую очередь тем, что способна по текстовому запросу, наподобие запроса к поисковику, выдавать картинки в высоком качестве, с проработанными деталями и, во многих случаях, отражающими суть текстового запроса. И, что немало важно, часто превосходя качество генерации с помощью GAN.

Немного о подходах. Подход к генерации изображения из текста называется text2image. Кроме него Stable Diffusion (и аналогичные ей проприетарные модели: dalle2, imagen и midjourney) способны генерировать:

• image2image, когда входная картинка и запрос используются как отправная точка для генерации нового изображения,

• image inpainting, когда с картинкой вдобавок на вход подается маска, где на изображении требуется генерация, а где нет,

• upscaling, этот подход похож на image2image, тоже на вход картинка, но на выходе не модификация ее под текстовый запрос, а изображение в увеличенном разрешении.

Тут стоит отметить, что это не всё, что уже можно делать с помощью Stable Diffusion. В качестве отправной точки можно взглянуть на другие пайплайны по ссылке https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion.

Генерация с помощью Stable Diffusion

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

На этапе инференса, т.е. когда модель, уже обученная на огромном количестве пар текстовых описаний с картинками, генерирует изображения по тексту, мы имеем следующую схему работы. В ней два входа: эмбеддинг входного текста, т.е. обработанный с помощью текстовой модели CLIP текст, разбитый на токены, и latent seed, массив размера 64 на 64, состоящий либо из случайных чисел из гауссова шума, либо заполненный с помощью выхода энкодера VAE. "Смесь" этих двух входов подаётся в сеть архитектуры UNet, чьей задачей является предсказать, что в данном массиве 64x64 (тензоре) является шумом, чтобы затем, с помощью шага scheduler очистить текущий вход от предсказанного шума и повторить процедуру. В зависимости от выбора алгоритма планировщика, варьируется и количество повторений. После этого цикла результат отправляется в декодер VAE, который выдаёт окончательное изображение. Здесь примечательно, что очистка изображения от шума происходит сугубо в латентном пространстве, т.е. в пространстве тензоров 64x64, а декодирование в картинку происходит лишь однажды, но, как мы увидим позже, если декодировать latent на промежуточных этапах в картинку, то будет явно заметна очистка от шума и в изображении.

О роли Scheduler

Подробно о "планировщиках" в генеративных моделях описано в статье от NVLabs https://arxiv.org/pdf/2206.00364.pdf. Попробую, донести ту суть, которую я сам уловил из статьи и опыта работы со stable diffusion.

1. Scheduler используется и на этапе обучения, и на этапе инференса для восстановления зашумленного изображения.

2. Глобально существует два вида планировщиков: с детерминистическим сэмплированием и со стохастическим.

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

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

3. Scheduler в паре с UNet* образуют цикл, в котором решается ДУ. Где UNet – это нейронная сеть, которая учится минимизировать ошибку денойзинга, которая соответственно выполняет процедуру очистки от шума. Участвует в численном решении дифференциального уравнения [см. уравнения 2, 3 и алгоритм 1 из ссылки выше], путём того, что результат сети перевычисляется на каждом шаге схемы.
   * – помимо UNet могут быть сети и другой архитектуры.

4. Некоторые исследования сфокусированы на поиске оптимального планировщика, который имеет минимальное число NFE (neural function evaluation), т.е. количество шагов схемы решения ДУ, за которое можно добиться качественного результата.

Так, кратко в пунктах, могу описать роль планировщика в диффузионных моделях, если возникли вопросы в результате прочтения – давайте обсудим в комментариях!

Эксперимент

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

Суть эксперимента заключается в том, чтобы увидеть качественную разницу между генерациями различными подходами по одному заданному текстовому запросу, в нашем случае первый запрос, это гибрид Маска и Карателя, второй – кот из примера вверху. При этом фиксируя random seed. Увидеть, на каком этапе очистки от шума, при различных подходах, проявляется вменяемое изображение, и как результаты разнятся между собой.

Код исполнялся на Tesla T4 в облаке и генерация четырех изображений по запросу во всех случаях занимала 01:42 min. Номера шагов указаны красным в верхних строках.

EulerDiscreteScheduler

В случае Маск-Каратель, на 50-м шаге начинает проявляться изображение, которое мы четко видим на шаге 100. Этот scheduler создает довольно кинематографичные результаты, особенно второй и третий сверху, но на двух из четырех толком не видно лица.

EulerAncestralDiscreteScheduler

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

DDIMScheduler

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

UPDATE:
Если использовать дефолтный DDIMScheduler, инициализированный с помощью from_config, то мы увидим результаты, не соответствующие запросу, как показано выше. Но если создавать планировщик явно, используя параметры из этого обсуждения (https://github.com/huggingface/diffusers/issues/706), то результаты будут гораздо ближе, к тому, что мы видим при других планировщиках. Спасибо за бдительность: @trix

# change this for DDIMScheduler
scheduler = sched_class.from_config(config)
# to this
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)

DDPMScheduler

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

Другие подходы

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

DPMSolverMultistepScheduler почти не отличим от DPMSolverSinglestepScheduler. LMSDiscreteScheduler явно имеет иную структуру шума на ранних этапах. А PNDMScheduler отказался вызывать callback с сохранением изображения на первом шаге денойзинга.

Заключение

В этой статье мы совершили небольшой обзор архитектуры диффузионной генеративной модели stable diffusion. Кратко затронули основные пункты о роли scheduler в этой архитектуре, что за планировщиком кроется численное решение СДУ или ОДУ. Воспользовались библиотекой diffusers для работы со stable diffusion, применили механизм коллбэков для генерации изображений на промежуточных этапах очистки от шума. И, наконец, подменяли на ходу алгоритмы планировщика, для того, чтобы сравнить результаты генераций в зависимости от выбора алгоритма.

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

• Зарегистрироваться на вебинар