Diffusers 0.30.1: Merging LoRA adapters without modifying original model structure

[asfiyab-nvidia]

Is there a way to fully merge LoRA weights into the base model in newer versions of Diffusers, resulting in an unmodified model graph structure similar to how it worked in version 0.26.3?

Previous behavior (Diffusers 0.26.3):

1. LoRA weights were fully merged into the original model weights.
2. The model structure remained unchanged after merging.

Current behavior (Diffusers 0.30.1):

1. LoRA adapters are integrated using PEFT (Parameter-Efficient Fine-Tuning).
2. This results in additional layers for LoRA weights in the model structure.
3. For example, a simple Linear layer is replaced by a more complex lora.Linear structure as shown below

Before merging

...
(proj_out): Linear(in_features=1280, out_features=1280, bias=True)
...

After merging

...
(proj_out): lora.Linear(
          (base_layer): Linear(in_features=1280, out_features=1280, bias=True)
          (lora_dropout): ModuleDict(
            (ikea): Identity()
          )
          (lora_A): ModuleDict(
            (ikea): Linear(in_features=1280, out_features=32, bias=False)
          )
          (lora_B): ModuleDict(
            (ikea): Linear(in_features=32, out_features=1280, bias=False)
          )
          (lora_embedding_A): ParameterDict()
          (lora_embedding_B): ParameterDict()
          (lora_magnitude_vector): ModuleDict()
        )
...

I did follow the recommendations in https://huggingface.co/docs/diffusers/en/using-diffusers/merge_loras but observed the same behavior as listed above.

Goal

My use case requires merging LoRA weights into the model in a way that preserves the original model structure.

Reproduce

Script to reproduce the above behavior

import torch
from diffusers import UNet2DConditionModel
from diffusers.loaders import LoraLoaderMixin

model = UNet2DConditionModel.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    subfolder="unet",
    use_safetensors="True",
    variant="fp16",
    torch_dtype=torch.float16
).to("cuda")
print(model) # original model structure

lora_loader = LoraLoaderMixin()
lora_state_dict, network_alphas = lora_loader.lora_state_dict("ostris/crayon_style_lora_sdxl")
model.load_attn_procs(lora_state_dict)
model.fuse_lora(lora_scale=0.6)

[asfiyab-nvidia]

@sayakpaul @yiyixuxu any input on this?

[sayakpaul]

You need to call unload_lora_weights() to complete unload them. Detailed guide: https://huggingface.co/docs/diffusers/main/en/using-diffusers/merge_loras