Loading .safetensors Lora

[adhikjoshi]

Describe the bug

I have downloaded lora from civitai which is in .safetensor format.

When i load it using below code,

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe.unet.load_attn_procs("lora.safetensors")

It throws error : KeyError: 'to_k_lora.down.weight'

File "/workspace/server/tasks.py", line 346, in txt2img
self.pipe.unet.load_attn_procs(embd, use_safetensors=True)
File "/opt/conda/envs/ldm/lib/python3.8/site-packages/diffusers/loaders.py", line 224, in load_attn_procs
rank = value_dict["to_k_lora.down.weight"].shape[0]
KeyError: 'to_k_lora.down.weight'

Reproduction

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe.unet.load_attn_procs("lora.safetensors")

Logs

No response

System Info

Diffusers Version: 0.15.0.dev0

[patrickvonplaten]

Hey @adhikjoshi,

Thanks for the issue we should indeed try to support also A1111 loading of LoRA tensors soon. cc @sayakpaul here

[alejobrainz]

Kohya-ss/sd-scripts has a nice mechanism for it, but it broke with 0.15, but you can for sure load A1111 LoRA Tensors with the function below on 0.14.0:

def apply_lora(pipe, lora_path, weight:float = 1.0):
    from safetensors.torch import load_file
    from sd-scripts.networks.lora import create_network_from_weights
    import torch
    
    vae = pipe.vae
    text_encoder = pipe.text_encoder
    unet = pipe.unet

    sd = load_file(lora_path)
    lora_network, sd = create_network_from_weights(weight, None, vae, text_encoder, unet, sd)
    lora_network.apply_to(text_encoder, unet)
    lora_network.load_state_dict(sd)
    lora_network.to("cuda", dtype=torch.float16)

but as of 0.15 it fails:

assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
AssertionError: duplicated lora name: lora_unet_down_blocks_0_attentions_0_transformer_blocks_0_attn1_to_q

[adhikjoshi]

Kohya-ss/sd-scripts has a nice mechanism for it, but it broke with 0.15, but you can for sure load A1111 LoRA Tensors with the function below on 0.14.0:

def apply_lora(pipe, lora_path, weight:float = 1.0):

    from safetensors.torch import load_file

    from sd-scripts.networks.lora import create_network_from_weights

    import torch

    

    vae = pipe.vae

    text_encoder = pipe.text_encoder

    unet = pipe.unet



    sd = load_file(lora_path)

    lora_network, sd = create_network_from_weights(weight, None, vae, text_encoder, unet, sd)

    lora_network.apply_to(text_encoder, unet)

    lora_network.load_state_dict(sd)

    lora_network.to("cuda", dtype=torch.float16)

but as of 0.15 it fails:

assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"

AssertionError: duplicated lora name: lora_unet_down_blocks_0_attentions_0_transformer_blocks_0_attn1_to_q

CC @haofanwang @sayakpaul

[sayakpaul]

Can someone provide LoRA file in the A1111 format? Providing as many relevant details associated to the file as possible would be great too.

[adhikjoshi]

Can someone provide LoRA file in the A1111 format? Providing as many relevant details associated to the file as possible would be great too.

I have downloaded offset noise trained lora and uploaded its .safetensor on huggingface

https://huggingface.co/adhikjoshi/epi_noiseoffset

[alejobrainz]

@sayakpaul here you go. This Lora was trained using ss-kohya's scripts and works fine in A1111. I can load it on diffusers 0.14.0 with the snippet above using the lora.py from sd-scripts:

caAos-000001.zip

Thanks,

Alejandro.

[sayakpaul]

Cc: @patrickvonplaten ^

[adhikjoshi]

Here is function i made from convert_lora_safetensor_to_diffusers.py to load lora on inference time.

import torch
from safetensors.torch import load_file

def load_lora_weights(pipeline, checkpoint_path):
    # load base model
    pipeline.to("cuda")
    LORA_PREFIX_UNET = "lora_unet"
    LORA_PREFIX_TEXT_ENCODER = "lora_te"
    alpha = 0.75
    # load LoRA weight from .safetensors
    state_dict = load_file(checkpoint_path, device="cuda")
    visited = []

    # directly update weight in diffusers model
    for key in state_dict:
        # it is suggested to print out the key, it usually will be something like below
        # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"

        # as we have set the alpha beforehand, so just skip
        if ".alpha" in key or key in visited:
            continue

        if "text" in key:
            layer_infos = key.split(".")[0].split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_")
            curr_layer = pipeline.text_encoder
        else:
            layer_infos = key.split(".")[0].split(LORA_PREFIX_UNET + "_")[-1].split("_")
            curr_layer = pipeline.unet

        # find the target layer
        temp_name = layer_infos.pop(0)
        while len(layer_infos) > -1:
            try:
                curr_layer = curr_layer.__getattr__(temp_name)
                if len(layer_infos) > 0:
                    temp_name = layer_infos.pop(0)
                elif len(layer_infos) == 0:
                    break
            except Exception:
                if len(temp_name) > 0:
                    temp_name += "_" + layer_infos.pop(0)
                else:
                    temp_name = layer_infos.pop(0)

        pair_keys = []
        if "lora_down" in key:
            pair_keys.append(key.replace("lora_down", "lora_up"))
            pair_keys.append(key)
        else:
            pair_keys.append(key)
            pair_keys.append(key.replace("lora_up", "lora_down"))

        # update weight
        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)
        else:
            weight_up = state_dict[pair_keys[0]].to(torch.float32)
            weight_down = state_dict[pair_keys[1]].to(torch.float32)
            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)

        # update visited list
        for item in pair_keys:
            visited.append(item)

    return pipeline

Can use it like,

lora_model = lora_models + "/" + opt.lora + ".safetensors"
self.pipe = load_lora_weights(self.pipe, lora_model)

@sayakpaul @patrickvonplaten

[sayakpaul]

Thanks! Do you have the checkpoints with which we could test this?

[adhikjoshi]

Can someone provide LoRA file in the A1111 format? Providing as many relevant details associated to the file as possible would be great too.

I have downloaded offset noise trained lora and uploaded its .safetensor on huggingface

https://huggingface.co/adhikjoshi/epi_noiseoffset

This uploaded safetensor lora and others work well