TL;DR

Это новая модель, похожая на Stable-Diffusion (доработанная из SD2) для изображений аниме, поддерживающая разрешение 768x768 и значительно превосходящая некоторые популярные модели аниме, согласно нашим оценкам. Чекпоинты общедоступны (инструкция по установке), демо доступно некоторое время.

вступление

После моих экспериментов с обучением моделей диффузии с нуля несколько человек указали, что существует множество доработанных версий Стабильная диффузия на одних и тех же/похожих наборах данных. В частности:

  • Waifu Diffusion: он использует относительно небольшое подмножество данных и просто передает приглашение тега в CLIP для обработки) без каких-либо изменений кода исходной стабильной диффузии (это проблематично для ИМО). Тем не менее, на данный момент это самая популярная модель распространения аниме. Он выпустил КПП для публики, в отличие от следующей модели;
  • NovelAI: этим занимается настоящая компания, а не независимые исследователи/инженеры. Они не выпускают модель, а предоставляют к ней платный интерфейс. По-видимому, была также некоторая драма с их контрольной точкой и утечкой кода, из-за чего, похоже, у них были небольшие модификации оригинальной Stable Diffusion;
  • Anything-v3: этот появился чуть позже, по модели мало информации, но есть подозрение, что она доработана из WaifuDiffusion (по крайней мере, она использует аналогичные быстрые трансформации перед подачей в CLIP). Он производит очень подробные и красивые изображения, но не очень точно соответствует подсказке. Его выходное пространство очень маленькое (отсутствие разнообразия, все изображения имеют одинаковый стиль и похожи друг на друга) и напоминает модели GAN до диффузии. В отличие от других моделей, эта также не способна генерировать NSFW — поэтому, возможно, она была получена путем тонкой настройки небольшого набора высококачественных изображений одним художником.

У меня есть некоторые проблемы с архитектурой всего этого — самое главное тот факт, что они используют замороженный CLIP для кондиционирования. Это кажется очень неоптимальным — CLIP обучен для текста и плохо подходит для использования с наборами тегов. Играя с этими моделями, проблемы становятся совершенно очевидными — если по качеству они лучше, чем Анифузия, то точность подсказок во многих случаях хуже.

Однако, учитывая обилие моделей в этой сфере, поначалу я решил не заморачиваться и позволить прогрессу идти своим чередом. Но сейчас вышла Stable Diffusion 2, а видимого прогресса в аниме-моделях до сих пор нет. Поэтому я решил попробовать — в этом посте я описываю, как я настроил Stable Diffusion 2 с помощью специального кондиционера подсказок, и покажу результаты. Разумеется, модель тоже публикую — из нескольких опробованных мной подсказок она оказалась лучше перечисленных выше моделей. Следующие несколько разделов посвящены техническим деталям, если вам это неинтересно, вы можете сразу перейти к Результатам.

Фон

Чтобы понять, что происходит и что изменилось, давайте сначала кратко рассмотрим существующие модели.

Диффузионные модели

Не буду вдаваться в детали диффузионных моделей, хороших отзывов о них в инете уже полно. Для наших целей достаточно знать, что идея состоит в том, чтобы разложить процесс генерации изображения на последовательность шагов шумоподавления. Во время обучения магистраль модели (в случае изображений — обычно UNet) получает поврежденное изображение и шаг шумоподавления и пытается отменить один шаг повреждения. Во время логического вывода мы начинаем со случайного шума и многократно применяем магистраль, чтобы отменить все шаги воображаемого искажения. Точная математика, стоящая за этим, должна быть точной. Во время вывода большинство современных диффузионных моделей модифицируют процесс, чтобы уменьшить количество шагов выборки (это стало возможным благодаря тому, что кто-то заметил, что процесс подозрительно похож на решение стохастического дифференциального уравнения с помощью численного метода, и применил более умный решатель DE). вместо).

Скрытая диффузия

Это, пожалуй, самая важная идея, используемая в большинстве моделей открытой диффузии. Такие компании, как Google или OpenAI, могут позволить себе обучать и использовать модели распространения в исходном пространстве пикселей, потому что у них много оборудования. Но для неспециалиста обучение (или даже запуск) UNet в пространстве 512х512 пикселей очень дорого. Идея состоит в том, чтобы обучить автоэнкодер, который сжимает изображение из исходного пространства 512x512x3 в скрытое пространство (в случае Stable Diffusion — 64x64x4). Автоэнкодер обучен как VAE с регуляризацией на основе KL, чтобы сохранить скрытое пространство хорошим.

Затем модель выполняет диффузию в этом пространстве 64x64x4, что, очевидно, намного эффективнее. Недостатком является то, что VAE неидеален, и часть информации теряется при сжатии изображения. Но на практике это работает очень хорошо.

FWIW, возможны альтернативные подходы, например. Imagen генерирует низкоразмерное изображение в пространстве пикселей, а затем постепенно масштабирует его. Однако это не относится к модели, которую мы описываем в этой статье.

Стабильная диффузия

Это был скорее социальный, чем технический прорыв. Модель была ванильной латентной диффузионной моделью, обученной на LAION, и, в отличие от большинства предыдущих моделей генерации изображений, она была опубликована. Это сделало технологию доступной для широкой аудитории, не знакомой с машинным обучением, что вызвало много творчества, дискуссий и людей, которые настраивали ее на различных наборах данных (например, упомянутые выше аниме-модели).

Один интересный выбор в архитектуре, который очень актуален для нас, — это то, как модель выполняет кондиционирование. Обычно модели обучаются как безусловным (UNet(поврежденное изображение) ~ изображение), так и условным (UNet(поврежденное изображение, условие) ~ изображение), а при выборке используется управление без классификатора: вместо использования UNet(изображение повреждено, состояние) при создании изображений , он использует C * UNet(изображение повреждено, условие)—(C — 1)* UNet(изображение повреждено), а C значительно выше 1 (например, 7). Это не очень интуитивно понятно, но значительно улучшает как качество изображения, так и оперативную достоверность сгенерированных изображений.

Нам важно, как достигается условие. В Stable Diffusion используют последний скрытый слой кодировщика текста модели CLIP. Он заморожен, т.е. градиенты не распространяются на него во время обучения. Его можно рассматривать как модель языкового кодировщика с некоторыми вторичными вариантами изображения, смещающими понимание языка в сторону визуального понимания.

Диффузия Вайфу

Это стабильная диффузия, настроенная на аниме-изображения. Похоже, они использовали подмножество Danbooru2021 (или что-то подобное) в качестве своего тренировочного набора. Они использовали Stable Diffusion как есть, без каких-либо изменений в коде (однако они настроили часть декодера VAE, чтобы он немного лучше работал с аниме-изображениями).

В результате они аналогично используют CLIP в качестве кондиционера, передавая ему строку с тегами в качестве подсказки. Я предполагаю, что это сработало довольно плохо, поскольку они также предварительно обрабатывают теги (например, заменяют символы подчеркивания пробелами). Но даже при этом оказывается, что CLIP понимает подсказку неоптимально.

Стабильная диффузия 2

Это инкрементальный релиз, который переобучил модель (того же размера) с нуля, а также заменил кондиционер на OpenCLIP. Он также имеет некоторые интересные дополнительные функции, такие как img2img на основе глубины и модель с более высоким разрешением (768x768 пикселей).

Наша модель: обучение

Цель состоит в том, чтобы точно настроить Stable Diffusion 2 на аниме-изображениях, заменив кодировщик пользовательской моделью, специализированной для набора тегов в качестве входных данных, в отличие от естественного языка CLIP. Помимо улучшения кондиционирования, это также позволило бы использовать технику быстрой аугментации от Anifusion, которая так помогла качеству и разнообразию.

Кондиционер

Мы используем тот же подход, что и в Anifusion: модель кондиционирования представляет собой преобразователь размера BERT без позиционных вложений (чтобы уловить симметрию подсказок тегов: порядок тегов не должен иметь значения). Мы просто заменяем кондиционер, сохраняя прежнюю часть перекрестного внимания диффузионной магистрали.

Важным вопросом здесь является то, как инициализировать преобразователь весов. Если начать с необученного кондиционера и обученного UNet и обучить их совместно, это может привести к разрушению или значительному снижению весов UNet в процессе обучения. К счастью, это нетрудно исправить: мы можем заморозить UNet и запустить обучение только для трансформаторной части на некоторое время.

Технические проблемы и оптимизации

Чтобы сделать обучение эффективным, важно втиснуть как можно больший размер пакета в видеопамять (например, потому что стоимость применения логики оптимизатора оплачивается только один раз за пакет). Здесь есть принципиально два направления.

Изменения конвейера

Канонический обучающий код SD запускает кодировщик VAE в режиме онлайн, анализируя необработанное изображение и кодируя его. Это самый простой плод, от которого нужно избавиться: веса VAE потребляют VRAM, промежуточные переменные от применения этого кодировщика также потребляют VRAM, процесс кодирования требует времени. В частности, если мы делаем несколько проходов по данным, нет смысла платить за многократное кодирование одного и того же изображения.

К счастью, эту часть легко оптимизировать: мы просто предварительно кодируем все изображения в тренировочном наборе. Дополнительным преимуществом является то, что результирующий обучающий набор намного меньше (с точки зрения размера диска), чем исходные изображения 512x512 в пространстве пикселей.

Еще один способ сэкономить немного видеопамяти — отключить EMA. Экспоненциальное скользящее среднее — это уловка, позволяющая немного улучшить качество модели путем (грубого) усреднения весов из множества последних контрольных точек, а отказ от хранения EMA — это простой способ использовать меньше видеопамяти, не сильно навредив полученной модели.

Изменения в обучении

Популярный способ ускорить обучение и сэкономить видеопамять — использовать более низкую точность: вместо обычного fp32 (32-битные числа с плавающей запятой) используйте fp16 или bf16 (вариации 16-битных чисел с плавающей запятой) или смешанную точность (выполнение некоторые операции в fp16).

Я попробовал это, и это действительно помогает производительности, но это слишком опасно. В частности, в какой-то момент я добавил «с torch.autocast():» (смешанная точность) для привлечения внимания, а затем потратил довольно много времени на выяснение того, почему модель медленно расходится с течением времени.

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

Обновление тренировочного набора

В Anifusion использовалась более старая версия датасета Danbooru, поэтому я решил, что пора его обновить, заодно расширив размер обучающего набора и набор поддерживаемых тегов в подсказке. Однако возникла пара сложностей: (1) новая версия набора данных не была простым надмножеством предыдущей, и (2) я решил сделать это, пока модель уже обучалась, поэтому потребовалась дополнительная работа по переносу существующего контрольно-пропускного пункта.

Разницы в наборах данных

Это вроде мелочь, но после загрузки метаданных для новой версии датасета я заметил, что куча изображений (пару процентов) пропала. Некоторые теги, которые раньше попадали в топ-2,5к, также исчезли (в том числе некоторые из самых популярных), более того, удаленные изображения соответствовали удаленным тегам.

Дальнейшая проверка показала, что все исчезнувшие теги соответствовали шаблону «странный NSFW» (некоторые из них были настолько странными, что я даже не знал, что такие вещи существуют…). Это было несколько непоследовательно, так как другие похожие по своей природе теги и изображения оставались на своих местах. Я предполагаю, что либо сопровождающий набор данных, либо сайт размещения исходных изображений внедрили его как быстрый хак для решения некоторых проблем с общественностью. Во всяком случае, я просто объединил старые и новые наборы данных, получив около 2 миллионов изображений для обучения. Кроме того, я взял из них 800 000 изображений с разрешением ≥ 786x768 для точной настройки модели с более высоким разрешением позже.

Расширение набора тегов

Этот раздел на самом деле довольно нетривиальный. На момент создания нового набора данных модель довольно долго обучалась, имея хорошо сходящийся преобразователь кондиционера, получая в качестве входных данных подсказки от тегов top-2,5k. Для более новой версии я решил вместо этого взять топ-12 тысяч тегов. Как добавить их в модель?

Что нужно доработать, так это стол встраивания трансформера (с 2500х1024 до 12000х1024). Наивный подход — рандомно инициализировать новые эмбеддинги, и какое-то время обучать модель с замороженным UNet, обновляя только трансформер. Но это (1) расточительно с точки зрения времени и вычислений, (2) скучно. Можем ли мы инициализировать их лучше?

Посмотрим, какая информация нам доступна. Мы уже натренировали 2,5 тыс. эмбеддингов, и эмбеддинги для новых тегов должны быть чем-то похожи на них. Но к каким? Наивная идея — использовать вложения семантически схожих тегов. Но как получить семантическое сходство? Старый добрый word2vec очень хорош в этом, поэтому мы можем просто взять наборы тегов и обучить на них вложения word2vec. Ниже представлена ​​t-SNE визуализация результата (тёмные точки соответствуют тегам NSFW, чтобы их было легко пропустить людям, которые не хотят на них смотреть): интерактивная ссылка.

Давайте посмотрим на ту же визуализацию для 2,5 тыс. обученных эмбеддингов от трансформера: интерактивная ссылка.

Так что это не совсем то же самое, и больше ориентировано на внешний вид, чем на семантику. Однако он должен быть достаточно близким, поэтому мы можем вычислить начальные вложения для новых тегов следующим образом: взять ближайших соседей в пространстве word2vec из первых 2,5k и агрегировать их вложения-преобразователи. На самом деле, мы можем даже немного улучшить его: обучить небольшую NN, которая преобразует word2vec так, чтобы агрегация ближайших соседей совпадала с истинным встраиванием на первых 2,5k, и применить ее к остальным тегам.

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

Обучение с более высоким разрешением

Версия Stable Diffusion 2 генерирует изображения размером 768x768 (т. е. в 2,25 раза больше, чем 512x512), обладающая некоторыми интересными свойствами (например, она лучше справляется с рисованием пальцев, чем версия с более низким разрешением). Поэтому мы решили сделать то же самое с нашей моделью. В частности, у нас есть два этапа обучения:

  • Низкое разрешение: обучение на изображениях 512x512 (скрытое пространство 64x64x4), размер партии 11, скорость обучения 5e-6, 1 миллион шагов;
  • Высокое разрешение: продолжая предыдущий этап, тренируйтесь на изображениях 768 x 768 (скрытое пространство 96 x 96 x 4), с размером пакета 5, скоростью обучения 5e-6, 80 тыс. шагов.

FWIW, похоже, что обе ступени недостаточно приспособлены, поэтому возможны дальнейшие улучшения, если тренироваться дольше.

Выборочные изменения

В Anifusion добавление подсказки отдельной моделью резко улучшило результаты. Интуитивно легче отбирать материал в дискретном концептуальном пространстве, а не в неявном непрерывном пространстве внутри процесса распространения.

Здесь мы применяем ту же технику, переобучая преобразователь быстрого увеличения на новом обучающем наборе. Мы наблюдаем тот же эффект — качество выборки и разнообразие резко увеличиваются при увеличении подсказок, особенно для коротких подсказок.

Полученные результаты

Получается, что модель 768х768 дает лучшее изображение, чем модель 512х512, и не только из-за разрешения: лучше общая согласованность изображений и (особенно) анатомические детали персонажей. Однако похоже, что достоверность подсказок несколько хуже (правда, у меня нет точных замеров, только мое впечатление).

Следует отметить, что по сравнению с Anifusion модель кажется более предвзятой к изображениям NSFW / racing. Частично это можно отнести к быстрому расширению, но часть UNet также несет ответственность. Я не уверен, почему это происходит (изменения тренировочного набора были довольно небольшими; может быть, это вызвано размером/началом со стабильной диффузии?).

Давайте рассмотрим различные типы изображений шаг за шагом.

Портреты

Это самый простой вариант использования, и аниме-портреты успешно генерировались даже в эпоху GAN. Так что неудивительно, что у модели нет проблем с этим (несколько примеров ниже).

Один символ, популярные теги

Это более сложный тест, и, как и ожидалось, время от времени (редко) можно получать плохие изображения. Распространенный способ отказа — плохая анатомия — дополнительная рука или нога. И, конечно же, пальцы довольно сложны для модели.

Один символ, хвостовые теги

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

Взаимодействия персонажей

Наконец-то мы добрались до чего-то интересного. Оригинальный Anifusion не преуспел, когда персонажи взаимодействовали. Новая модель лучше?

Несколько символов, просто стоящих рядом друг с другом, кажутся нормальными:

Простые действия тоже немного работают:

Взаимодействия с объектами тоже более-менее нормальные:

Однако, если мы попытаемся конкретизировать сложные взаимодействия, это чаще всего вызывает ужас тела (я не буду размещать эти фотографии здесь, поскольку они могут вызывать беспокойство). Почему это происходит? Я предполагаю, что это связано с тем, что подсказки представляют собой просто наборы тегов и недостаточно точно определяют взаимодействия. Что заставило диффузионные модели создавать хорошие изображения, в первую очередь, так это руководство без классификатора: чрезмерное выделение подсказок резко улучшило качество по сравнению с выборкой без подсказок, потому что выборка сложных концепций внутри диффузии затруднена. Таким образом, мы можем ожидать, что все, что недвусмысленно указано в подсказке, будет сгенерировано хорошо, но неоднозначности будут разрешены не очень хорошим способом. В этом смысле подсказки на естественном языке могут быть лучше.

FWIW, в комментариях к исходному сообщению Anifusion на Reddit, Gwern предложил использовать BLIP для создания подсказок на естественном языке из изображений; Я пробовал, но качество подсказок было крайне низким, так что это не так просто.

Стили

В оригинальной версии Stable Diffusion люди нашли множество способов указать стили модели, в основном добавляя имена некоторых художников. То же самое работает для аниме-модели?

Пробуем — работает очень хорошо (это преимущество по сравнению с AnythingV3). Ниже приведены примеры в различных случайных стилях.

Сравнение с другими аниме-моделями

Мы не проводим сравнение с AnythingV3, так как он на самом деле не учитывает приглашение, генерируя похожие изображения практически для всего (как следует из названия). Остаются WaifuDiffusion и NovelAI. Мы отказались от сравнения с NovelAI, так как это потребует использования просочившейся контрольной точки, и лучше держаться подальше от этой драмы. Это оставляет нас с WaifuDiffusion.

Следует отметить, что сравнение коротких подсказок было бы слишком несправедливым по отношению к WaifuDiffusion: он не включает увеличение подсказок (в отличие от Anifusion), что дает очень плохие результаты на коротких подсказках. Поэтому вместо этого мы делаем два разных сравнения. В одном из них мы отбираем наборы тегов из исходного набора данных, эффективно исключая быстрое увеличение из рассмотрения и сравнивая только диффузную часть. Во втором мы вручную строим подсказки среднего размера и оцениваем на них обе модели (включая эффект увеличения подсказок).

Наборы тегов из набора данных

Мы выбираем 97 наборов тегов из набора данных (равномерно, фильтруя изображения с безопасным рейтингом). Для Anifusion мы предварительно обрабатываем их, добавляя в приглашение rating_s, score_perc_100, Adjust_score_perc_100. Для WaifuDiffusion мы предварительно обрабатываем их, заменяя символы подчеркивания пробелами. В обоих случаях мы используем выборку DDIM с 50 шагами и шкалой наведения без классификатора 9.

После получения некоторых человеческих оценок (со случайным переворачиванием сторон для сравнения) мы получаем следующие результаты по сравнению с вайфу:

  • Anifusion лучше: 30%
  • Вайфу лучше: 24%
  • Примерно столько же: 46%

Полные данные можно найти здесь, ниже несколько примеров (слева — Anifusion, справа — Waifu):

Один паттерн сбоя, который я заметил в Waifu: он часто генерирует часть изображения вместо полного изображения, обрезая голову персонажа. Паттерн неудачи Anifusion — это непроизвольный NSFW (и связанный с этим — плохая анатомия; по какой-то причине, чем ближе изображение к NSFW, тем хуже модель получает анатомию).

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

Так что с выборкой в ​​распределении ситуация не намного лучше. Но что, если мы используем версию модели с высоким разрешением?

  • Anifusion лучше: 41%
  • Вайфу лучше: 13%
  • Примерно столько же: 46%

Хорошо, это имеет большое значение! (сырые данные) Соотношение выигрышей к проигрышам больше 3, что очень много. Правда, на это могли повлиять какие-то личные предпочтения оценщиков (например, было много потерь вайфу из-за обрезанных изображений), но все равно это довольно убедительно.

Можно заметить нетривиальное количество «примерно одинаково». В большинстве этих примеров обе модели одинаково хороши, но есть и такие, где обе модели ужасно плохи (в этом они каким-то образом коррелируют).

Подсказки среднего размера вручную

Мы используем ту же предварительную обработку подсказок, что и в предыдущем разделе. Мы вручную готовим 33 разумных подсказки среднего размера (конечно, поскольку они подготовлены одним человеком, они могут не отражать общее распределение).

Здесь результаты совсем другие:

  • Anifusion лучше: 73%
  • Вайфу лучше: 9%
  • Примерно столько же: 18%

Вот — это сравнение. Ниже приведены несколько примеров (слева — Anifusion).

Как мы видим, быстрое увеличение имеет значение. Что происходит, когда мы используем модель высокого разрешения? Результаты очень похожи (необработанные данные):

  • Anifusion лучше: 73%
  • Вайфу лучше: 3%
  • Примерно столько же: 24%

Обсуждение

Доступность для непрофессионалов

Одно интересное наблюдение заключается в том, что даже после миллиона обучающих мини-пакетов модель продолжает улучшаться. Неясно, помнит ли модель что-нибудь из исходной контрольной точки Stable Diffusion после стольких шагов. Возможно, набор аниме-данных достаточно прост, чтобы можно было обучать диффузионные модели SD-размера с нуля на стандартном оборудовании без горячего старта с существующего SD? Это направление, которое следует исследовать: если это так, тестирование новых архитектур для отдельных лиц должно стать проще без необходимости тратить миллионы долларов на вычисления, как в случае с обучением на основе LAION.

Дальнейшие улучшения качества

Интересно посмотреть, как далеко можно продвинуться с текущей моделью, поэтому я хочу еще немного ее настроить, возможно, используя полный набор данных, а не подмножество 2M. Но я подозреваю, что ограничивающим фактором является качество набора данных, а не объем тонкой настройки (т. е. для улучшения качества модели нужно либо скармливать только более качественные изображения, либо добавлять в кондиционер специальный идентификатор, который помогает модели различать лучшие и худшие изображения). Насколько мне известно, для аниме нет надежных моделей, подобных этой (для обычных изображений люди используют «эстетическое» подмножество LAION). Хотя есть некоторые модели, которые дают отличные результаты, которые выглядят иначе, чем эта модель или WaifuDiffusion (например, аниме-модель Миджорни, которая доступна только в виде выбранных ими изображений), так что это может быть возможно.

Важно отметить, что генерация с более высоким разрешением улучшает качество в аспектах, отличных от разрешения, например. печально известные ИИ-руки становятся лучше. Интересно, даст ли дальнейшее увеличение до 1024x1024 больше улучшений? Можно попробовать использовать каскадную диффузию (например, Imagen) в скрытом пространстве (например, 64x64 -> 128x128 -> 256x256 в скрытом пространстве, что дает 2048x2048 в исходном пространстве).

Удобство использования

Одна проблема с тем, что я выпускаю, заключается в том, что это доказательство концепции. Код для модели отличается от SD, поэтому он не совсем совместим с большинством инструментов (особенно с частью токенизатора подсказок). Таким образом, несмотря на то, что эта модель выглядит лучше, чем другие доступные аниме-модели, простота использования не так велика. Потенциальным решением проблемы может быть правильная интеграция сообщества (или меня, если у меня есть на это время) с существующими инструментами или обеспечение совместимости с SD на уровне интерфейса.

Рекомендации

указатели