Что такое VAE для Stable Diffusion? Обязательно добавляем!

Stable Diffusion принадлежит к классу диффузионных моделей. Идея состояла в том, чтобы смешать картинку и гауссовский шум. А далее обучить нейросеть из зашумленных изображений восстанавливать оригиналы. Если подать такой нейросети на вход чистый шум, то постепенно она превратит его в симпатичную картинку. В Stable Diffusion этим занимается UNet.
Нюанс в том, что картинка с разрешением 512х512 состоит из 262 144 пикселей. Если применять диффузионный процесс к ней напрямую, то на вычисления будет уходить много времени и памяти, что усложняет процесс тренировки и инференса. Мы же хотим генерировать картинки быстро и на относительно недорогих видеокартах. Поэтому изображения отображаются в пространство меньшей размерности (latent space), там производятся вычисления, а результат обратно разжимается с помощью Variational Autoencoder (VAE).

Если вы используете стандартную модель от SD, то при процессе раскодировки SD использует дефолтный VAE, который входит в сборку автоматика или другие сборки, которые вы используете. Если же вы скачали какую-то другую модель, то тут начинаются нюансы. Стандартная VAE от SD хорошо работает в целом, но у нее проблемы с глазами, носом, губами и пальцами. Стандартную VAE можно использовать, но больше подойдет для пейзажей, архитектуры и паинтов. Авторы новых моделей, которые можно найти на сайтах каталогах, иногда, добавляют VAE заточенные под их модели. Они тренируют модель VAE под себя, под свою модель, тем самым улучшают те детали, на которые была заточена модель, это может быть космический сет, аниме, аналоговое фото и т.д. Например, есть VAE заточенные под аниме, они натренированы под стиль аниме. С таким VAE отлично получается генерация анимешных артов, но с помощью аниме VAE вы уже не разожмете, например, реалистичное фото человека или животного.
Вместо стандартного SD VAE энтузиасты тренируют новые модели, одна из таких sd-vae-ft-ema-original натренированная на базе LAION-Aesthetics и LAION-Humans для SWF контента. Она дает более улучшенный результат при генерации человеческих тел и считается более универсальной. Есть еще варианты для NSWF, но их ищите отдельно.

Поэтому советуют скачать универсальную VAE, которая будет разжимать генерацию в приемлемом виде на всех моделях. В основном рекомендуют VAE sd-vae-ft-ema-original — скачиваем по ссылке https://huggingface.co/stabilityai/sd-vae-ft-ema-original/blob/main/vae-ft-ema-560000-ema-pruned.safetensors и закидываем файл в папку ВашSD/models/VAE
Далее перезапускаемся, во вкладке Setting выбираем настройки Stable Diffusion и выбираем vae-ft-ema-560000-ema-pruned.safetensors — жмем Apply settings. Готово.

Что такое VAE для Stable Diffusion? Обязательно добавляем!

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

Пояснение на английском языке.
A VAE is a variational autoencoder.
An autoencoder is a model (or part of a model) that is trained to produce its input as output. By giving the model less information to represent the data than the input contains, it’s forced to learn about the input distribution and compress the information. A stereotypical autoencoder has an hourglass shape — let’s say it starts with 100 inputs and reduces it to 50 then 20 then 10 (encoder) and then 10 to 20 to 50 to 100 (decoder). The 10 dimensions that the encoder produces and the decoder consumes are called the latent representation.
Autoencoders can be a powerful paradigm and can be trained in an unsupervised way (without needing to label data since we only need the input data). However, if we want to sample from the input distribution, a vanilla autoencoder makes this difficult or impossible. One variation on the autoencoder is the variational autoencoder where the latent is normally distributed, which allows for the output distribution to be sampled from.
SD is somewhat unique in the vision class of diffusion models in that the diffusion process operates in the autoencoder space instead of pixel space. This makes the diffusion process more computationally efficient / memory efficient compared to a vanilla pixel space diffusion model. One other related technique some models use is to start the diffusion at a lower spatial resolution and progressively upscale to save compute.
In practice, in SD, the VAE is pretty aggressive and the dataset is filtered (indirectly through the aesthetic score) which removes images with a lot of text. This combined with the autoencoder is a significant reason SD struggles more with producing text than models like Dall-e.
From the above, an autoencoder is essential in SD. Generally speaking, there’s no reason to modify the autoencoder unless the image distribution you’re training on is dramatically different than the natural images given to SD. In this case, you’d likely need to retrain all parts of the model (or at least the unet). One example case where this might be useful is if you wanted to train an audio diffuser using the same components as SD but on «pixel» data from a spectrogram.
Image generation is done in a «compressed» representation, the VAE takes the compressed results and turn them into full sozed images. SD comes with a VAE already but certain models may supply a custom VAE that works better for that model.