load_lora_weights function fails to load custom lora

[akameswa]

Describe the bug

Base model: segmind/tiny-sd
LoRA: akameswa/lcm-lora-tiny-sd

Error:
Got the following error using load_lora_weights.
ValueError: Target modules {'base_model.model.up_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q', ...} not found in the base model. Please check the target modules and try again.

Steps taken:

1. Checked if the target modules are absent in the base model. But, they are present.
2. Tried unet.load_attn_procs from here. The results are bad and blurry.

Fix:
Using the following code worked.

from peft import PeftModel, PeftConfig
config = PeftConfig.from_pretrained(adapter_id)
lora_unet = PeftModel.from_pretrained(pipe.unet, adapter_id)

Training:
Follow LCM-LoRA at tutorial with
MODEL_NAME='segmind/tiny-sd' and DATASET_DIR='akameswa/improved_aesthetics_6.5plus_webdataset'

Additional comments:
Did not face any issues during training. Attaching training log

Note:
Please refer to the attached colab. Walks through the steps taken, the errors, and the results in a clear manner.

Reproduction

https://colab.research.google.com/drive/135v4cBVFu_zxXhcyao5lVZFlfscde0n0?usp=sharing

Logs

No response

System Info

• diffusers version: 0.25.0.dev0
• Platform: Windows-10-10.0.22635-SP0
• Python version: 3.10.13
• PyTorch version (GPU?): 2.1.2+cu121 (True)
• Huggingface_hub version: 0.19.4
• Transformers version: 4.36.2
• Accelerate version: 0.25.0
• peft version: 0.7.0
• xFormers version: 0.0.23.post1

Who can help?

@sayakpaul @patrickvonplaten @DN6

[sayakpaul]

Follow how it's done here:

diffusers/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py

Line 253 in 9f283b0

def log_validation(vae, unet, args, accelerator, weight_dtype, step):

But if you follow the serialization logic of https://github.com/huggingface/diffusers/blob/9f283b01d20aa8d9e5251dbaad19432bfb2ffe11/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py#L1334C1-L1336C108, it should all work right of box. We recommend repurposing this variant because:

• It's tested.
• Follows best practices from peft by not initializing two UNets. Hence more memory savings.

[akameswa]

I need clarification on how to load the lcm-lora unet.

diffusers/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py

Line 256 in 9f283b0

unet = accelerator.unwrap_model(unet)

Replaced above line with unet = UNet2DConditionModel.from_pretrained("akameswa/lcm-lora-tiny-sd", use_safetensors=True). But resulted in a similar error.

ValueError: Cannot load <class 'diffusers.models.unet_2d_condition.UNet2DConditionModel'> from akameswa/lcm-lora-tiny-sd because the following keys are missing:
down_blocks.0.attentions.1.transformer_blocks.0.norm2.bias, ..

Attaching colab for reference

[sayakpaul]

I think there's some misunderstanding here.

If you want to have the UNet already populated with the LCM LoRA parameters, there are ways, including the one you wrote using peft:

from peft import PeftModel, PeftConfig

lora_unet = PeftModel.from_pretrained(pipe.unet, adapter_id)

What it does is load the adapter and then merge it in the UNet. The adapter config file has the base model ID inside of it, so, the base model is first correctly loaded and then the loaded adapter is merged. Here the assumption is that the directory or the repository corresponding to adapter_id has two files: the adapter state dict and the related config.

However, if you don't want to do it this way i.e., without using peft, you can follow how it's done here:

diffusers/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py

Line 99 in bf725e0

def log_validation(vae, args, accelerator, weight_dtype, step, unet=None, is_final_validation=False):

• Here the UNet is initialized from the pretrained_teacher_model.
• And then the adapter parameters are loaded using load_lora_weights().
• Here the adapter checkpoint is NOT in the peft format but rather is in the diffusers format:

  diffusers/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py

  Line 116 in bf725e0

  state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))

  .

I have tried my best to explain all the bits :D If something's not clear, let me know. However, I also recommend you to scan through the code minutely as the details might pop up when doing so.

[akameswa]

Thanks for the explanation. I used the sd_wds training script and not sdxl.
https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py

I'm with you till the second bullet point. But I believe the adapter checkpoint here is not in peft and diffusers but in kohya

diffusers/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py

Lines 249 to 250 in 61d223c

	lora_state_dict = get_module_kohya_state_dict(unet, "lora_unet", weight_dtype)
	pipeline.load_lora_weights(lora_state_dict)

To populate the lora_state_dict, I believe the unet should be initialized from the adapter checkpoint and not pretrained_teacher_model as mentioned here. So, I used unet = UNet2DConditionModel.from_pretrained("akameswa/lcm-lora-tiny-sd", use_safetensors=True) to initialize the adapter checkpoint, causing the following error. Please let me know if my understanding is wrong.

I need clarification on how to load the lcm-lora unet.

diffusers/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py

Line 256 in 9f283b0

unet = accelerator.unwrap_model(unet)

Replaced above line with unet = UNet2DConditionModel.from_pretrained("akameswa/lcm-lora-tiny-sd", use_safetensors=True). But resulted in a similar error.

ValueError: Cannot load <class 'diffusers.models.unet_2d_condition.UNet2DConditionModel'> from akameswa/lcm-lora-tiny-sd because the following keys are missing: down_blocks.0.attentions.1.transformer_blocks.0.norm2.bias, ..

Attaching colab for reference

I am fine using the method I suggested using peft, but I wanted it to work right out of the box by following how it's done here.

diffusers/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py

Line 235 in 61d223c

def log_validation(vae, unet, args, accelerator, weight_dtype, step):

Kindly note that the loading of state_dict is different in sd_wds and sdxl

[sayakpaul]

Cc: @patil-suraj

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[sayakpaul]

@patil-suraj a gentle bump.

[rohit901]

I have the same question on loading trained LoRA weights using the script [for SD, WDS format] in https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py

For example, in this pre-trained SD1.5 LCM LoRA , the code is as follows:

import torch
from diffusers import LCMScheduler, AutoPipelineForText2Image

model_id = "Lykon/dreamshaper-7"
adapter_id = "latent-consistency/lcm-lora-sdv1-5"

pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")

# load and fuse lcm lora
pipe.load_lora_weights(adapter_id)
pipe.fuse_lora()

If we directly replace adapter_id in the above code with path to our local directory containing the trained LoRA UNet, it does not work and gives error saying target module not present in base model.

It would be nice if the above loading code could directly work with the given script since it seems to be more generic?

@sayakpaul @patil-suraj

[AgainstEntropy]

It seems that the key mismatching problem is caused by PeftModel wrapping the base model into its self.base_model.model (see its source code)

So the simpliest way is removing the "base_model.model" from all keys in the state_dict. Below is a sample function to do this.

def unwrap_peft_state_dict(peft_model: PeftModel, dtype: torch.dtype, adapter_name: str = "default"):
    unwrapped_state_dict = {}
    for peft_key, weight in get_peft_model_state_dict(peft_model, adapter_name=adapter_name).items():
        key = peft_key.replace("base_model.model.", "")
        unwrapped_state_dict[key] = weight.to(dtype)

    return unwrapped_state_dict

Hope the following code snippet works for you guys.

pipeline = StableDiffusionPipeline.from_pretrained(
    ...
)

unwrapped_unet = accelerator.unwrap_model(unet)
lora_state_dict = unwrap_peft_state_dict(unwrapped_unet, weight_dtype)
pipeline.load_lora_weights(lora_state_dict)
pipeline.fuse_lora()