Как работает Stable Diffusion: объяснение в картинках
@designbaza![](/file/3e2ec7c9f1e587b472617.png)
Генерация изображений при помощи ИИ — одна из самых новых возможностей искусственного интеллекта, поражающая людей (в том числе и меня). Способность создания потрясающих изображений на основании текстовых описаний похожа на магию; компьютер стал ближе к тому, как творит искусство человек. Выпуск Stable Diffusion стал важной вехой в этом развитии, поскольку высокопроизводительная модель оказалась доступной широкой публике (производительная с точки зрения качества изображения, скорости и относительно низких требований к ресурсам и памяти).
Поэкспериментировав в генерацией изображений, вы можете задаться вопросом, как же она работает.
В этой статье я вкратце расскажу, как функционирует Stable Diffusion.
![](/file/d43bdc01ddd2d76236765.png)
Stable Diffusion гибка, то есть может использоваться множеством разных способов. Давайте сначала рассмотрим генерацию изображений на основе одного текста (text2img). На картинке выше показан пример текстового ввода и получившееся сгенерированное изображение. Кроме превращения текста в изображение, другим основным способом применения модели является изменение изображений (то есть входными данными становятся текст + изображение).
![](/file/73d86ccae4d3f8c4d99c2.png)
Давайте начнём разбираться со внутренностями модели, потому что это поможет нам объяснить её компоненты, их взаимодействие и значение опций/параметров генерации изображений.
Компоненты Stable Diffusion
Stable Diffusion — это система, состоящая из множества компонентов и моделей. Это не единая монолитная модель.
Изучая внутренности, мы первым делом заметим, что в модели есть компонент понимания текста, преобразующий текстовую информацию в цифровой вид, который передаёт заложенный в текст смысл.
![](/file/fd6fa51f052955aadffce.png)
Мы начнём с общего обзора, а позже углубимся в подробности машинного обучения. Однако для начала можно сказать, что этот кодировщик текста — это специальная языковая модель Transformer (технически её можно описать как текстовый кодировщик модели CLIP). Она получает на входе текст и выдаёт на выходе список чисел (вектор), описывающий каждое слово/токен в тексте.
Далее эта информация передаётся генератору изображений, который состоит из двух компонентов.
![](/file/90a7f7383a9c93e88a7de.png)
Генератор изображений выполняет два этапа:
1. Создание информации изображения.
Этот компонент является секретным ингредиентом Stable Diffusion. Именно благодаря нему возник такой рост качества по сравнению с предыдущими моделями.
Этот компонент выполняется в несколько шагов (step), генерируя информацию изображения. Это параметр steps в интерфейсах и библиотеках Stable Diffusion, который часто по умолчанию имеет значение 50 или 100.
Этап создания информации изображения действует полностью в пространстве информации изображения (или в скрытом пространстве). Подробнее о том, что это значит, мы расскажем ниже. Это свойство ускоряет работу по сравнению с предыдущими моделями диффузии, работавшими в пространстве пикселей. Этот компонент состоит из нейросети UNet и алгоритма планирования.
Слово «диффузия» (diffusion) описывает происходящее в этом компоненте. Это пошаговая обработка информации, приводящая в конечном итоге к генерации высококачественного изображения (при помощи следующего компонента — декодера изображений).
![](/file/d27c018cb2c188a4d73d8.png)
2. Декодер изображений.
Декодер изображений рисует картину на основе информации, которую он получил на этапе создания информации. Он выполняется только один раз в конце процесса и создаёт готовое пиксельное изображение.
![](/file/8f2f0ef50c090e8e8eeaf.png)
На изображении выше мы видим три основных компонента (каждый со своей собственной нейросетью), из которых состоит Stable Diffusion:
- ClipText для кодирования текста.
Входные данные: текст.
Выходные данные: 77 векторов эмбеддингов токенов, каждый в 768 измерениях. - UNet + Scheduler для постепенной обработки/диффузии информации в пространстве информации (скрытом пространстве).
Входные данные: эмбеддинги текста и исходный многомерный массив (структурированные списки чисел, также называемые тензором), состоящий из шума.
Выходные данные: массив обработанной информации - Декодер автокодировщика, рисующий готовое изображение при помощи массива обработанной информации.
Входные данные: массив обработанной информации (размеры: (4,64,64))
Выходные данные: готовое изображение (размеры: (3, 512, 512) — (красный/зелёный/синий, ширина, высота))
![](/file/e7ab3f63ab098e6ed721b.png)
Что такое диффузия?
Диффузия — это процесс, выполняемый внутри розового компонента «image information creator» (этапа создания информации изображения). Имея эмбеддинги токенов, описывающие введённый текст, и случайный начальный массив информации изображения (также они называются latent), процесс создаёт массив информации, который декодер изображения использует для рисования готового изображения.
![](/file/d42804180d3bd42911dd8.png)
Это процесс выполняется поэтапно. Каждый шаг добавляет больше релевантной информации. Чтобы представить процесс в целом, мы можем изучить массив случайных latent, и увидеть, что он преобразуется в визуальный шум. В данном случае визуальное изучение — это прохождение данных через декодер изображений.
![](/file/19ce100c83736fb2ad4a9.png)
Диффузия выполняется в несколько шагов, каждый из которых работает с входным массивом latent и создаёт ещё один массив latent, ещё больше напоминающий введённый текст, а вся визуальная информация модели собирается из всех изображений, на которых была обучена модель.
![](/file/2df68615a94efd823c549.png)
Мы можем визуализировать набор таких latent, чтобы увидеть, какая информация добавляется на каждом из шагов.
![](/file/ab6af75a5a473471fbdbc.png)
Наблюдать за этим процессом довольно увлекательно.
Как работает диффузия
Основная идея генерации изображений при помощи диффузионной модели использует тот факт, что у нас есть мощные модели компьютерного зрения. Если им передать достаточно большой массив данных, эти модели могут обучаться сложным операциям. Диффузионные модели подходят к задаче генерации изображений, формулируя задачу следующим образом:
Допустим, у нас есть изображение, сделаем первый шаг, добавив в него немного шума.
![](/file/8d0cd70f9d3c6639c5e08.png)
Назовём «срез» (slice) добавленного нами шума «noise slice 1». Сделаем ещё один шаг, добавив к шумному изображению ещё шума («noise slice 2»).
![](/file/3e0ecc70dd634def2b30c.png)
На этом этапе изображение полностью состоит из шума. Теперь давайте возьмём их в качестве примеров для обучения нейронной сети компьютерного зрения. Имея номер шага и изображение, мы хотим, чтобы она спрогнозировала, сколько шума было добавлено на предыдущем шаге.
![](/file/7260229b0fd206ebf7292.png)
Хотя этот пример показывает два шага от изображения к полному шуму, мы можем управлять тем, сколько шума добавляется к изображению, поэтому можно распределить его на десятки шагов, создав десятки примеров для обучения на каждое изображение для всех изображений в обучающем массиве данных.
![](/file/08c91d436dd3fd6d20d93.png)
Красота здесь в том, что после того, как эта сеть прогнозирования шума начнёт работать правильно, она, по сути, сможет рисовать картины, удаляя шум на протяжении множества шагов.
Примечание: это небольшое упрощение алгоритма диффузии. На ресурсах по ссылкам в конце статьи представлено более подробное математическое описание.
Рисование изображений устранением шума
Обученный предсказатель шума может взять шумное изображение и количество шагов устранения шума, и на основании этого способен спрогнозировать срез шума.
![](/file/92ef81494bc8972e74fb4.png)
Срез шума прогнозируется таким образом, что если мы вычтем его из изображения, то получим изображение, которое ближе к изображениям, на которых обучалась модель.
![](/file/fe6867c4f13ac1338c105.png)
Если обучающий массив данных состоял из эстетически приятных изображений (например, LAION Aesthetics, на котором обучалась Stable Diffusion), то получившееся изображение будет иметь склонность к эстетической приятности.
![](/file/1706a0492f0ae79756d28.png)
В этом по большей мере и заключается описание генерации изображений диффузионными моделями, представленное в статье Denoising Diffusion Probabilistic Models. Теперь, когда мы понимаем, что такое диффузия, нам понятно, как работают основные компоненты не только Stable Diffusion, но и Dall-E 2 с Google Imagen.
Обратите внимание, что описанный выше процесс диффузии генерирует изображения без использования текстовых данных. В последующих разделах мы расскажем, как в процесс внедряется текст.
Увеличение скорости: диффузия сжатых (скрытых) данных, а не пиксельного изображения
Для ускорения процесса генерации изображений Stable Diffusion (по информации из исследовательской статьи) выполняет процесс диффузии не с самими пиксельными изображениями, а со сжатой версией изображения. В статье это называется «переходом в скрытое пространство».
Это сжатие (и последующая распаковка/рисование) выполняется при помощи автокодировщика. Автокодировщик сжимает изображение в скрытое пространство при помощи своего кодировщика, а затем воссоздаёт его при помощи декодера на основе только сжатой информации.
![](/file/1b453314ec0cbda45774e.png)
Далее со сжатыми latent выполняется прямой процесс диффузии. Используются срезы шума, применяемые к этим latent, а не к пиксельному изображению. То есть предсказатель шума на самом деле обучается прогнозировать шум в сжатом описании (в скрытом пространстве).
![](/file/88b31a6fc0fd7fa1ab1cc.png)
При помощи прямого процесса (с использованием кодировщика автокодировщика) мы генерируем данные для обучения предсказателя шума. После его обучения мы можем генерировать изображения, выполняя обратный процесс (при помощи декодера автокодировщика).
![](/file/1b03ccfa8bd4c09e3404d.png)
Эти два потока показаны на рисунке 3 статьи про LDM/Stable Diffusion:
![](/file/ab9f4f8f6aaa940b861aa.png)
Также на этом рисунке показаны компоненты «согласования», которые в данном случае являются текстовыми строками, описывающими изображение, которое должна генерировать модель. Поэтому давайте рассмотрим эти текстовые компоненты.
Текстовый кодировщик: языковая модель Transformer
Языковая модель Transformer используется в качестве компонента понимания языка, она получает текстовую строку и создаёт эмбеддинги токенов. В опубликованной модели Stable Diffusion используется ClipText (модель на основе GPT), а в статье применяется BERT.
В статье, посвящённой Imagen, показано, что выбор языковой модели важен. Замена на более объёмные языковые модели сильнее влияет на качество генерируемого изображения, чем более объёмные компоненты генерации изображений.
![](/file/2a40db157af812c50c6a2.png)
Улучшение/увеличение языковых моделей существенно влияет на качество моделей генерации изображений. Источник: Статья про Google Imagen, написанная Saharia и соавторами. Рисунок A.5.
В первых моделях Stable Diffusion просто подключалась предварительно обученная модель ClipText, выпущенная OpenAI. Возможно, будущие модели перейдут на новые и гораздо более объёмные OpenCLIP-варианты CLIP. В эту новую группу входных векторов включены текстовые модели размерами до 354 миллионов параметров, в отличие от 63 миллионов параметров в ClipText.
Как обучается CLIP
CLIP обучается на массиве изображений и подписей к ним. Массив данных выглядит примерно так, только состоит из 400 миллионов изображений и подписей:
![](/file/9f6d63b74137fa91e8309.png)
CLIP — это сочетание кодировщика изображений и кодировщика текста. Обучающий процесс модели можно упрощённо представить как кодирование изображения и его подписи кодировщиками изображений и текста.
![](/file/dba82bbb1c798c5b3e10f.png)
Затем мы сравниваем получившиеся эмбеддинги при помощи косинусного коэффициента. В начале процесса обучения схожесть будет низкой, даже если тест описывает изображение правильно.
![](/file/52c0c4e21c2e53fa1dbfe.png)
Мы обновляем две модели так, чтобы в следующий раз при создании эмбеддингов получившиеся эмбеддинги были схожими.
![](/file/1cf80d60e7a78bc3a6698.png)
Повторяя этот процесс со всем массивом данных и группами входных векторов большого размера, мы получаем кодировщики, способные создавать эмбеддинги, в которых изображение собаки и предложение «a picture of a dog» схожи. Как и в word2vec, процесс обучения также должен включать в себя отрицательные примеры изображений и подписей, которые не совпадают, а модель должна присваивать им низкую оценку схожести.
Передача текстовой информации в процесс генерации изображений
Чтобы сделать текст частью процесса генерации изображений, нам нужно модифицировать предсказатель шума так, чтобы он использовал в качестве входных данных текст.
![](/file/215759b3e1adbdbe9c794.png)
Теперь наш массив данных содержит закодированный текст. Так как мы работаем в скрытом пространстве, то входные изображения и прогнозируемый шум находятся в скрытом пространстве.
![](/file/3ce307665840416a9b192.png)
Чтобы лучше понять, как текстовые токены используются в Unet, давайте глубже разберёмся с Unet.
Слои предсказателя шума Unet (без текста)
Для начала рассмотрим диффузионную Unet, не использующую текст. Её входы и выходы выглядят так:
![](https://habrastorage.org/r/w1560/getpro/habr/post_images/bd0/800/406/bd0800406e5c9e8b436142fec2d0e0d5.png)
Внутри мы видим следующее:
- Unet — это последовательность слоёв, работающая над преобразованием массива latent.
- Каждый слой обрабатывает выходные данные предыдущего слоя.
- Часто выходных данных подается (через остаточные соединения) для обработки на дальнейших этапах сети.
- Шаг времени преобразуется в вектор эмбеддингов шага времени, который используется в слоях.
![](https://habrastorage.org/r/w1560/getpro/habr/post_images/762/e0a/e1e/762e0ae1e8403206f41a7329022cddc8.png)
Слои предсказателя шума Unet с текстом
Давайте посмотрим, как изменить эту систему, чтобы уделить внимание тексту.
![](https://habrastorage.org/r/w1560/getpro/habr/post_images/419/112/425/419112425321728bfb3718b7e1a6dd5c.png)
Основное изменение системы, которое необходимо для добавления поддержки текстового ввода (техническое название: text conditioning) — добавление слоя attention между блоками ResNet.
![](https://habrastorage.org/r/w1560/getpro/habr/post_images/717/cf0/103/717cf0103b0ce5ece016412cd542aea7.png)
Обратите внимание, что блок resnet не смотрит непосредственно на текст. Слои attention объединяют эти текстовые описания в latent. И теперь следующий ResNet может использовать эту встроенную текстовую информацию в своей обработке.
Заключение
Надеюсь, это даст вам поверхностное понимание работы Stable Diffusion. В ней задействовано множество других концепций, но я считаю, что их проще понять, если вы знаете описанные выше строительные блоки. Для дальнейшего изучения можно воспользоваться представленными ниже полезными ресурсами.
![](/file/ac965c2048a24458030d1.png)