Victarry / stable-dreambooth

Dreambooth implementation based on Stable Diffusion with minimal code.

Stable DreamBooth

This is an implementation of DreamBooth based on Stable Diffusion.

Update

• This repository has been migrated into diffusers! Less than 14 GB memeory is required!
  See more in https://github.com/huggingface/diffusers/tree/main/examples/dreambooth

Results

Dreambooth results from original paper:

The reproduced results:

Requirements

Hardware

• A GPU with at least 30G Memory.
• The training requires about 10 minites on A100 80G GPU with batch_size set to 4.

Environment Setup

Create conda environment with pytorch>=1.11.

conda env create -f environment.yaml
conda activate stable-diffusion

Quick Start

python sample.py # Generate class samples.
python train.py # Finetune stable diffusion model.

The generation results are in logs/dog_finetune.

Finetune with your own data.

1. Data Preparation

1. Collect 3~5 images of an object and save into data/mydata/instance folder.
2. Sample images of the same class as specified object using sample.py.
   1. Change corresponding variables in sample.py. The prompt should be like "a {class}". And the save_dir should be changed to data/mydata/class.
   2. Run the sample script.

   python sample.py

2. Finetuning

1. Change the TrainConfig in train.py.
2. Start training.

   python train.py

3. Inference

python inference.py --prompt "photo of a [V] dog in a dog house" --checkpoint_dir logs/dogs_finetune

Generated images are in outputs by default.

Acknowledgement

• Stable Diffusion by CompVis https://github.com/CompVis/stable-diffusion
• DreamBooth https://dreambooth.github.io/
• Diffusers https://github.com/huggingface/diffusers