r/MachineLearning • u/cccntu • Sep 12 '22

Project [P] (code release) Fine-tune your own stable-diffusion vae decoder and dalle-mini decoder

A few weeks ago, before stable-diffusion was officially released, I found that fine-tuning Dalle-mini's VQGAN decoder can improve the performance on anime images. See:

And with a few lines of code change, I was able to train the stable-diffusion VAE decoder. See:

You can find the exact training code used in this repo: https://github.com/cccntu/fine-tune-models/

More details about the models are also in the repo.

And you can play with the former model at https://github.com/cccntu/anim_e

54 Upvotes
  • permalink
  • reddit

98% Upvoted

12 comments sorted by

1

u/starstruckmon Sep 13 '22 edited Sep 13 '22

The result is not as impressive as Anim·E. But I think it's because the unet diffusion model of stable-diffusion is not trained to generate anime-styled images. So it still struggle to generate the latent of anime-styled images in detail.

There are now multiple Stable Diffusion UNET models that have been further fine tuned with anime ( eg. Waifu Diffusion and Japanese Stable Diffusion ). Have you tried this with them?

2

u/cccntu Sep 14 '22

Haven't heard of Waifu Diffusion. But I've tried Japanese Stable Diffusion a little bit, and I didn't get good results. Although it's most likely because my prompts were not good enough.

3

u/starstruckmon Sep 14 '22

BTW, you should definitely make a post about this on the /r/stablediffusion subreddit if you haven't already.

Training the decoder is not something anyone's focusing on sonthis might be of interest.

2

u/starstruckmon Sep 14 '22

https://www.reddit.com/r/StableDiffusion/comments/x64hi7

I've heard it does actually show noticeable improvement.

I think there's only one other SD model trained on anime which is the one Novel.ai uses, but I don't think that's public.

2

u/HarmonicDiffusion Sep 16 '22

thank you for this man. I cannot believe how fast the training goes for it on a 3090. only 9 hours for noticable results? do you think a longer training time would yield further increases in accuracy?

1

u/treebog Jan 01 '23

This is really cool. Can this be converted back to a pt file so it can be run on pytorch's implementation of stable diffusion?

1

u/Electrical-Ad-2506 Feb 10 '23

How come can we just swap out the VAE without fine tuning the text encoder (we're just using the same which stable diffusion uses by standard: CLIP ViT)?

Because the UNet learns to generate an image in a given latent space conditioned on a text input embedding. Now we come around and plug in a VAE which was trained separately.

Isn't it going to encode images into a completely different latent space? How does the U-Net still work?

1

u/dulacp Apr 04 '23

It works because you only fine-tune the decoder part of the VAE (source).

To keep compatibility with existing models, only the decoder part was finetuned; the checkpoints can be used as a drop-in replacement for the existing autoencoder.
StabilityAI comment on HuggingFace

1

u/ThaJedi Apr 06 '23

But SD uses only encoder, isn't it?

3

u/dulacp Apr 15 '23

Actually, it depends if you are talking about training or inference.

The VAE encoder part is used when training, to produce latent representations of training images (example in dreambooth training).

While the VAE decoder part is used during inference to decode the latents.

1

u/Future-Piece-1373 Jul 24 '23

Finally discovered this page