diffusers/scripts/convert_animatediff_motion_lora_to_diffusers.py at main · zsc/diffusers · GitHub

History

69 lines (52 loc) · 2.1 KB

Raw

import argparse

import os

import torch

from huggingface_hub import create_repo, upload_folder

from safetensors.torch import load_file, save_file

def convert_motion_module(original_state_dict):

converted_state_dict = {}

for k, v in original_state_dict.items():

if "pos_encoder" in k:

continue

else:

converted_state_dict[

k.replace(".norms.0", ".norm1")

.replace(".norms.1", ".norm2")

.replace(".ff_norm", ".norm3")

.replace(".attention_blocks.0", ".attn1")

.replace(".attention_blocks.1", ".attn2")

.replace(".temporal_transformer", "")

] = v

return converted_state_dict

def get_args():

parser = argparse.ArgumentParser()

parser.add_argument("--ckpt_path", type=str, required=True, help="Path to checkpoint")

parser.add_argument("--output_path", type=str, required=True, help="Path to output directory")

parser.add_argument(

"--push_to_hub",

action="store_true",

default=False,

help="Whether to push the converted model to the HF or not",

)

return parser.parse_args()

if __name__ == "__main__":

args = get_args()

if args.ckpt_path.endswith(".safetensors"):

state_dict = load_file(args.ckpt_path)

else:

state_dict = torch.load(args.ckpt_path, map_location="cpu")

if "state_dict" in state_dict.keys():

state_dict = state_dict["state_dict"]

conv_state_dict = convert_motion_module(state_dict)

# convert to new format

output_dict = {}

for module_name, params in conv_state_dict.items():

if type(params) is not torch.Tensor:

continue

output_dict.update({f"unet.{module_name}": params})

os.makedirs(args.output_path, exist_ok=True)

filepath = os.path.join(args.output_path, "diffusion_pytorch_model.safetensors")

save_file(output_dict, filepath)

if args.push_to_hub:

repo_id = create_repo(args.output_path, exist_ok=True).repo_id

upload_folder(repo_id=repo_id, folder_path=args.output_path, repo_type="model")

diffusers/scripts/convert_animatediff_motion_lora_to_diffusers.py at main · zsc/diffusers · GitHub

Expand file tree

Breadcrumbs

convert_animatediff_motion_lora_to_diffusers.py

Latest commit

History

Breadcrumbs

convert_animatediff_motion_lora_to_diffusers.py

File metadata and controls