Bug: lora trained with train_dreambooth_lora_sdxl_advanced.py, the weight cannot be loaded

[garychan22]

Describe the bug

I have trained a sdxl lora with the script https://github.com/huggingface/diffusers/blob/7bc0a07b1947bcbe0e84bbe9ecf8ae2d234382c4/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py , and a loading error was encountered during the validation in the end of the training whenc calling load_lora_weights

diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

Line 2426 in 7bc0a07

pipeline.load_lora_weights(args.output_dir)

Loading adapter weights from state_dict led to missing keys in the model: down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q.lora_A.default_0.weigh
t, down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q.lora_B.default_0.weight, down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k.lora_A.de
fault_0.weight, down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k.lora_B.default_0.weight, down_blocks.1.attentions.0.transformer_blocks.0.attn1.t
o_v.lora_A.default_0.weight, down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_v.lora_B.default_0.weight, down_blocks.1.attentions.0.transformer_blo
cks.0.attn1.to_out.0.lora_A.default_0.weight, down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.0.lora_B.default_0.weight, down_blocks.1.attenti
ons.0.transformer_blocks.0.attn2.to_q.lora_A.default_0.weight, down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_q.lora_B.default_0.weight, 
....

Reproduction

running the https://github.com/huggingface/diffusers/blob/7bc0a07b1947bcbe0e84bbe9ecf8ae2d234382c4/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

by changing the target_modules = ["to_k", "to_q", "to_v", "to_out.0","proj_in", "proj_out","conv1", "conv2"] for the unet

Logs

System Info

linux python 3.9.21
diffusers==0.34.0.dev0
peft==0.15.0

Who can help?

No response

[nsabir2011]

I am also facing the same issue. I've noticed that it occurs only when "proj_in", "proj_out" layers are added.
Most probably caused by:

diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

Line 2382 in 7bc0a07

unet_lora_layers = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))

Going through the code in the conversion file seems like it is only focusing on the self-attention modules of the UNet. But it is common to include other modules in the community.

That being said, I am not 100% sure if the above is causing this issue with my limited knowledge.

[garychan22]

I am also facing the same issue. I've noticed that it occurs only when "proj_in", "proj_out" layers are added. Most probably caused by:

diffusers/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

Line 2382 in 7bc0a07

unet_lora_layers = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
Going through the code in the conversion file seems like it is only focusing on the self-attention modules of the UNet. But it is common to include other modules in the community.

That being said, I am not 100% sure if the above is causing this issue with my limited knowledge.

I have also noticed that the function convert_state_dict_to_diffusers or get_peft_model_state_dict produce incorrect state_dict keys when including proj_in, proj_out, conv* in the target_modules in the script. The similar layers corresponding to these `incorrect state_dict keys`` are also included in the weights pretrained by koyha_ss, which has been supported by the latest diffusers. Thus, there is no reason that the pretrained weights cannot be loaded by diffusers if the saved state_dict is correct.

At the time, you can just use the PeftModel instead of unet.add_adapter(*) to train and load the pretrained weights. However, this solution will prevent you from loading the other pretrained loras from the community which are usually not trained with PeftModel.