Skip to content

Commit a6fb940

Browse files
pcuencapcuenca
authored
Dreambooth docs: minor fixes (huggingface#1758)
* Section header for in-painting, inference from checkpoint. * Inference: link to section to perform inference from checkpoint. * Move Dreambooth in-painting instructions to the proper place.
1 parent 261a448 commit a6fb940

File tree

3 files changed

+100
-95
lines changed

3 files changed

+100
-95
lines changed

‎docs/source/training/dreambooth.mdx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,3 +283,5 @@ image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
283283

284284
image.save("dog-bucket.png")
285285
```
286+
287+
You may also run inference from [any of the saved training checkpoints](#performing-inference-using-a-saved-checkpoint).

‎examples/dreambooth/README.md‎

Lines changed: 5 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -232,8 +232,11 @@ image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
232232
image.save("dog-bucket.png")
233233
```
234234

235+
### Inference from a training checkpoint
235236

236-
## Running with Flax/JAX
237+
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](https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint) to see how to do it.
238+
239+
## Training with Flax/JAX
237240

238241
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.
239242

@@ -314,96 +317,4 @@ python train_dreambooth_flax.py \
314317
--max_train_steps=800
315318
```
316319

317-
### Training with prior-preservation loss
318-
319-
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.
320-
According to the paper, it's recommended to generate `num_epochs * num_samples` images for prior-preservation. 200-300 works well for most cases.
321-
322-
```bash
323-
export MODEL_NAME="runwayml/stable-diffusion-inpainting"
324-
export INSTANCE_DIR="path-to-instance-images"
325-
export CLASS_DIR="path-to-class-images"
326-
export OUTPUT_DIR="path-to-save-model"
327-
328-
accelerate launch train_dreambooth_inpaint.py \
329-
--pretrained_model_name_or_path=$MODEL_NAME \
330-
--instance_data_dir=$INSTANCE_DIR \
331-
--class_data_dir=$CLASS_DIR \
332-
--output_dir=$OUTPUT_DIR \
333-
--with_prior_preservation --prior_loss_weight=1.0 \
334-
--instance_prompt="a photo of sks dog" \
335-
--class_prompt="a photo of dog" \
336-
--resolution=512 \
337-
--train_batch_size=1 \
338-
--gradient_accumulation_steps=1 \
339-
--learning_rate=5e-6 \
340-
--lr_scheduler="constant" \
341-
--lr_warmup_steps=0 \
342-
--num_class_images=200 \
343-
--max_train_steps=800
344-
```
345-
346-
347-
### Training with gradient checkpointing and 8-bit optimizer:
348-
349-
With the help of gradient checkpointing and the 8-bit optimizer from bitsandbytes it's possible to run train dreambooth on a 16GB GPU.
350-
351-
To install `bitandbytes` please refer to this [readme](https://github.com/TimDettmers/bitsandbytes#requirements--installation).
352-
353-
```bash
354-
export MODEL_NAME="runwayml/stable-diffusion-inpainting"
355-
export INSTANCE_DIR="path-to-instance-images"
356-
export CLASS_DIR="path-to-class-images"
357-
export OUTPUT_DIR="path-to-save-model"
358-
359-
accelerate launch train_dreambooth_inpaint.py \
360-
--pretrained_model_name_or_path=$MODEL_NAME \
361-
--instance_data_dir=$INSTANCE_DIR \
362-
--class_data_dir=$CLASS_DIR \
363-
--output_dir=$OUTPUT_DIR \
364-
--with_prior_preservation --prior_loss_weight=1.0 \
365-
--instance_prompt="a photo of sks dog" \
366-
--class_prompt="a photo of dog" \
367-
--resolution=512 \
368-
--train_batch_size=1 \
369-
--gradient_accumulation_steps=2 --gradient_checkpointing \
370-
--use_8bit_adam \
371-
--learning_rate=5e-6 \
372-
--lr_scheduler="constant" \
373-
--lr_warmup_steps=0 \
374-
--num_class_images=200 \
375-
--max_train_steps=800
376-
```
377-
378-
### Fine-tune text encoder with the UNet.
379-
380-
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.
381-
Pass the `--train_text_encoder` argument to the script to enable training `text_encoder`.
382-
383-
___Note: Training text encoder requires more memory, with this option the training won't fit on 16GB GPU. It needs at least 24GB VRAM.___
384-
385-
```bash
386-
export MODEL_NAME="runwayml/stable-diffusion-inpainting"
387-
export INSTANCE_DIR="path-to-instance-images"
388-
export CLASS_DIR="path-to-class-images"
389-
export OUTPUT_DIR="path-to-save-model"
390-
391-
accelerate launch train_dreambooth_inpaint.py \
392-
--pretrained_model_name_or_path=$MODEL_NAME \
393-
--train_text_encoder \
394-
--instance_data_dir=$INSTANCE_DIR \
395-
--class_data_dir=$CLASS_DIR \
396-
--output_dir=$OUTPUT_DIR \
397-
--with_prior_preservation --prior_loss_weight=1.0 \
398-
--instance_prompt="a photo of sks dog" \
399-
--class_prompt="a photo of dog" \
400-
--resolution=512 \
401-
--train_batch_size=1 \
402-
--use_8bit_adam \
403-
--gradient_checkpointing \
404-
--learning_rate=2e-6 \
405-
--lr_scheduler="constant" \
406-
--lr_warmup_steps=0 \
407-
--num_class_images=200 \
408-
--max_train_steps=800
409-
```
320+
You can also use Dreambooth to train the specialized in-painting model. See [the script in the research folder for details](https://github.com/huggingface/diffusers/tree/main/examples/research_projects/dreambooth_inpaint).

‎examples/research_projects/dreambooth_inpaint/README.md‎

Lines changed: 93 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,96 @@ accelerate launch train_dreambooth_inpaint.py \
2323
--max_train_steps=400
2424
```
2525

26-
The script is also compatible with prior preservation loss and gradient checkpointing
26+
### Training with prior-preservation loss
27+
28+
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.
29+
According to the paper, it's recommended to generate `num_epochs * num_samples` images for prior-preservation. 200-300 works well for most cases.
30+
31+
```bash
32+
export MODEL_NAME="runwayml/stable-diffusion-inpainting"
33+
export INSTANCE_DIR="path-to-instance-images"
34+
export CLASS_DIR="path-to-class-images"
35+
export OUTPUT_DIR="path-to-save-model"
36+
37+
accelerate launch train_dreambooth_inpaint.py \
38+
--pretrained_model_name_or_path=$MODEL_NAME \
39+
--instance_data_dir=$INSTANCE_DIR \
40+
--class_data_dir=$CLASS_DIR \
41+
--output_dir=$OUTPUT_DIR \
42+
--with_prior_preservation --prior_loss_weight=1.0 \
43+
--instance_prompt="a photo of sks dog" \
44+
--class_prompt="a photo of dog" \
45+
--resolution=512 \
46+
--train_batch_size=1 \
47+
--gradient_accumulation_steps=1 \
48+
--learning_rate=5e-6 \
49+
--lr_scheduler="constant" \
50+
--lr_warmup_steps=0 \
51+
--num_class_images=200 \
52+
--max_train_steps=800
53+
```
54+
55+
56+
### Training with gradient checkpointing and 8-bit optimizer:
57+
58+
With the help of gradient checkpointing and the 8-bit optimizer from bitsandbytes it's possible to run train dreambooth on a 16GB GPU.
59+
60+
To install `bitandbytes` please refer to this [readme](https://github.com/TimDettmers/bitsandbytes#requirements--installation).
61+
62+
```bash
63+
export MODEL_NAME="runwayml/stable-diffusion-inpainting"
64+
export INSTANCE_DIR="path-to-instance-images"
65+
export CLASS_DIR="path-to-class-images"
66+
export OUTPUT_DIR="path-to-save-model"
67+
68+
accelerate launch train_dreambooth_inpaint.py \
69+
--pretrained_model_name_or_path=$MODEL_NAME \
70+
--instance_data_dir=$INSTANCE_DIR \
71+
--class_data_dir=$CLASS_DIR \
72+
--output_dir=$OUTPUT_DIR \
73+
--with_prior_preservation --prior_loss_weight=1.0 \
74+
--instance_prompt="a photo of sks dog" \
75+
--class_prompt="a photo of dog" \
76+
--resolution=512 \
77+
--train_batch_size=1 \
78+
--gradient_accumulation_steps=2 --gradient_checkpointing \
79+
--use_8bit_adam \
80+
--learning_rate=5e-6 \
81+
--lr_scheduler="constant" \
82+
--lr_warmup_steps=0 \
83+
--num_class_images=200 \
84+
--max_train_steps=800
85+
```
86+
87+
### Fine-tune text encoder with the UNet.
88+
89+
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.
90+
Pass the `--train_text_encoder` argument to the script to enable training `text_encoder`.
91+
92+
___Note: Training text encoder requires more memory, with this option the training won't fit on 16GB GPU. It needs at least 24GB VRAM.___
93+
94+
```bash
95+
export MODEL_NAME="runwayml/stable-diffusion-inpainting"
96+
export INSTANCE_DIR="path-to-instance-images"
97+
export CLASS_DIR="path-to-class-images"
98+
export OUTPUT_DIR="path-to-save-model"
99+
100+
accelerate launch train_dreambooth_inpaint.py \
101+
--pretrained_model_name_or_path=$MODEL_NAME \
102+
--train_text_encoder \
103+
--instance_data_dir=$INSTANCE_DIR \
104+
--class_data_dir=$CLASS_DIR \
105+
--output_dir=$OUTPUT_DIR \
106+
--with_prior_preservation --prior_loss_weight=1.0 \
107+
--instance_prompt="a photo of sks dog" \
108+
--class_prompt="a photo of dog" \
109+
--resolution=512 \
110+
--train_batch_size=1 \
111+
--use_8bit_adam \
112+
--gradient_checkpointing \
113+
--learning_rate=2e-6 \
114+
--lr_scheduler="constant" \
115+
--lr_warmup_steps=0 \
116+
--num_class_images=200 \
117+
--max_train_steps=800
118+
```

0 commit comments

Comments
 (0)