Sumith0102/model / convert_checkpoint.py
Sumith0102's picture
download
raw
1.45 kB
"""Convert the owner's checkpoint without enabling unrestricted pickle loading."""
import argparse
from pathlib import Path
def convert(checkpoint, output):
import torch
from safetensors.torch import save_file
import os
import tempfile
checkpoint, output = Path(checkpoint), Path(output)
if output.exists():
raise ValueError("Output already exists; choose another path.")
obj = torch.load(checkpoint, map_location="cpu", weights_only=True)
state = obj.get("state_dict", obj) if isinstance(obj, dict) else None
if not isinstance(state, dict) or not state:
raise ValueError("No model tensors were found.")
if any(not isinstance(k, str) or not isinstance(v, torch.Tensor) for k, v in state.items()):
raise ValueError("Unexpected checkpoint content. Inspect it before conversion.")
fd, temp_name = tempfile.mkstemp(dir=output.parent, suffix=".safetensors")
os.close(fd)
try:
save_file({k: v.detach().contiguous().clone() for k, v in state.items()}, temp_name)
os.replace(temp_name, output)
finally:
Path(temp_name).unlink(missing_ok=True)
return output
def main():
parser = argparse.ArgumentParser()
parser.add_argument("checkpoint", type=Path)
parser.add_argument("output", type=Path)
args = parser.parse_args()
convert(args.checkpoint, args.output)
print("Saved", args.output)
if __name__ == "__main__":
main()

Xet Storage Details

Size:
1.45 kB
·
Xet hash:
5d729893eb795b96625e6452b5e6e2e0db18a27197c1141095410b78aca5010e

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.