Исходная публикация с описаниеим метода [Mar 2023] — DreamBooth: Fine Tuning Text‑to‑Image Diffusion Models for Subject‑Driven Generation

Что такое stable-diffusion

Stable Diffusion — новаторская модель генерации изображений на основе искусственного интеллекта. Она использует метод диффузии для постепенного улучшения качества сгенерированных изображений.

В основеStable Diffusion лежит архитектура autoencoder. Сначала модель «кодирует» текстовое описание в вектор признаков. Затем декодер поэтапно создает изображение, уточняя его с каждым шагом. Этот процесс напоминает диффузию — плавное распространение частиц.

Что такое дообучение

Нейронные сети «получают информацию» об окружающем мире при обучении. Создатели модели генерируют огромный датасет картинок с их описанием, и загоняют их в модель. Однако если вы хотите генерировать какой то конкретный обьект — например лично вас, вашего питомца, родной дом, картины любимого автора в особом стиле итп.
Вы можете загнать ваши новые изображения в модель, придумав для них описание, содержащие названия вашего объекта (или лучше какой‑то специальный токен, условно [V]), и скоректировав веса градиентным спуском, так что бы модель понимала, что скрывается за этим токеном. Однако такое дообучение дифузионных моделей имеет свои недостатки

Проблемы

1. Language drift — модель предварительно обученная на большом корпусе текста и позже точно настроенная для конкретной задачи, постепенно теряет синтаксические и семантические знания языка.

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

Основная идея решения

Авторы статьи решают эту задачу дополнительным лоссом — class‑specific prior preservation loss. Вначале обучаемая модель пытаться сгенерировать нашу собаку, и обучается по лоссу с фотографиями из выборки для дообучения. Затем она же генерирует произвольную собаку, и лосс считается с фотографиями, сгенерированными исходной моделью. Таким образом, мы наказываем модель, за то что она забывает абстрактные признаки собак, концентрируясь на обучающей выборке. Более лаконично на схеме:

Или более строго прибавку к лосу можно записать в виде:

Насколько хорошо это работает?

Как видно из картинок ниже, добавление class‑specific prior preservation loss очень позитивно сказалось на возможности модели генерировать исходный объект в разнообразных вариациях

Просто сравните с тем что было выше!

Так же модель научилась видоизменять обьект в зависимости от промпта, генерируя полностью новые объекты на основе существующего:

При этом сама модель не потеряла возможности генерировать обычные объекты (произвольных собак, отличных от той, на которой происходило до обучение)

А так же такой подход очень хорошо сработал для задач style‑transfer конкретного объекта. Авторы смогли сгенерировать картинки в разных стилях, но с исходной собакой.

LoRA

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

Воспроизведение результатов

Я до обучил stable‑diffusion‑v1–5 на двух собственных датасетах — фотографии кошки (cloody‑cat) и декоративного гриба. Обучение проходило на видеокарте A100 (занимая примерно 1/8 видеопамяти) в течении 2х минут. Интерактивная демонстрация работы моей модели на Hugging Face — freQuensy23/diffusion‑cloody — там можно потыкать свои запросы бесплатно (за такое им огромный респект)! В таблице ниже показано как со временем обучения меняется результат генерации по промпту «A [V] cat», а более красочно вы можете посмотреть это в анимации в начале этой статьи или тут. Так же я проверил данную систему на генерации этого кота в очень разных условиях и стилях. Дальше будет просто фотогалерея

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

P. S. Если есть ошибка в правописании — сообщите мне через стандартную систему хабра — я исправлю. Мне бывает сложно писать на русском.