Введение

На сегодняшний день Stable Diffusion является источником вдохновения для сообщества любителей опенсорсного машинного обучения и в то же время источником расстройства для художников всего мира. Мне стало любопытно, что ещё может сделать эта важная технология кроме того, как подвергать угрозе рабочие места профессиональных художников и дизайнеров.

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




Сан-Франциско, сверху вниз: JPG (6,16 Кбайт), WebP (6,80 Кбайт), Stable Diffusion: (4,96 Кбайт)




Кондитерская, сверху вниз: JPG (5,68 Кбайт), WebP (5,71 Кбайт), Stable Diffusion (4,98 Кбайт)




Анна, сверху вниз: JPG (5,66 Кбайт), WebP (6,74 Кбайт), Stable Diffusion (4,97 Кбайт)

Из этих примеров достаточно очевидно, что сжатие изображений при помощи Stable Diffusion обеспечивает гораздо большее качество изображения при меньшем размере файла по сравнению с JPG и WebP. Это качество несёт с собой и некоторые важные особенности, которые нужно учитывать; подробнее я расскажу об этом ниже. Однако на первый взгляд это очень многообещающий вариант для агрессивного сжатия изображений с потерями.

Скрытое пространство Stable Diffusion

Stable Diffusion использует в тандеме три обученных искусственных нейронных сети:

• Variational Auto Encoder
• U-Net
• Text Encoder

Variational Auto Encoder (VAE) кодирует и декодирует изображения из пространства изображений в некое описание в скрытом пространстве. Описание в скрытом пространстве — это описание любого исходного изображения (512 x 512 с глубиной 3x8 или 4x8 бит) с уменьшенным разрешением (64 x 64) и бОльшей точностью (4x32 бита).

VAE кодирует изображение в это скрытое пространство следующим образом: она обучается в процессе тренировки, то есть описания скрытого пространства разных версий модели с большой вероятностью при дальнейшем обучении будут выглядеть по-разному. В Stable Diffusion v1.4 описание исходного изображения из начала статьи выглядит так (после преобразования и интерпретации в качестве изображения с четырёхканальным цветом):


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

Один цикл кодирования/декодирования при помощи этой VAE выглядит так:




Цикл VAE, сверху вниз: исходное изображение 512x512@24bpp, описание скрытого пространства 64x64@128bpp, декодированное изображение 512x512@24bpp

Обратите внимание, что этот цикл работает с потерями. Например, имя Анны на её наморднике после декодирования менее читаемо. VAE версии Stable Diffusion 1.4 в целом не очень хорошо справляется с описанием мелкого текста (а также лиц; надеюсь, в версии 1.5 обученной модели ситуация улучшится).

Основной алгоритм Stable Diffusion, генерирующий изображения из коротких текстовых описаний, работает с этим описанием изображений в скрытом пространстве. Он начинает со случайного шума в описании скрытого пространства, а затем итеративно избавляет это изображение в скрытом пространстве от шума при помощи обученной U-Net, которая как бы на выходе создаёт прогнозы того, что она «видит» в этом шуме, аналогично тому, как глядя на облака, мы иногда видим фигуры и лица. Когда Stable Diffusion используется для генерации изображений, этот итеративный этап избавления от шума управляется третьей моделью ML — кодировщиком текста, дающим U-Net информацию о том, что она должна попытаться увидеть в шуме. Для описываемого в моей статье экспериментального кодека изображений кодировщик текста не нужен. В показанном ниже коде Google Colab он всё равно используется, но только для создания единовременной кодировки пустой строки, сообщающей U-Net, что при воссоздании изображения нужно избавляться от шума без управления.

Методика сжатия

Чтобы использовать Stable Diffusion в качестве кодека сжатия изображений, я изучил возможности эффективного сжатия сгенерированного VAE описания в скрытом пространстве. В моих экспериментах даунсэмплинг данных скрытого пространства или применение к этим данным имеющихся способов сжатия с потерями существенно снижали качество воссоздаваемых изображений. Однако я выяснил, что декодирование VAE, похоже, чрезвычайно устойчиво к квантизации данных скрытого пространства.

Квантизация данных из значений с плавающей запятой в 8-битные беззнаковые integer при помощи масштабирования, усечения и преобразования приводит лишь к очень малозаметной погрешности воссоздания:




Сверху: декодировано из 32-битных данных с плавающей запятой скрытого пространства. Посередине: исходное изображение. Снизу: декодировано из 8-битных данных integer скрытого пространства

Для квантизации сгенерированных VAE изображений я сначала отмасштабировал их в соотношении 1 / 0,18215 — это число вы можете найти в исходном коде Stable Diffusion. Похоже, деление данных скрытого пространства на это число достаточно хорошо преобразует их в диапазон [-1, 1], однако частичное ограничение всё равно применяется.

После квантизации данных скрытого пространства до 8 бит описание изображения имеет размер 64*64*4*8 бит = 16 Кбайт (исходное изображение без сжатия имеет размер 512*512*3*8 бит = 768 Кбайт).

В моих экспериментах квантизация до менее чем 8 бит не давала хороших результатов, однако на удивление хорошо подошла дальнейшая квантизация при помощи палитр и дизеринга. Я создал палитризированное описание при помощи палитры скрытого пространства из 256 векторов 4*8 бит и дизеринга по методу Флойда-Стейнберга. Использование палитры из 256 элементов позволяет описать каждый вектор скрытого пространства при помощи одного 8-битного индекса, что снижает размер данных до 64*64*8+256*4*8 бит = 5 Кбайт

Однако палитризированное описание приводит к заметным артефактам при декодировании их непосредственно с помощью VAE:




Сверху: 32-битные данные скрытого пространства. Посередине: 8-битные квантизованные данные. Внизу: палитризированные 8-битные данные с дизерингом по методу Флойда-Стейнберга (обратите внимание на заметные искажения)

Дизеринг палитризированных данных добавил шума, что исказило декодированный результат. Но поскольку работа Stable Diffusion основана на устранении шума из данных скрытого пространства, мы можем использовать U-Net для устранения шума, добавленного дизерингом. Спустя всего четыре итерации результат воссоздания визуально очень близок к неквантизованной версии:




Сверху: после палитизирования и дизеринга данных скрытого пространства. Посередине: после четырёх этапов устранения шума. Внизу: исходное изображение

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

Далее я сжал без потери палитру и индексы при помощи zlib; в результате размер большинства образцов, на которых я выполнял тестирование, оказалось чуть меньше 5 Кбайт. Я подумывал об использовании кодирования длин серий (run-length encoding, RLE), однако из-за дизеринга в большинстве изображений оставались лишь очень короткие серии с одинаковым индексом (даже в таких изображениях, которые отлично подходят для RLE при сжатии в пространстве изображения, например, в графике без большого количества градиентов). Я практически уверен, что в этой области есть большой потенциал для дальнейших оптимизаций.

Оценка

Для оценки этого экспериментального кодека сжатия я не использовал стандартных тестовых изображений или изображений, найденных онлайн, чтобы точно не тестировать его на данных, которые могли бы использоваться в обучающем массиве данных модели Stable Diffusion (потому что такие изображения могут получить нечестное преимущество при сжатии, поскольку часть их данных уже закодирована в обученной модели). Кроме того, чтобы сделать сравнение как можно более справедливым, я использовал самые высокие параметры качества кодировщика алгоритмов сжатия JPG и WebP в библиотеке Image языка Python, а также дополнительно применил к сжатым данным JPG сжатие без потерь при помощи библиотеки mozjpeg. Затем я использовал силу сжатия на единицу больше той, которая бы привела к созданию результатов с размером данных меньше, чем у результатов Stable Diffusion, или в противном случае максимальную силу (многие изображения не становятся меньше, чем результат SD даже при сжатии JPG или WebP с максимальными потерями).

Важно заметить, что, хотя результаты Stable Diffusion субъективно выглядят гораздо лучше изображений, сжатых JPG и WebP, они ненамного лучше (но и не хуже) с точки зрения стандартных метрик наподобие PSNR или SSIM. Добавляемые SD артефакты намного менее заметны, потому что они больше влияют на содержимое изображения, а не на его качество. В этом и заключается небольшая опасность данной методики: не нужно обманываться качеством воссозданных результатов, так как на содержимое могли повлиять артефакты сжатия, даже если оно выглядит очень чётко. Например, посмотрим на деталь тестового изображения с видом Сан-Франциско:




Сверху: сжатие JPG. Посередине: исходное изображение. Снизу: сжатие Stable Diffusion

Как видите, хотя Stable Diffusion в качестве кодека гораздо лучше сохраняет качественные аспекты изображения, вплоть до зернистости камеры (с которой испытывают трудности большинство традиционных алгоритмов сжатия), на содержимое всё равно влияют артефакты сжатия, а такие мелкие детали, как форма зданий, могут меняться. Хотя в сжатом JPG совершенно точно нельзя распознать больше исходного изображения по сравнению со сжатием Stable Diffusion, высокое визуальное качество результата SD может быть обманчивым, так как артефакты сжатия в JPG и WebP гораздо проще выявить как артефакты.

Ограничения

На данный момент некоторые признаки не очень хорошо сохраняются VAE модели Stable Diffusion; в частности, возникают проблемы с мелким текстом и лицами (это согласуется с «Ограничениями и искажениями» версии 1.4 модели SD, перечисленными в карточке модели). Вот три примера, иллюстрирующих это:




Сверху: исходное изображение. Посередине: непосредственный цикл VAE (32-битные данные скрытого пространства). Снизу: декодировано из палитризованных и избавленных от шума 8-битных данных скрытого пространства




Сверху: исходное изображение. Посередине: непосредственный цикл VAE (32-битные данные скрытого пространства). Снизу: декодировано из палитризованных и избавленных от шума 8-битных данных скрытого пространства




Сверху: исходное изображение. Посередине: непосредственный цикл VAE (32-битные данные скрытого пространства). Снизу: декодировано из палитризованных и избавленных от шума 8-битных данных скрытого пространства

Как видно по изображениям посередине, актуальная VAE модели Stable Diffusion 1.4 не очень хорошо сохраняет в скрытом пространстве мелкий текст и лица, поэтому данная схема сжатия наследует эти ограничения. Но поскольку версия 1.5 Stable Diffusion, пока доступная только в Dreamstudio компании stability.ai, гораздо лучше справляется с лицами, я хочу дополнить эту статью после публикации модели 1.5 для общего доступа.

Чем эта методика отличается от восстановления изображений с помощью ML?

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

Важно понимать, что Stable Diffusion не используется для восстановления изображения, которое было испорчено сжатием в пространстве изображения, а применяет сжатие с потерями к внутреннему пониманию моделью Stable Diffusion этого изображения и пытается «починить» повреждения, которые причинило внутреннему описанию сжатие с потерями (что отличается от восстановления испорченного изображения).

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

Аналогия этого различия: допустим, у нас есть художник, обладающий фотографической памятью. Мы показываем ему изображение, а затем просим воссоздать его, и он создаёт почти идеальную копию по памяти. Фотографическая память этого художника и есть VAE модели Stable Diffusion.

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

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

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

Затем мы просим его снова нарисовать изображение, но на этот раз, если он помнит некоторые аспекты очень странными и не имеющими смысла, он должен воспользоваться своим опытом и нарисовать эти объекты так, чтобы они были логичными (стоит заметить, что «не имеющий смысла» не всегда означает «плохой»: размытое или поцарапанное фото является плохим, но эти дефекты вполне логичны).

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

Благодаря этой аналогии очевидно, что данная схема сжатия ограничена тем, насколько хорошо работает фотографическая память художника. В случае Stable Diffusion v1.4 он не очень хорошо справляется с запоминанием лиц, а также страдает от дислексии.

Код

Если вы хотите попробовать этот экспериментальный кодек самостоятельно, то это можно сделать при помощи Google Colab:

https://colab.research.google.com/drive/1Ci1VYHuFJK5eOX9TB0Mq4NsqkeDrMaaH?usp=sharing

Заключение

Stable Diffusion кажется очень перспективной в качестве основы для схемы сжатия изображений с потерями.

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

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

Приложение: дополнительные результаты

Сверху: JPG. Посередине: WebP. Снизу:: Stable Diffusion. Во всех случаях изображение SD сжимается в меньший размер, чем соответствующие JPG и WebP. Результаты специально не выбирались.