diffusers/examples/dreambooth at main · huggingface/diffusers · GitHub
DreamBooth training example
DreamBooth is a method to personalize text2image models like stable diffusion given just a few(3~5) images of a subject. The
train_dreambooth.py script shows how to implement the training procedure and adapt it for stable diffusion.Running locally with PyTorch
Installing the dependencies
Before running the scripts, make sure to install the library's training dependencies:
Important
To make sure you can successfully run the latest versions of the example scripts, we highly recommend installing from source and keeping the install up to date as we update the example scripts frequently and install some example-specific requirements. To do this, execute the following steps in a new virtual environment:
Plain Text
Then cd in the example folder and run
Plain Text
And initialize an
environment with:
Plain Text
Or for a default accelerate configuration without answering questions about your environment
Plain Text
Or if your environment doesn't support an interactive shell e.g. a notebook
Plain Text
Dog toy example
Now let's get our dataset. Download images from here and save them in a directory. This will be our training data.
And launch the training using
Note: Change the
resolution to 768 if you are using the stable-diffusion-2 768x768 model.Plain Text
Training with prior-preservation loss
Prior-preservation is used to avoid overfitting and language-drift. Refer to the paper to learn more about it. For prior-preservation we first generate images using the model with a class prompt and then use those during training along with our data. According to the paper, it's recommended to generate
num_epochs * num_samples images for prior-preservation. 200-300 works well for most cases. The num_class_images flag sets the number of images to generate with the class prompt. You can place existing images in class_data_dir, and the training script will generate any additional images so that num_class_images are present in class_data_dir during training time.Plain Text
Training on a 16GB GPU:
With the help of gradient checkpointing and the 8-bit optimizer from bitsandbytes it's possible to run train dreambooth on a 16GB GPU.
Plain Text
Training on a 8 GB GPU:
By using DeepSpeed it's possible to offload some tensors from VRAM to either CPU or NVME allowing to train with less VRAM.
DeepSpeed needs to be enabled with
accelerate config. During configuration answer yes to "Do you want to use DeepSpeed?". With DeepSpeed stage 2, fp16 mixed precision and offloading both parameters and optimizer state to cpu it's possible to train on under 8 GB VRAM with a drawback of requiring significantly more RAM (about 25 GB). See documentation for more DeepSpeed configuration options.Changing the default Adam optimizer to DeepSpeed's special version of Adam
deepspeed.ops.adam.DeepSpeedCPUAdam gives a substantial speedup but enabling it requires CUDA toolchain with the same version as pytorch. 8-bit optimizer does not seem to be compatible with DeepSpeed at the moment.Plain Text
Fine-tune text encoder with the UNet.
The script also allows to fine-tune the
text_encoder along with the unet. It's been observed experimentally that fine-tuning text_encoder gives much better results especially on faces. Pass the --train_text_encoder argument to the script to enable training text_encoder.Note: Training text encoder requires more memory, with this option the training won't fit on 16GB GPU. It needs at least 24GB VRAM.
Plain Text
Using DreamBooth for other pipelines than Stable Diffusion
Altdiffusion also support dreambooth now, the runing comman is basically the same as abouve, all you need to do is replace the
MODEL_NAME like this: One can now simply change the pretrained_model_name_or_path to another architecture such as AltDiffusion.Plain Text
Inference
Once you have trained a model using above command, the inference can be done simply using the
StableDiffusionPipeline. Make sure to include the identifier(e.g. sks in above example) in your prompt.Plain Text
Inference from a training checkpoint
You can also perform inference from one of the checkpoints saved during the training process, if you used the
--checkpointing_steps argument. Please, refer to the documentation to see how to do it.Training with Flax/JAX
For faster training on TPUs and GPUs you can leverage the flax training example. Follow the instructions above to get the model and dataset before running the script.
_Note: The flax example don't yet support features like gradient checkpoint, gradient accumulation etc, so to use flax for faster training we will need >30GB cards.
Before running the scripts, make sure to install the library's training dependencies:
Plain Text
Training without prior preservation loss
Plain Text
Training with prior preservation loss
Plain Text
Fine-tune text encoder with the UNet.
Plain Text
Training with xformers:
You can enable memory efficient attention by installing xFormers and padding the
--enable_xformers_memory_efficient_attention argument to the script. This is not available with the Flax/JAX implementation.You can also use Dreambooth to train the specialized in-painting model. See the script in the research folder for details.