convert_animatediff_motion_lora_to_diffusers.py

importargparse
importos

importtorch
fromhuggingface_hubimportcreate_repo, upload_folder
fromsafetensors.torchimportload_file, save_file


defconvert_motion_module(original_state_dict):
converted_state_dict= {}
fork, vinoriginal_state_dict.items():
if"pos_encoder"ink:
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

returnconverted_state_dict


defget_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",
 )

returnparser.parse_args()


if__name__=="__main__":
args=get_args()

ifargs.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"instate_dict.keys():
state_dict=state_dict["state_dict"]

conv_state_dict=convert_motion_module(state_dict)

# convert to new format
output_dict= {}
formodule_name, paramsinconv_state_dict.items():
iftype(params) isnottorch.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)

ifargs.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")