import torch
from safetensors.torch import save_file

input_file = "AnimateDiffV3MotionModule____training_NOISEWARP_72x40_Test_Sep14_FromV3_RandomlyDegraded-2024-09-19T12-37-23____checkpoint_latest.ckpt"

model = torch.load(input_file, weights_only=True)

mm_sd = {}
for key in model["state_dict"].keys():
    if "motion_modules" in key:
        new_key = key.replace("module.", "")
        mm_sd[new_key] = model["state_dict"][key]

save_file(mm_sd, "motion_model.safetensors")

unet_sd = {}
for key in model["state_dict"].keys():
    if "motion_modules" not in key:
        new_key = key.replace("module.", "")
        unet_sd[new_key] = model["state_dict"][key]

save_file(unet_sd, "unet.safetensors")