Face-Landmark-ControlNet / prune_ckpt.py
georgefen's picture
init commit for huggingface
45b4aa7
raw
history blame
395 Bytes
import torch
import argparse
def prune_ckpt(ckpt_path, save_path):
raw = torch.load(ckpt_path, map_location=torch.device('cpu'))
state_dict = raw["state_dict"]
torch.save(state_dict, save_path)
if __name__ == '__main__':
args = argparse.ArgumentParser()
args.add_argument('--ckpt_path', type=str)
args.add_argument('--save_path', type=str)
args = args.parse_args()