Convert LoRA safetensor to diffusers form

[smko77]

Describe the bug

The error occured when converting the custom LoRA weight in safetensor to diffuesrs form. I used this script

I was converting the LoRA weight from civitai. I tried debugging to find out why this error occured. It turns out that this LoRA weight has additional structure such as

# lora unet down block 0
lora_unet_down_blocks_0_downsamplers_0_conv.alpha	<class 'torch.Tensor'>
lora_unet_down_blocks_0_downsamplers_0_conv.lora_down.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_downsamplers_0_conv.lora_up.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_0_conv1.alpha	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_0_conv1.lora_down.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_0_conv1.lora_up.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_0_conv2.alpha	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_0_conv2.lora_down.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_0_conv2.lora_up.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_0_time_emb_proj.alpha	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_0_time_emb_proj.lora_down.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_0_time_emb_proj.lora_up.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_1_conv1.alpha	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_1_conv1.lora_down.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_1_conv1.lora_up.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_1_conv2.alpha	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_1_conv2.lora_down.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_1_conv2.lora_up.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_1_time_emb_proj.alpha	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_1_time_emb_proj.lora_down.weight	<class 'torch.Tensor'>
lora_unet_down_blocks_0_resnets_1_time_emb_proj.lora_up.weight	<class 'torch.Tensor'>

and the error occured exactly when it converts lora_unet_down_blocks_0_downsamplers_0_conv.lora_down.weight.
At this time, each tensor has size

• weight_up : torch.Size([320, 32])
• weight_down : torch.Size([32, 320, 3, 3])

Reproduction

convert Lora weight using this script

Logs

curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
RuntimeError: mat2 must be a matrix

System Info

• diffusers version: 0.15.1
• Platform: Linux-4.15.0-47-generic-x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.7.5
• PyTorch version (GPU?): 1.13.1+cu117 (True)
• Huggingface_hub version: 0.14.0
• Transformers version: 4.25.1
• Accelerate version: 0.13.1
• xFormers version: 0.0.16
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

[hafriedlander]

Diffusers only supports Lora on linear layers of the UNet

I haven't checked but nearly positive those are Conv2D modules, which Diffusers doesn't (currently) support.

(Script could maybe throw a warning when the Lora you're trying to convert includes unsupported models?)

Diffusers team, the way I've handled Lora (and Lycoris) in Diffusers is with an accelerator ModelHook - https://github.com/stablecabal/gyre/blob/test/gyre/pipeline/lora.py#L99. This works much more generically than your current Attention processor subclass method. Not sure if that's something you'd want to adopt or not.

[alexblattner]

@hafriedlander how do you use this with a pipeline? Do you have an example with compel? Can you simulate using this prompt "concept art, matte, masterpiece, 8k, art by lora:mngstle:1, black and white, monochrome" by any chance?

Thank you for producing something this valuable!

[LiaoFJ]

I changed a little from:

    if len(state_dict[pair_keys[0]].shape) == 4:
        weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
        weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
        curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)

to

    if len(state_dict[pair_keys[0]].shape) == 4:
        weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
        weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
        if len(weight_up.shape) == len(weight_down.shape):
            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
        else:
            curr_layer.weight.data += alpha * torch.einsum('a b, b c h w -> a c h w', weight_up, weight_down)

and it may work. Hope it works too in your code.

[sayakpaul]

What @hafriedlander suggested in #3238 (comment) is true:

I haven't checked but nearly positive those are Conv2D modules, which Diffusers doesn't (currently) support.

Diffusers team, the way I've handled Lora (and Lycoris) in Diffusers is with an accelerator ModelHook - https://github.com/stablecabal/gyre/blob/test/gyre/pipeline/lora.py#L99. This works much more generically than your current Attention processor subclass method. Not sure if that's something you'd want to adopt or not.

I think this introduces a hard dependency on accelerate during inference which I am not sure we want. Therefore I am seeking further opinions from @patrickvonplaten.

[patrickvonplaten]

Think adding a new hook next to attention_processor is a bit too much at the moment. Also I don't think we need to support LoRA for Conv layers at the moment

[hafriedlander]

Cool, figured maybe so, hence why no PR. I'm also not completely sure if it works with training - gyre.ai is inference only, although it does work with CLIP guidance which involves a backwards step.

I can probably pull just that hook out as a lib for pure Diffusers users.

[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.