New features on LORA break lora+dreambooth model loading

[tcapelle]

Describe the bug

Unable to load the attn_proc on the freshly trained dreambooth+lora. Just using the supplied example. It was working on 7b0ba48

Reproduction

Train a model and save the new weights in "pytorch_lora_weights.bin". That works fine, then try to create inference with that model by doing this:

model_base = "runwayml/stable-diffusion-v1-5"

pipeline = DiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16)
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)

lora_model_path = "pytorch_lora_weights.bin"
pipeline.unet.load_attn_procs(lora_model_path)

Logs

The following command is not working anymore:

pipeline.unet.load_attn_procs(lora_model_path)

The error:

---

KeyError Traceback (most recent call last)
Cell In[4], line 1
----> 1 pipeline.unet.load_attn_procs(lora_model_path)
2 pipeline.to("cuda")

File ~/Apps/diffusers/src/diffusers/loaders.py:279, in UNet2DConditionLoadersMixin.load_attn_procs(self, pretrained_model_name_or_path_or_dict, **kwargs)
276 attn_processors = {k: v.to(device=self.device, dtype=self.dtype) for k, v in attn_processors.items()}
278 # set layers
--> 279 self.set_attn_processor(attn_processors)

File ~/Apps/diffusers/src/diffusers/models/unet_2d_condition.py:513, in UNet2DConditionModel.set_attn_processor(self, processor)
510 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
512 for name, module in self.named_children():
--> 513 fn_recursive_attn_processor(name, module, processor)

File ~/Apps/diffusers/src/diffusers/models/unet_2d_condition.py:510, in UNet2DConditionModel.set_attn_processor..fn_recursive_attn_processor(name, module, processor)
507 module.set_processor(processor.pop(f"{name}.processor"))
509 for sub_name, child in module.named_children():
--> 510 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

File ~/Apps/diffusers/src/diffusers/models/unet_2d_condition.py:510, in UNet2DConditionModel.set_attn_processor..fn_recursive_attn_processor(name, module, processor)
507 module.set_processor(processor.pop(f"{name}.processor"))
509 for sub_name, child in module.named_children():
--> 510 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

[... skipping similar frames: UNet2DConditionModel.set_attn_processor.<locals>.fn_recursive_attn_processor at line 510 (3 times)]

File ~/Apps/diffusers/src/diffusers/models/unet_2d_condition.py:510, in UNet2DConditionModel.set_attn_processor..fn_recursive_attn_processor(name, module, processor)
507 module.set_processor(processor.pop(f"{name}.processor"))
509 for sub_name, child in module.named_children():
--> 510 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

File ~/Apps/diffusers/src/diffusers/models/unet_2d_condition.py:507, in UNet2DConditionModel.set_attn_processor..fn_recursive_attn_processor(name, module, processor)
505 module.set_processor(processor)
506 else:
--> 507 module.set_processor(processor.pop(f"{name}.processor"))
509 for sub_name, child in module.named_children():
510 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

KeyError: 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.processor'

System Info

• diffusers version: 0.16.0.dev0
• Platform: Linux-5.15.0-1031-gcp-x86_64-with-glibc2.31
• Python version: 3.10.8
• PyTorch version (GPU?): 1.13.1 (True)
• Huggingface_hub version: 0.13.4
• Transformers version: 4.28.1
• Accelerate version: 0.18.0
• xFormers version: 0.0.18
• Using GPU in script?: V100
• Using distributed or parallel set-up in script?: NO

[tcapelle]

cc. @sayakpaul it's related to PR #3130

[yiyixuxu]

@tcapelle
ohh I think I misunderstood the issue - it's with the new training script. ignore my comments here

would love to take a look but I can't reproduce on my end. Below code works for me

from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
import torch

model_base = "runwayml/stable-diffusion-v1-5"

pipeline = DiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16)
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)

lora_model_path = "patrickvonplaten/lora_dreambooth_dog_example"
pipeline.unet.load_attn_procs(lora_model_path)

[tcapelle]

Yeah the new training scripts (or changes) make the trained model break after loading for inference and trying to do the load_attn_procs

[yiyixuxu]

@tcapelle

so I think you have to load this instead (since now the lora weights include not only unet but also text_encoder)

pipeline.load_attn_procs(lora_model_path)

cc @sayakpaul here: what should we do here? it will cause a lot of confusion - is updating doc sufficient?

[sayakpaul]

Yeah, it's my bad that I didn't include enough documentation about it. Sorry.

I will prioritize and fix the documentation :)

[sayakpaul]

It should work with the following:

from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
import torch

model_base = "runwayml/stable-diffusion-v1-5"

pipeline = DiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16)
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
pipeline.load_lora_weights("sayakpaul/dreambooth")

Where "sayakpaul/dreambooth" was obtained using #3130 (comment).

But this is not documented anywhere. I am currently working on it.

[sayakpaul]

#3180 should address these concerns better. PTAL.

[Hzj199]

I also encountered this problem and fixed it by changing line 765 in loaders.py.

 # unet_lora_state_dict = {k: v for k, v in state_dict.items() if k.startswith(self.unet_name)}
 unet_lora_state_dict = {k.replace(f"{self.unet_name}.", ""): v for k, v in state_dict.items() if k.startswith(self.unet_name)}

@tcapelle

[sing817]

我也遇到了這個問題並通過更改 loaders.py 中的第 765 行修復了它。

 # unet_lora_state_dict = {k: v for k, v in state_dict.items() if k.startswith(self.unet_name)}
 unet_lora_state_dict = {k.replace(f"{self.unet_name}.", ""): v for k, v in state_dict.items() if k.startswith(self.unet_name)}

@tcapelle

it work.

[sayakpaul]

@tcapelle feel free to re-open the issue if the problems persist after #3180

[jelling]

For anyone else hitting this issue now, the PR has been merged but the package has not hit pip yet. But installing diffusers from Github will fix it:

pip install git+https://github.com/huggingface/diffusers

[jelling]

@sayakpaul you may want to update the docs here too: https://huggingface.co/docs/diffusers/training/lora#dreambooth-inference