Skip to content

Commit

Permalink
convert safetensors
Browse files Browse the repository at this point in the history
convert safetensors
  • Loading branch information
cgisky1980 committed Aug 25, 2023
1 parent 3120154 commit aff4b57
Showing 1 changed file with 44 additions and 0 deletions.
44 changes: 44 additions & 0 deletions convert_safetensors.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
import json
import os
import sys
import copy
import torch
from safetensors.torch import load_file, save_file

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, help='Path to input pth model')
parser.add_argument('--output', type=str, default='./converted.st',
help='Path to output safetensors model')
args = parser.parse_args()


def convert_file(
pt_filename: str,
sf_filename: str,
):
loaded = torch.load(pt_filename, map_location="cpu")
if "state_dict" in loaded:
loaded = loaded["state_dict"]

loaded = {k: v.clone().half() for k, v in loaded.items()}
for k, v in loaded.items():
print(f'{k}\t{v.shape}\t{v.dtype}')

# For tensors to be contiguous
loaded = {k: v.contiguous() for k, v in loaded.items()}

dirname = os.path.dirname(sf_filename)
os.makedirs(dirname, exist_ok=True)
save_file(loaded, sf_filename, metadata={"format": "pt"})
reloaded = load_file(sf_filename)
for k in loaded:
pt_tensor = loaded[k]
sf_tensor = reloaded[k]
if not torch.equal(pt_tensor, sf_tensor):
raise RuntimeError(f"The output tensors do not match for key {k}")


if __name__ == "__main__":
convert_file(args.input, args.output)
print(f"Saved to {args.output}")

0 comments on commit aff4b57

Please sign in to comment.