Skip to content

brian6091/Dreambooth

 
 

Repository files navigation

Dreambooth-style fine tuning of Stable Diffusion models

NOTE: This branch exists for the sole purpose of not breaking old links. Go to the main branch, which has already merged in these updates, and probably does cooler stuff by now.

For classic Dreambooth Train In Colab

For Low-rank Adaptation (LoRA) Train In Colab

Some notebooks for fine-tuning Stable Diffusion models.

Tested with Tesla T4 and A100 GPUs on Google Colab (some settings will not work on T4 due to limited memory)

Tested with Stable Diffusion v1-5 and Stable Diffusion v2-base.

There are lots of notebooks for Dreambooth-style training. This one borrows elements from ShivamShrirao's implementation, but is distinguished by some additional features:

  • based on Hugging Face Diffusers🧨 implementation so it's easy to stay up-to-date
  • Low-rank Adaptation (LoRA) for fast text-to-image fine-tuning (using cloneofsimo's implementation)
  • exposes lesser-explored parameters for experimentation (ADAM optimizer parameters, cosine_with_restarts learning rate scheduler, etc), all of which are dumped to a json file so you can remember what you did
  • possibility to drop some text-conditioning to improve classifier-free guidance sampling (e.g., how SD V1-5 was fine-tuned)
  • training loss and prior class loss are tracked separately (can be visualized using tensorboard)
  • option to generate exponentially-weighted moving average (EMA) weights for the unet
  • easily switch in different variational autoencoders (VAE) or text encoders
  • inference with trained models is done using Diffusers🧨 pipelines, does not rely on any web-apps

Buy Me A Coffee