dreambooth training script dose not make "unet" forler for each checkpointed directory

[yoon28]

Describe the bug

according to this document, dreambooth training script should make "unet" folder for each checkpointed directory.
but it is not.

https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint

Reproduction

launch dreambooth training script
eg) CUDA_VISIBLE_DEVICES=0,1 accelerate launch train_dreambooth.py --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4" --instance_data_dir=/mnt/volatile/data --class_data_dir=/mnt/volatile/dreambooth/gen/person --output_dir=/mnt/volatile/dreambooth/log/ --with_prior_preservation --prior_loss_weight=1.0 --instance_prompt="ykj person" --class_prompt="person" --resolution=512 --train_batch_size=4 --gradient_accumulation_steps=1 --learning_rate=2e-6 --lr_scheduler="constant" --lr_warmup_steps=0 --num_class_images=250 --max_train_steps=1200 --train_text_encoder --checkpointing_steps=100

After that, check the folders that were checkpointed. And you can not see the "unet" folder.
Finally, I can not sample from checkpoints unlike this

Logs

No response

System Info

• diffusers version: 0.14.0.dev0
• Platform: Linux-3.10.0-1160.59.1.el7.x86_64-x86_64-with-glibc2.31
• Python version: 3.10.9
• PyTorch version (GPU?): 1.13.1+cu117 (True)
• Huggingface_hub version: 0.12.1
• Transformers version: 4.26.1
• Accelerate version: 0.16.0
• xFormers version: not installed
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

[patrickvonplaten]

@williamberman can you take a look here?

[patrickvonplaten]

Maybe @sayakpaul ?

[sayakpaul]

Hi.

I was unable to reproduce this problem.

Here's the output of my diffusers-cli env:

- `diffusers` version: 0.15.0.dev0
- Platform: Linux-4.19.0-23-cloud-amd64-x86_64-with-glibc2.10
- Python version: 3.8.16
- PyTorch version (GPU?): 1.13.1+cu116 (True)
- Huggingface_hub version: 0.13.2
- Transformers version: 4.26.1
- Accelerate version: 0.18.0
- xFormers version: 0.0.16
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

Here's the command I tested train_dreambooth.py with:

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export INSTANCE_DIR="dog_images"
export OUTPUT_DIR="dreambooth"

accelerate launch --mixed_precision="fp16" train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --mixed_precision="fp16" \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=800 \
  --checkpointing_steps=50

After the command has finished execution, I do get a dreambooth folder structured like the following:

dreambooth
├── checkpoint-100
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-150
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-200
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-250
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-300
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-350
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-400
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-450
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-50
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-500
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-550
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-600
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-650
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-700
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-750
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── checkpoint-800
│   ├── optimizer.bin
│   ├── random_states_0.pkl
│   ├── scaler.pt
│   ├── scheduler.bin
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.bin
├── feature_extractor
│   └── preprocessor_config.json
├── logs
│   └── dreambooth
│       ├── 1680580171.5524163
│       │   └── events.out.tfevents.1680580171.sayak-experiments.19237.1
│       ├── 1680580171.5542052
│       │   └── hparams.yml
│       └── events.out.tfevents.1680580171.sayak-experiments.19237.0
├── model_index.json
├── safety_checker
│   ├── config.json
│   └── pytorch_model.bin
├── scheduler
│   └── scheduler_config.json
├── text_encoder
│   ├── config.json
│   └── pytorch_model.bin
├── tokenizer
│   ├── merges.txt
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   └── vocab.json
├── unet
│   ├── config.json
│   └── diffusion_pytorch_model.bin
└── vae
    ├── config.json
    └── diffusion_pytorch_model.bin

Also, I am able to do the following:

from diffusers import DiffusionPipeline, UNet2DConditionModel
from transformers import CLIPTextModel
import torch

# Load the pipeline with the same arguments (model, revision) that were used for training
model_id = "CompVis/stable-diffusion-v1-4"

unet = UNet2DConditionModel.from_pretrained("dreambooth/checkpoint-200/unet")


pipeline = DiffusionPipeline.from_pretrained(model_id, unet=unet, dtype=torch.float16)
pipeline.to("cuda")

# Perform inference, or save, or push to the hub
pipeline.save_pretrained("dreambooth-pipeline")

[yoon28]

thanks! i will try again!

[sayakpaul]

@yoon28 did you find some time to check it?

[amitgurintelcom]

Hi,
I have the same issue.
@yoon28 trained with --train_text_encoder, and the checkpoint saved only text_encoder/ folder. It didn't save unet folder.
While @sayakpaul didn't use --train_text_encoder option.
This is my command line:
accelerate launch train_dreambooth.py
--mixed_precision="fp16"
--revision="fp16"
--instance_data_dir="./XYZ"
--output_dir="checkpoints"
--instance_prompt="XYZ"
--resolution=512
--train_batch_size=2
--gradient_accumulation_steps=1
--learning_rate=1.5e-6
--lr_scheduler="constant"
--lr_warmup_steps=0
--max_train_steps=1200
--num_class_images=260
--with_prior_preservation
--prior_loss_weight=1.0
--class_prompt="man"
--class_data_dir="./man"
*--train_text_encoder*
--validation_prompt="Portrait of XYZ"
--validation_steps=100
--checkpointing_steps=500
--pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4"

This is my config:

• diffusers version: 0.17.0.dev0
• Platform: Linux-5.15.0-71-generic-x86_64-with-glibc2.29
• Python version: 3.8.10
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Huggingface_hub version: 0.13.3
• Transformers version: 4.28.1
• Accelerate version: 0.18.0
• xFormers version: not installed
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

[amitgurintelcom]

I think I found the bug in train_dreambooth.py

    def save_model_hook(models, weights, output_dir):
        for model in models:
            sub_dir = "unet" if type(model) == type(unet) else "text_encoder"
            model.save_pretrained(os.path.join(output_dir, sub_dir))

            # make sure to pop weight so that corresponding model is not saved again
            weights.pop()

The issue is that:
type(model) = <class 'diffusers.models.unet_2d_condition.UNet2DConditionModel'>
type(unet) = <class 'torch.nn.parallel.distributed.DistributedDataParallel'>

Fix should be:

    def save_model_hook(models, weights, output_dir):
        for model in models:
            sub_dir = "unet" if "unet" in str(type(model)) else "text_encoder"
            model.save_pretrained(os.path.join(output_dir, sub_dir))

            # make sure to pop weight so that corresponding model is not saved again
            weights.pop()

[sayakpaul]

Replied over the other thread.