mirror of
https://github.com/THUDM/CogVideo.git
synced 2025-04-05 03:04:56 +08:00
84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
from typing import Any, Dict
|
|
import torch
|
|
import argparse
|
|
from diffusers.loaders.lora_base import LoraBaseMixin
|
|
from diffusers.models.modeling_utils import load_state_dict
|
|
|
|
|
|
def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
|
|
state_dict = saved_dict
|
|
if "model" in saved_dict.keys():
|
|
state_dict = state_dict["model"]
|
|
if "module" in saved_dict.keys():
|
|
state_dict = state_dict["module"]
|
|
if "state_dict" in saved_dict.keys():
|
|
state_dict = state_dict["state_dict"]
|
|
return state_dict
|
|
|
|
LORA_KEYS_RENAME = {
|
|
|
|
'attention.query_key_value.matrix_A.0': 'attn1.to_q.lora_A.weight',
|
|
'attention.query_key_value.matrix_A.1': 'attn1.to_k.lora_A.weight',
|
|
'attention.query_key_value.matrix_A.2': 'attn1.to_v.lora_A.weight',
|
|
'attention.query_key_value.matrix_B.0': 'attn1.to_q.lora_B.weight',
|
|
'attention.query_key_value.matrix_B.1': 'attn1.to_k.lora_B.weight',
|
|
'attention.query_key_value.matrix_B.2': 'attn1.to_v.lora_B.weight',
|
|
'attention.dense.matrix_A.0': 'attn1.to_out.0.lora_A.weight',
|
|
'attention.dense.matrix_B.0': 'attn1.to_out.0.lora_B.weight'
|
|
}
|
|
|
|
|
|
|
|
PREFIX_KEY = "model.diffusion_model."
|
|
SAT_UNIT_KEY = "layers"
|
|
LORA_PREFIX_KEY = "transformer_blocks"
|
|
|
|
|
|
|
|
def export_lora_weight(ckpt_path,lora_save_directory):
|
|
|
|
merge_original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True))
|
|
|
|
|
|
lora_state_dict = {}
|
|
for key in list(merge_original_state_dict.keys()):
|
|
new_key = key[len(PREFIX_KEY) :]
|
|
for special_key, lora_keys in LORA_KEYS_RENAME.items():
|
|
if new_key.endswith(special_key):
|
|
new_key = new_key.replace(special_key, lora_keys)
|
|
new_key = new_key.replace(SAT_UNIT_KEY, LORA_PREFIX_KEY)
|
|
|
|
lora_state_dict[new_key] = merge_original_state_dict[key]
|
|
|
|
|
|
|
|
# final length should be 240
|
|
if len(lora_state_dict) != 240:
|
|
raise ValueError("lora_state_dict length is not 240")
|
|
|
|
lora_state_dict.keys()
|
|
|
|
LoraBaseMixin.write_lora_layers(
|
|
state_dict=lora_state_dict,
|
|
save_directory=lora_save_directory,
|
|
is_main_process=True,
|
|
weight_name=None,
|
|
save_function=None,
|
|
safe_serialization=True
|
|
)
|
|
|
|
|
|
def get_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--sat_pt_path", type=str, required=True, help="Path to original sat transformer checkpoint"
|
|
)
|
|
parser.add_argument("--lora_save_directory", type=str, required=True, help="Path where converted lora should be saved")
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = get_args()
|
|
|
|
export_lora_weight(args.sat_pt_path, args.lora_save_directory)
|