Stable diffusion VAE fine tuning (backport AutoencoderKL and its config.yaml to taming-transformers)

[rbbb]

Can we have stable diffusion VAE fine-tuning directly from taming-transformers ?

The code seems to work (obviously, AutoencoderKL was taken from taming-transformers).
Both the AutoencoderKL code and the config snippet were taken from stable-diffusion.
Usage is strictly identical to 'VQGAN with your own data'.

There is some safetensors loading code, but it doesn't work with torch 1.7 that is recommended with taming-transformers.

Related discussions:
lllyasviel/ControlNet#500

Related code:
https://github.com/cccntu/fine-tune-models/blob/main/run_finetune_vae.py
(adapted from Patil Suraj's stable-diffusion-jax)

[rbbb]

Added a colab notebook in commit efb20eb.

I'm slightly confused about the actual objective function.

• some terms are sometimes dropped from the objective [Community] Training AutoencoderKL huggingface/diffusers#894
• this gives additional objective functions used https://huggingface.co/stabilityai/sd-vae-ft-mse-original , as well as training setup
• my best guess is that it boils down to aesthetic choices instead of correct math (clip skip tweaks, hand choosing Lora checkpoints, etc ...)

[sgw-ite]

hi, i use the code in https://github.com/CompVis/taming-transformers/pull/222/files, I would like to ask why you used VQLPIPS as the loss function in line 20(configs/finetune_vae.yaml), and also thank you very much for your code!

[rbbb]

Hi.

In the original pull request, I wrote 'it is an aesthetic choice'

There are no rules in what metrics are used to fine-tune the VAE (the VAE police will not come to get you if you change the loss function). It is usual to drop the discriminator in fine tuning, and by usual, I mean 'common practice that people usually do without formal verification or peer reviewed paper'. Dropping or including LPIPS is the same, the result will give you an aesthetically different result.

To take a classical example: https://en.wikipedia.org/wiki/Dither
After reducing color space, having only pixel loss will produce bands in the image
Having perceptive loss should achieve dithering, with the occasional bad pixel.

So you should run the two (with and without LPIPS), look very closely at you images, and see which one you prefer.

More formally, if you take a paper at random (say, the StableDiff3 paper https://arxiv.org/pdf/2403.03206.pdf ), you'll notice that model evaluation is all human preference.

HTH