File size: 1,203 Bytes
9d71531 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import argparse
from pathlib import Path
import torch
from safetensors.torch import save_file
def convert(path: Path):
state_dict = torch.load(path, map_location="cpu")
if "state_dict" in state_dict:
state_dict = state_dict["state_dict"]
to_remove = []
for k, v in state_dict.items():
if not isinstance(v, torch.Tensor):
to_remove.append(k)
for k in to_remove:
del state_dict[k]
output_path = path.with_suffix(".safetensors").as_posix()
save_file(state_dict, output_path)
def main(path: str):
path_ = Path(path).resolve()
if not path_.exists():
raise ValueError(f"Invalid path: {path}")
if path_.is_file():
to_convert = [path_]
else:
to_convert = list(path_.glob("*.ckpt"))
for file in to_convert:
if file.with_suffix(".safetensors").exists():
continue
print(f"Converting... {file}")
convert(file)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("path", type=str, help="Path to checkpoint file or directory.")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
main(args.path)
|