[Community] Training AutoencoderKL

[AI-Guru]

Hi!

I am working on latent diffusion for audio and music. It seems to me that Diffusers 🧨 is the place to be! There is a feature I would like to request: Training AutoencoderKL (Variational Autoencoder).

What I would love to do, is training AutoencoderKL on square and non-square images, either with one or more than one channels. I checked the implementation, and it seems to me that due to its fully convolutional nature, this would be perfectly possible.

A good start would be a script/notebook that shows how to train AutoencoderKL on a Hugging Face dataset. On the long run it could be even a Trainer.

[Yogiraj587]

@AI-Guru Can you please assign this issue.....????

[patrickvonplaten]

Hey @AI-Guru,

Cool idea! I'm not sure how much time our training experts will have for this (cc'ing @patil-suraj and @anton-l) but this would indeed be a very useful addition.

Also opening this one up to the Community in case anybody is interested :-)

[patil-suraj]

Nice idea,

But note that, training the AutoencoderKL is bit complicated and outside the scope of diffusers. It's complicated because the VAE in SD is trained with GAN objective and has multiple losses like lpips, discriminator loss which requires implementing extra modules. For training VAE I think the best resource is the taming-transformers repo, it's the repo which is uses to train the VAE in SD has all the components required for training implemented.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[jbmaxwell]

Did you get this running @AI-Guru? I'd be very interested if you did... and if so, any script you can share? :)

[eeyrw]

Nice idea,

But note that, training the AutoencoderKL is bit complicated and outside the scope of diffusers. It's complicated because the VAE in SD is trained with GAN objective and has multiple losses like lpips, discriminator loss which requires implementing extra modules. For training VAE I think the best resource is the taming-transformers repo, it's the repo which is uses to train the VAE in SD has all the components required for training implemented.

According https://huggingface.co/stabilityai/sd-vae-ft-mse, it seems to finetune VAE used in Stable Diffusion, only LPIPS and MSE loss are required and no discriminator loss required. Is that true?

[patil-suraj]

Yeah for fine-tuning think it didn't use the discriminator, but the vae is pre-trained using all those 3 losses, as you can see in taming-transformers

[jbmaxwell]

Just wondering; is there a way to load AutoencoderKL from a config.json, but without the pretrained weights? Or a way to initialize it? I've been trying to fine-tune it to my specific data, but it's struggling to improve beyond a certain level so I'd like to give it a shot "from scratch"... I can't find any way to do that (no AutoConfig -> from_config() for Diffusers).

Any help appreciated.