spacepxl

/

Wan2.1-control-loras

like 143

Hi！I have recently been trying to reproduce your work and attempted to implement it using Python code. While loading the LoRA model with diffusers, I encountered an error. After using the following script for conversion, the same error still occurred. Could you please advise on what might be the issue?
script：
'''
import torch
from safetensors.torch import load_file, save_file
from typing import Dict, Any
import os

TRANSFORMER_KEYS_RENAME_DICT = {
"time_embedding.0": "condition_embedder.time_embedder.linear_1",
"time_embedding.2": "condition_embedder.time_embedder.linear_2",
"text_embedding.0": "condition_embedder.text_embedder.linear_1",
"text_embedding.2": "condition_embedder.text_embedder.linear_2",
"time_projection.1": "condition_embedder.time_proj",
"head.modulation": "scale_shift_table",
"head.head": "proj_out",
"modulation": "scale_shift_table",
"ffn.0": "ffn.net.0.proj",
"ffn.2": "ffn.net.2",
"norm2": "norm__placeholder",
"norm3": "norm2",
"norm__placeholder": "norm3",
"img_emb.proj.0": "condition_embedder.image_embedder.norm1",
"img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj",
"img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2",
"img_emb.proj.4": "condition_embedder.image_embedder.norm2",
"img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed",
"self_attn.q": "attn1.to_q",
"self_attn.k": "attn1.to_k",
"self_attn.v": "attn1.to_v",
"self_attn.o": "attn1.to_out.0",
"self_attn.norm_q": "attn1.norm_q",
"self_attn.norm_k": "attn1.norm_k",
"cross_attn.q": "attn2.to_q",
"cross_attn.k": "attn2.to_k",
"cross_attn.v": "attn2.to_v",
"cross_attn.o": "attn2.to_out.0",
"cross_attn.norm_q": "attn2.norm_q",
"cross_attn.norm_k": "attn2.norm_k",
"attn2.to_k_img": "attn2.add_k_proj",
"attn2.to_v_img": "attn2.add_v_proj",
"attn2.norm_k_img": "attn2.norm_added_k",
}

def convert_lora_keys(state_dict: Dict[str, Any]) -> Dict[str, Any]:

new_state_dict = {}
for key, value in state_dict.items():
    new_key = key
    for old_pattern, new_pattern in TRANSFORMER_KEYS_RENAME_DICT.items():
        new_key = new_key.replace(old_pattern, new_pattern)
    

    if "lora" in new_key.lower():

        if not any(x in new_key for x in ["lora_down", "lora_up", "alpha"]):
            new_key = new_key.replace("weight", "lora_weight")  
    new_state_dict[new_key] = value
return new_state_dict

def process_lora_file(input_path: str, output_dir: str):

if input_path.endswith(".safetensors"):
    state_dict = load_file(input_path)
else:
    state_dict = torch.load(input_path)


converted_state_dict = convert_lora_keys(state_dict)


os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, os.path.basename(input_path))

if input_path.endswith(".safetensors"):
    save_file(converted_state_dict, output_path)
else:
    torch.save(converted_state_dict, output_path)

Errors：
'''
Traceback (most recent call last):
File "/home/shaochengshen/Wan2.1/control-lora.py", line 18, in
pipe.load_lora_weights(
File "/root/miniconda3/envs/Wan/lib/python3.10/site-packages/diffusers/loaders/lora_pipeline.py", line 4869, in load_lora_weights
state_dict = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
File "/root/miniconda3/envs/Wan/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
return fn(*args, **kwargs)
File "/root/miniconda3/envs/Wan/lib/python3.10/site-packages/diffusers/loaders/lora_pipeline.py", line 4796, in lora_state_dict
state_dict = _convert_non_diffusers_wan_lora_to_diffusers(state_dict)
File "/root/miniconda3/envs/Wan/lib/python3.10/site-packages/diffusers/loaders/lora_conversion_utils.py", line 1564, in _convert_non_diffusers_wan_lora_to_diffusers
num_blocks = len({k.split("blocks.")[1].split(".")[0] for k in original_state_dict})
File "/root/miniconda3/envs/Wan/lib/python3.10/site-packages/diffusers/loaders/lora_conversion_utils.py", line 1564, in
num_blocks = len({k.split("blocks.")[1].split(".")[0] for k in original_state_dict})
IndexError: list index out of range
'''