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