Что делать, если хочешь отредактировать картинку?
На фотошоп не хватает знаний, времени или сил. Inpaint-модели подводят. Конструкторы вроде ComfyUI требуют столько всего, что проще изучить фотошоп.
Как было бы здорово показать картинку нейросети и написать: «Пожалуйста, смени юбку на джинсы. И сзади ещё одну клумбу добавь… И погода пусть солнечной будет!»
И, вроде бы, это не выглядит нереальным. Современные мультимодальные нейросети уже довольно хорошо работают с изображениями: находят и описывают объекты, понимают визуальный контекст, проводят сложные рассуждения — но не умеют создавать изображения. Генеративные же модели создают потрясающие изображения, но очень плохо понимают человеческие формулировки и не улавливают контекст.
Кажется, что два подхода друг друга отлично дополняют, нужно только их правильно собрать. Но всё не так просто. Как бы ни ломали голову исследователи, задача произвольного нейросетевого редактирования изображений, увы, на сегодняшний день в целом не решена.
Однако не всё так плохо! В попытках обучить нейросети манипулировать изображениями уже есть заметные достижения, о которых нельзя не рассказать.
Начинаем путь
Всем привет, с вами команда The Layer!
Это первая из цикла статей, где мы пройдём полный путь задачи редактирования изображений посредством нейросетей: от «а что, если просто перегенерить с фиксированным random seed?» до огромных мультимодальных моделей.
В этой статье мы рассмотрим тривиальные попытки редактирования (фиксирование random seed, инпейнтинг, зашумление-расшумление), разберем методы DDIM Inversion и Null-Text Inversion.
Мы постарались сделать эту статью понятной для всех, потому добавили напоминалку в начало, где коротко и ясно расписали релевантную теорию: диффузия, механизм внимания и так далее. Будет полезно как новичкам, так и подзабывшим. А кто в теме — смело пропускайте!
Задача редактирования изображений
Итак, есть стремление обучить модель, способную выполнять сложные, осмысленные изменения в структуре, стиле, семантике изображения по указаниям пользователя. Указания бывают разными, но нас интересует самый захватывающий кейс — инструкции в естественных формулировках. Например, вот так:
Задача на самом деле очень сложная. Модель должна «понять»:
где на картинке передний план, где стол, а где — справа от стола;
как выглядит фарфоровая кукла, как выглядит платье в стиле 18 века, как будет выглядеть миниатюра такого платья на кукле;
как корректно «вписать» куклу в изображение с точки зрения и законов физики, и эстетических соображений, при этом не изменив остальные детали изображения;
где источники света на изображении, какие части куклы будут освещены / затенены, какую тень (и куда) будет отбрасывать кукла.
Так много условий нужно выполнить для корректного добавления крупного объекта в центр изображения — и это ещё сравнительно простая операция. Получается, чтобы справляться с произвольными задачами редактирования, модель должна обладать глубинными познаниями о самых разных аспектах нашего мира и о физических законах. И уметь этими знаниями пользоваться.
В задаче редактирования речь, конечно, не только о добавлении объектов. От моделей-редакторов ожидается выполнение:
семантических изменений (добавление, удаление, замена объекта или фона, смена выражения лица);
стилистических изменений (замена цвета / текстуры объекта, общая смена стиля);
структурных изменений (перемещение объекта, смена размера / формы, замена действия / позы, смена ракурса).
Релевантная теория
Всё, что потребуется для дальнейшего понимания нашего обзора. Здесь можно освежить память и подготовиться к погружению в задачу редактирования.
Диффузия
Можем ли мы создать генератор изображений?
Конечно, можно написать функцию, задающую каждый пиксель произвольным значением… Но разве мы хотим генерировать бессмысленные наборы пикселей, простой шум? Осмысленные изображения формируются по огромному множеству правил и закономерностей, отражающих как физические законы нашего мира, так и эстетические, смысловые особенности.
Конечно, это распределение очень многомерно
Говоря языком статистики, всё множество осмысленных изображений, которые в принципе могут существовать, формирует некое невероятно сложное распределение.
Знай мы закон этого распределения, могли бы из него семплировать — генерировать такие произвольные значения, которые гарантированно будут распределены по этому закону.
В целом это кажется возможным: если собрать из интернета несколько миллиардов изображений, то получится довольно разнородная, репрезентативная выборка. Такой выборки достаточно для некоторого приближения искомого распределения (для любознательных: генеративные модели на основе VAE), однако у такого подхода много проблем, в том числе — с качеством результатов.
Но есть и другой способ, называемый диффузионной моделью. Существуют ведь распределения с известным нам законом: стандартное нормальное, например. Семплировать произвольные значения такого распределения (генерировать гауссов шум) довольно легко. Так давайте выучим отображение из гауссова шума в распределение изображений! Тогда, семплировав произвольное значение из стандартного нормального распределения и выполнив отображение, мы будто бы семплируем из искомого нами распределения изображений:
Однако простое отображение из чистого шума в готовое изображение было бы слишком сложным. Основное преимущество диффузии в постепенности, итеративности процесса. Мы учим модель двигаться от шума к изображению небольшими шагами через множество промежуточных состояний.
Синтезировать такие промежуточные состояния для обучения можно, всего лишь накладывая в произвольной степени шум на изображения из обучающего датасета:
Процесс обучения теперь несложно представить: учим модель «расшумлять» эти промежуточные состояния обратно в сторону оригиналов. Считаем лосс между предсказаниями модели и оригиналами, обновляем веса модели:
После долгого обучения подобной модели на большом и разнородном наборе изображений можно надеяться, что бесчисленное множество траекторий «шум-изображение» позволит модели обобщить некие законы движения из гауссова шума в распределение осмысленных изображений. Тогда, семплировав произвольный гауссов шум и подав обученной диффузии на вход, мы можем надеяться, что через множество итераций расшумления модель приведёт этот шум к какому-либо осмысленному изображению (почти наверняка не существовавшему ранее):
Обратим внимание ещё раз, что процесс генерации итеративный. Даже когда модель уже обучена, не выйдет просто «прыгнуть» от случайного шума к готовому изображению. Чем больше шума в изображении, тем хуже «предположение» модели о том, как должен выглядеть финальный, полностью расшумлённый вариант:
Ремарка: есть способы, позволяющие генерировать изображения в 1–4 шага (например, дистилляция), но это за пределами нашего ликбеза.
Обобщённый процесс обучения выглядит так:
Берём несколько изображений из датасета (желательно большого и разнородного).
Зашумляем их в произвольной степени. Важно, чтобы модель хорошо справлялась с устранением как незначительного, так и сильного шума. Такие промежуточные шаги помогут модели «проложить путь» от распределения чистого шума к распределению чистых изображений.
Подаём зашумлённые изображения в модель.
Для каждого из них модель предсказывает шум, который нужно вычесть, чтобы получить чуть менее зашумлённое изображение.
Вычитаем предсказанный шум, сравниваем результат с оригиналом из датасета, обновляем веса модели.
Архитектура модели обычно представляет собой автокодировщик U-Net:
Латентное пространство и латентная диффузия
Классическая диффузия приближает распределение пикселей. То есть, например, её можно обучить на работу с трёхканальными картинками 32 на 32 пикселя: модель приблизит 3072-мерное распределение (32 × 32 × 3). Результирующее изображение обычно увеличивается специализированной моделью.
Получается, мы ограничены размерами 32 на 32 и вынуждены подсчитывать градиенты для каждого пикселя. Это не только добавляет много бессмысленного труда, но и лишает возможности масштабироваться: при увеличении разрешения картинок вычислительная сложность возрастает экспоненциально.
Можно ли как-то сжать изображение, не потеряв при этом смысл его содержимого? Вспомним свёрточные нейросети: каскадом последовательных свёрток из изображения извлекается малое число признаков, оптимальных для той или иной задачи.
Однако что если поставить эту задачу как максимальное сжатие с минимальными потерями? То есть обучить модель так уменьшать размерность изображения, чтобы его затем можно было (почти) идеально восстановить.
Подобная задача решается автокодировщиками:
Автокодировщик обучается последовательными свёртками извлекать из изображения такие ключевые признаки, которых (при знании общих закономерностей внутри изображений датасета) будет достаточно для удовлетворительного восстановления исходного изображения.
Эти признаки называются латентными, то есть скрытыми, так как они представляют собой ключевые черты данных, «скрытые» в них. И ещё потому, что зачастую они неинтерпретируемы человеком. Пространство таких признаков называется латентным пространством. А автокодировщик — это отображение из пространства изображений в латентное пространство и затем обратно в пространство изображений.
Например, если обучить особый тип автокодировщика на «сжатие» всего датасета MNIST в 2 признака и изобразить полученное двумерное пространство, то получится нечто подобное:
Обученный автокодировщик приближает некоторое непрерывное распределение ключевых признаков изображения, позволяя затем эффективно восстанавливать изображения из такого распределения.
С учётом этого ключевая идея латентной диффузии — давайте «обернём» диффузию в обученный автокодировщик, чтобы модель работала не с изображениями, а с их латентными представлениями. Латентная диффузия приближает не распределение пикселей, а распределение извлекаемых из них признаков.
Механизм внимания, обусловленная генерация
Допустим, мы обучили модель диффузии. Можно ли как-нибудь теперь обусловить генерацию? Семплировать не произвольные изображения, а с заданным нами смыслом — например, по текстовому описанию?
Да, обучить модель с подобным функционалом вполне реально. Из текстов, как и из картинок, также можно формировать латентные представления — эмбеддинги. Теперь достаточно на каждом шаге обучения подавать на вход не только зашумлённое изображение, но и эмбеддинг его текстового описания. Можно надеяться, что модель сформирует связь между текстовыми смыслами и конкретными участками искомого распределения латентов изображений, — и тогда обусловленное текстом семплирование «приведёт» процесс генерации куда-то близко к искомому смыслу.
Обучаемый механизм связи между текстовыми эмбеддингами и латентными представлениями изображений представляет собой механизм внимания, а точнее, cross-attention.
Если абстрагироваться от формул, механизм внимания — это связь между эмбеддингами, где каждый обладает тремя обучаемыми параметрами: Q (query, запрос к прочим эмбеддингам), K (key, насколько сильно эмбеддинг «откликается» на тот или иной запрос), V (value, значение). Популярной аналогией механизма внимания является поиск данных в бумажном архиве: query — карточка с нашим запросом, key — релевантность каждой папки относительно нашего query. Пропорционально значениям key мы достаём из папок бумаги — value.
Допустим, студент готовит презентацию по обработке естественного языка (NLP). Здесь карточка NLP — это запрос студента, пришедшего в картотеку. Разные папки бумаг сильнее или слабее резонируют с NLP, что измеряется их ключами (обучаясь, механизм внимания приближает значения ключей). Пропорционально значениям ключей студент набирает из папок бумаги — это значения.
Пускай тогда запрос Q будет получаться из латента изображения, а ключ K и значение V — из эмбеддингов текста. Всё, механизм обусловленной генерации изображений готов к обучению.
Classifier-Free Guidance
Итак, мы научились обуславливать диффузию! Действительно, изображения теперь совпадают по содержанию с заданными текстовыми условиями:
Как видим, качество оставляет желать лучшего: изображения слишком разнообразны. Модель будто бы «понимает» запрос, но генерирует со слишком большим разбросом.
В качестве решения проблемы было предложено дополнительно к диффузии обучать классификатор на зашумлённых картинках. Исследователи выяснили, что для корректировки процесса семплирования можно использовать градиенты классификатора (производные; вычисляются в ходе обучения, указывают в сторону оптимального решения). Это уменьшает разнообразие генерируемых изображений, при этом повышая их качество. Приятный бонус — исходная диффузионная модель никак не меняется.
Поскольку обучаемый классификатор как бы ведёт (guide) процесс генерации, метод называется Classifier Guidance.
По изображениям видим, что иногда наличие guidance глобально меняет изображения, иногда только едва корректирует — но в среднем заметно улучшает качество.
Проблемой такого подхода является необходимость дополнительно к модели диффузии обучать и модель классификатора. Классификатор должен быть обучен на зашумлённых данных, поэтому предобученный взять не выйдет.
В поисках альтернатив был предложен следующий метод: давайте обучать условную диффузию, но иногда условие (текст, метка класса, …) оставлять пустым. Так мы будем учить одну и ту же модель как на безусловную, так и на условную генерацию (будто бы две модели). С такой моделью генерацию изображения можно провести в два параллельных процесса: безусловный и условный. Так, на каждом шаге можно вычитать условное предположение из безусловного в некоторой контролируемой пропорции, что приведёт к схожему с Classifier Guidance эффекту: меньше разнообразия, выше качество. Поскольку обучение классификатора теперь не требуется, метод носит название Classifier-Free Guidance.
В дальнейшем было доказано, что Classifier-Free Guidance способствует большему улучшению качества, нежели Classifier Guidance.
Вот и подошёл к концу наш ликбез! В конце статьи мы собрали список полезных источников для желающих копнуть поглубже.
Тривиальные попытки редактирования
Итак, можно приступать к задаче редактирования! Сначала рассмотрим самые примитивные идеи, и почему они не работают:
Идея: перегенерировать с фиксированным random seed, но с другим промптом
Идея заключается в фиксировании random seed генератора шума. Так каждый процесс генерации будет начинаться с одного и того же шума. Это позволяет надеяться, что изображения будут похожи, а единственное значимое изменение траектории семплирования будет вызвано редактированием промпта — и приведёт к желаемым изменениям в изображении.
Однако на практике, к сожалению, так не выходит: изменяется всё изображение.
Фиксирование random seed действительно частично сохраняет некоторые участки изображения (трава в самом низу картинок), но в целом задний план сильно меняется. И это ещё очень простая структура картинки. Кроме того, если нужно отредактировать реальное изображение, seed зафиксировать не получится.
Идея: инпейнтинг
«Вырезать» кусок изображения (выделить его маской) и сгенерировать новое наполнение, опираясь на оставшуюся часть.
Эта методика действительно почти не меняет не выделенные участки, но удаление участка изображения зачастую отбрасывает полезную информацию: поза, форма, цвет, детали. Это делает любое редактирование, требующее сохранения структуры, невозможным (смена цвета или текстуры, деформация, частичное добавление, частичное удаление, вращение, стилизация).
При этом даже простые задачи редактирования выполняются очень нестабильно или не выполняются вовсе:
Более того, рисовать маску вокруг сложных объектов может быть довольно долгим и трудозатратным занятием.
Однако стоит отдать должное инпейнтингу, для простых изображений (или очень настойчивых пользователей) это всё ещё один из лучших подходов к некоторым задачам редактирования.
Идея: зашумить исходное изображение и расшумить с изменённым промптом
Слабого зашумления недостаточно для внесения изменений. При более сильном зашумлении, наоборот, теряется слишком много информации — изменится всё изображение.
Заметьте, что не только кот вышел жутковатым, но и лужайка с фоном сильно изменились.
Однако идея воспользоваться шумоподавляющей природой диффузии для редактирования изображений оказалась достаточно перспективна — и при некоторой доработке сформировала один из первых нетривиальных методов редактирования. Давайте изучим его.
DDIM Inversion
Если не считать примитивный инпейнтинг и его вариации редактированием, то DDIM Inversion — первопроходец в данной задаче.
Идея инверсии — «обратить вспять» генерацию изображения латентной диффузией. Обычно мы генерируем изображение, начиная с чистого шума. Мы итеративно предсказываем шум и «вычитаем» его, постепенно формируя изображения.
В инверсии же мы начинаем с готового изображения и предсказываем для него шум, но вместо «вычитания» его, наоборот, «прибавляем» этот шум, как бы возвращая изображение на предыдущие шаги генерации.
После множества итераций инверсии мы хотим получить латентный шум, который при использовании в качестве отправной точки для обычной процедуры генерации приведет к созданию оригинала.
Получилось немного не то, что мы ожидали…
Дело в том, что для корректной работы инверсии предполагается, что предсказанный шум будет одинаков для любых шагов t и t-1 — а это будет тем ближе к правде, чем больше итераций обращения мы выполним (не первые 50–100).
Но можно схитрить и начать со, скажем, пятидесятого шага генерации с теми промежуточными латентными значениями, которые мы сохранили:
Результат очень близок к нашему исходному изображению.
Таким образом, инверсия — это процесс «обращения» генерации изображения. Мы «возвращаем» изображение к такому латентному шуму, из которого модель сможет восстановить это изображение с точностью.
Зачем мы это делаем? Мы надеемся, что если выполнить инверсию над изображением, а затем произвести из полученного шума генерацию с немного изменённым промптом, то получится изображение, соответствующее исходному везде, кроме изменённых в промпте мест. Например, заменив puppy на cat, мы надеемся увидеть кошку с идентичной лужайкой и фоном.
Это и есть DDIM Inversion.
Так почему же тривиальный способ зашумления-расшумления не работает, а предсказанный моделью работает весьма неплохо?
Случайно распределённый шум искажает латенты случайным образом, и модель не может идеально обратить такие изменения. А шум, предсказанный самой моделью, получается, модель обратить может?
Смотря какая модель. Например, классические диффузионные модели не могут. Однако одной из ключевых особенностей моделей DDIM является детерминированный процесс генерации (финальное изображение зависит только от изначального шума). Это и делает инверсию возможной, поэтому и DDIM Inversion (хотя можно сказать, что идея инверсии перекочевала в диффузию из GAN-ов).
Хоть результаты DDIM Inversion на простых изображениях сопоставимы с хорошим инпейнтингом, в общем случае инверсия часто заставляет желать лучшего. Проблема в том, что между предсказанным в процессе инверсии шумом и ожидаемым при генерации шумом всё же есть некоторая разница, — и зачастую модель неспособна удовлетворительно эту разницу преодолеть.
Мы рассматривали этот метод на очень простых изображениях. Да и то, кошка в результате получилась с некоторыми артефактами. Вот как инверсия работает на произвольных, более сложных изображениях:
И это просто инверсия, без редактирования. А при редактировании подобные несоответствия зачастую возрастают в разы за счет усиления Classifier-Free Guidance.
Исследователи предпринимали неоднократные попытки разрешить подобные проблемы. Давайте вкратце рассмотрим одно из ключевых усовершенствований.
Null-text Inversion
Данный метод адресует ключевую проблему инверсии, изображённую на схеме: вместо генерации «по направлению» инверсии траектория искажается и смещается в силу сложности идеального обращения шумов. Метод Null-text Inversion акцентируется на приближении траектории генерации (светло-серой) к желаемой (тёмно-серой).
Как это происходит?
Рассмотрим следующую схему:
Если не обращать внимания на красные стрелки, это схема классического Classifier-Free Guidance (CFG). Напомним идею CFG. Диффузия генерирует два изображения:
одно условно, с текстовым описанием (верхняя часть схемы);
другое — безусловно, с «пустым» текстовым описанием (null-text, нижняя часть схемы).
Разница между этими изображениями позволяет направить генерацию в сторону более строгого следования тексту с силой w (правая часть схемы).
Обычно любые усовершенствования CFG сводятся к оптимизации условной генерации, но при данном подходе решили, наоборот, оптимизировать безусловную часть, заменив Null-text эмбеддинги на оптимизированные посредством MSE между предсказанным шумом и реальным.
Таким образом удалось значительно улучшить качество инверсии:
Очевидно, и качество редактирования возросло:
Однако даже на этих черрипиках из статьи заметно, что сложные запросы выполнены совсем не идеально:
Хотелось бы более строгих масок, более естественных изменений — и всё это без надобности вручную вырисовывать эти самые маски.
Как раз во всём этом и заключается совершенно иной подход, оперирующий актуальными и по сей день идеями. Ему мы и посвятим следующую часть. Расскажем про становление foundation-подходов к редактированию: Prompt-to-Prompt, InstructPix2Pix.
Спасибо, что дочитали до конца!
Если статья понравилась, да и в целом вам интересны качественные разборы интересных научных работ, без воды, флуда и рекламы — то подписывайтесь на наш Телеграм-канал!
Полезные материалы:
https://fleuret.org/public/lbdl.pdf — карманный справочник, если нужно освежить память по глубокому обучению.
Математическая постановка задачи диффузии очень сложна, потому желающим в ней разобраться рекомендуем до знакомства с оригинальными статьями вроде DDPM и DDIM углубиться в следующие материалы:
https://mhsung.github.io/kaist-cs492d-fall-2024/ — свежий всеобъемлющий курс по диффузионным моделям от Корейского института KAIST.
https://lilianweng.github.io/posts/2021-07-11-diffusion-models/ — подробное математическое объяснение диффузионных моделей от Lilian Weng, former Lead Safety Researcher в OpenAI.
https://udlbook.github.io/udlbook/ — свежая серьезная книга по глубокому обучению от профессора Лондонского университета (доступна бесплатно) — в том числе есть раздел про диффузию.
Конкретно по рассмотренным нами методам редактирования:
https://huggingface.co/learn/diffusion-course — в целом полезный курс, часть изображений и объяснений по DDIM Inversion взята отсюда.
https://arxiv.org/pdf/2211.09794 — статья Null-text Inversion, авторы сами очень доходчиво объяснили свой метод.
