Does load_attn_procs work with safetensors?
Hi,I’m trying to load a model and checkpoint with 0.14.0 version of the Diffusers (safetensorsalso installed) and am not able to get it working. The code looks like this:pipeline = pipeline_type.from_pretrained(
        "Lykon/DreamShaper",
        revision="main",
        torch_dtype=torch.float16,
    )

    pipeline.unet.load_attn_procs("Lykon/DreamShaper", revision="main", weight_name="DreamShaper_4BakedVae_fp16.safetensors", use_safetensors=True)It downloads the file but then errors withinvalid load key, '\x9f'.which kind of implies that diffusers thinks that it is a pickle file.I am not having much luck fining working examples ofload_attn_procsthat work with anything other than apytorch.binfile so maybe it just not there yet.If anyone has gotten this to work, I would love to know it!Don
2023-03-29T03:06:05Z
I’m running into the same issue.It appears that the KeyError is caused by an assumption in the code that the keys with ‘.lora_down.weight’ and ‘.lora_up.weight’ will always be present for each block in thelora_grouped_dict. However, it seems that this may not be the case.To fix this issue, try adding a check to ensure that the keys are present before accessing them. Modify thenew_load_attn_procsfunction like this:def new_load_attn_procs(pipe, tensors):
    lora_grouped_dict = defaultdict(dict)

    for key, value in tensors.items():
        if key.endswith(".lora_down.weight") or key.endswith(".lora_up.weight"):
            components = key.split("_")
            block = "_".join(components[:-1])
            lora_grouped_dict[block][components[-2] + "_" + components[-1]] = value

    for key, value_dict in lora_grouped_dict.items():
        if "to_k.lora_down.weight" not in value_dict or "to_k.lora_up.weight" not in value_dict:
            continue

        rank = value_dict["to_k.lora_down.weight"].shape[0]
        cross_attention_dim = value_dict["to_k.lora_down.weight"].shape[1]
        hidden_size = value_dict["to_k.lora_up.weight"].shape[0]

        # Assuming the `pipe.unet` has a method to get attention processor by its key
        attn_processor = pipe.unet.get_attention_processor_by_key(key)
        attn_processor.lora_rank = rank
        attn_processor.cross_attention_dim = cross_attention_dim
        attn_processor.hidden_size = hidden_size
        attn_processor.to_k.lora_down.load_state_dict({"weight": value_dict["to_k.lora_down.weight"]})
        attn_processor.to_k.lora_up.load_state_dict({"weight": value_dict["to_k.lora_up.weight"]})

# Replace the original call to `pipe.unet.load_attn_procs(tensors)` with the new function
new_load_attn_procs(pipe, tensors)Let me know how it goes.
2023-03-29T09:38:55Z
Thanks for reporting this! I just openeda GitHub issueto track this down.
2023-03-31T04:00:07Z
Thanks@slavakurilyak! Appreciate your input.Trying to understand this better – could you help me with a scenario of why the key mismatch would be caused?