diffusers: train_dreambooth_lora.py resume from checkpoint errors.

Describe the bug

I get the following issue when trying to resume from checkpoint. KeyError: 'unet.down_blocks.0.attentions.0.transformer_blocks.0.attn1.processor' There was also a naming issue where I had to change pytorch_lora_weights.bin to pytorch_model.bin in the checkpoint folders

Reproduction

accelerate launch train_dreambooth_lora.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
–checkpointing_steps=100
–learning_rate=1e-4
–report_to=“wandb”
–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”
–push_to_hub
–resume_from_checkpoint=“<checkpoint folder>” \

Logs

(most recent call last):
  File "/home/rip/git/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 1093, in <module>
    main(args)
  File "/home/rip/git/diffusers/examples/dreambooth/train_dreambooth_lora.py", line 881, in main
    accelerator.load_state(os.path.join(args.output_dir, path))
  File "/home/rip/.local/lib/python3.11/site-packages/accelerate/accelerator.py", line 2396, in load_state
    load_accelerator_state(
  File "/home/rip/.local/lib/python3.11/site-packages/accelerate/checkpointing.py", line 140, in load_accelerator_state
    models[i].load_state_dict(torch_mod, **load_model_func_kwargs)
  File "/home/rip/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 2027, in load_state_dict
    load(self, state_dict)
  File "/home/rip/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 2009, in load
    module._load_from_state_dict(
  File "/home/rip/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1909, in _load_from_state_dict
    hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
  File "/home/rip/.local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 69, in __call__
    return self.hook(module, *args, **kwargs)
  File "/home/rip/git/diffusers/src/diffusers/loaders.py", line 84, in map_from
    new_key = key.replace(replace_key, f"layers.{module.rev_mapping[replace_key]}")
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'unet.down_blocks.0.attentions.0.transformer_blocks.0.attn1.processor'

System Info

  • diffusers version: 0.17.0.dev0
  • Platform: Linux-5.10.16.3-microsoft-standard-WSL2-x86_64-with-glibc2.35
  • Python version: 3.11.3
  • PyTorch version (GPU?): 2.0.0+cu117 (True)
  • Huggingface_hub version: 0.14.0
  • Transformers version: 4.28.1
  • Accelerate version: 0.18.0
  • xFormers version: not installed
  • Using GPU in script?: yes
  • Using distributed or parallel set-up in script?: no

About this issue

  • Original URL
  • State: closed
  • Created a year ago
  • Comments: 22 (7 by maintainers)

Most upvoted comments

We have made a couple of fixes here so, we should be good here. @williamberman as well.