TL;DR

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

• вступление
• Предыстория
  - Модели распространения
  - Скрытое распространение
  - Стабильное распространение - Распространение Waifu
  - Стабильное распространение 2
  
• Наша модель: обучение
  - Кондиционер
  - Технические вопросы и оптимизация
  - Обновление тренировочного набора
  - Обучение с более высоким разрешением
  - Изменения выборки
• Результаты
• Сравнение с другими аниме-моделями
• Обсуждение
• использованная литература
• указатели

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

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

Однако, учитывая обилие моделей в этой сфере, поначалу я решил не заморачиваться и позволить прогрессу идти своим чередом. Но сейчас вышла 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 тыс. шагов.

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

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

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

Результаты

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

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

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

Портреты

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

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

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

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

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

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

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

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

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

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

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

FWIW, в комментариях к исходному сообщению Anifusion на Reddit, Гверн предложил использовать 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 без классификатора.

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

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

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

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

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

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

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

Ручные подсказки среднего размера

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

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

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

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

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

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

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

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

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

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

Юзабилити

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

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

• Стабильная диффузия 2
• Скрытая диффузия
• Набор данных Danbooru2021

• Инструкции по вызову веб-интерфейса
• Контрольные точки модели (huggingface): 512x512 , 768x768 , подсказка
• Репозиторий SD Github адаптирован для модели
• Демо-версия 768x768 — ограничение «рейтинг: безопасный», я оставлю его в рабочем состоянии на какое-то время, но в конечном итоге не буду его поддерживать