train_dreambooth_lora.py Does not run on latest

[Ambrosiussen]

Describe the bug

train_dreambooth_lora.py example no longer runs when installing diffusers from source.
If I run the script using the provided snippet, I get the following error.

Command run:
accelerate launch train_dreambooth_lora.py --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" --instance_data_dir="./dog/" --output_dir="./out/" --instance_prompt="a photo of sks dog" --resolution=512 --train_batch_size=1 --gradient_accumulation_steps=1 --checkpointing_steps=100 --learning_rate=1e-4 --lr_scheduler="constant" --lr_warmup_steps=0 --max_train_steps=500 --validation_prompt="A photo of sks dog in a bucket" --validation_epochs=50 --seed="0"

I am feeding it the dataset of the dog example provided (5 images)

Traceback (most recent call last):
File "H:\Github\MLOPs\MLOPs\scripts\python\sdpipeline\train_dreambooth_lora.py", line 1351, in
main(args)
File "H:\Github\MLOPs\MLOPs\scripts\python\sdpipeline\train_dreambooth_lora.py", line 1142, in main
model_pred = unet(noisy_model_input, timesteps, encoder_hidden_states).sample
File "C:\Users\Paul\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "C:\Users\Paul\AppData\Local\Programs\Python\Python39\lib\site-packages\diffusers\models\unet_2d_condition.py", line 718, in forward
sample = self.conv_in(sample)
File "C:\Users\Paul\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "C:\Users\Paul\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\conv.py", line 463, in forward
return self._conv_forward(input, self.weight, self.bias)
File "C:\Users\Paul\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\conv.py", line 459, in _conv_forward
return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [320, 4, 3, 3], expected input[1, 3, 512, 512] to have 4 channels, but got 3 channels instead

Reproduction

Run: (The example on the Github page for Dreambooth training)

accelerate launch train_dreambooth_lora.py --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" --instance_data_dir="./dog/" --output_dir="./out/" --instance_prompt="a photo of sks dog" --resolution=512 --train_batch_size=1 --gradient_accumulation_steps=1 --checkpointing_steps=100 --learning_rate=1e-4 --lr_scheduler="constant" --lr_warmup_steps=0 --max_train_steps=500 --validation_prompt="A photo of sks dog in a bucket" --validation_epochs=50 --seed="0"

Logs

No response

System Info

• diffusers version: 0.17.0.dev0
• Platform: Windows-10-10.0.22621-SP0
• Python version: 3.9.10
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Huggingface_hub version: 0.14.1
• Transformers version: 4.29.0.dev0
• Accelerate version: 0.18.0
• xFormers version: 0.0.19
• Using GPU in script?: YES
• Using distributed or parallel set-up in script?: NO

[patrickvonplaten]

I cannot reproduce the problem (but I'm also testing on Linux, not on Windows).

I ran:

accelerate launch train_lora_dreambooth.py --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" --instance_data_dir="./dog/" --output_dir="./out/" --instance_prompt="a photo of sks dog" --resolution=512 --train_batch_size=1 --gradient_accumulation_steps=1 --checkpointing_steps=100 --learning_rate=1e-4 --lr_scheduler="constant" --lr_warmup_steps=0 --max_train_steps=500 --validation_prompt="A photo of sks dog in a bucket" --validation_epochs=50 --seed="0"

with accelerate config:

compute_environment: LOCAL_MACHINE
distributed_type: 'NO'
downcast_bf16: 'no'
gpu_ids: '0'
machine_rank: 0
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 1
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

and the ENV:

- `diffusers` version: 0.17.0.dev0
- Platform: Linux-5.15.0-69-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Huggingface_hub version: 0.14.1
- Transformers version: 4.30.0.dev0
- Accelerate version: 0.20.0.dev0
- xFormers version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

[Ambrosiussen]

You seem to be running train_dreambooth.py. Not the LORA variant. Could you try that one as well?

[patrickvonplaten]

Ah sorry, I was actually running train_lora_dreambooth.py (did a symlink). Will edit above

[patrickvonplaten]

Also this error:

RuntimeError: Given groups=1, weight of size [320, 4, 3, 3], expected input[1, 3, 512, 512] to have 4 channels, but got 3 channels instead

really sounds like the images are not encoded with a VAE which is surprising.
So the reason for this error is surely that you're in the wrong branch here:

diffusers/examples/dreambooth/train_dreambooth_lora.py

Line 1110 in 6070b32

if vae is not None:

[Ambrosiussen]

I will try updating Transformers and Accelerate to latest to match your env. Will report back

[patrickvonplaten]

There is a problem with this function:

diffusers/examples/dreambooth/train_dreambooth_lora.py

Line 592 in 6070b32

def model_has_vae(args):

Curious to find out what might be going on here

[Ambrosiussen]

If I pass in the path to the runwayml/sd1.5 model I downloaded manually somewhere instead of just the name (which will grab the huggingface cache) it works! Something may be wrong with the code for downloading the model?

[patrickvonplaten]

Actually, the script was not written in a very robust way here - the error definitely came from the script thinking that the model doesn't have a VAE. Changed it to a more robust way here: #3462 cc @williamberman

[Ambrosiussen]

Thanks so much! Great work on diffusers :)

[williamberman]

ah sorry, thanks for the fix!

[bghira]

for anyone stumbling upon this in the future, this check also fails to perform robustly.

for example, any failure to load the state dict is thrown as an OSError from the from_pretrained method, which then makes vae = None and then the same exact channel error occurs