Pytorch Model to Safetensors converter

pt_to_safe.py

import torch
from safetensors.torch import save_file
import argparse
import os

def convert_pt_to_safetensors(model_path, safe_path):
    model = torch.load(model_path, map_location="cpu", weights_only=False)
    metadata = {"format":"pt"}
    save_file(model, safe_path, metadata)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Converts a PyTorch model to Safetensors.")
    parser.add_argument("model_path", type=str, help="Path to the PyTorch model file")


    args = parser.parse_args()
    model_path = args.model_path
    file_path = os.path.splitext(model_path)[0]
    safe_path = file_path + ".safetensors"
    convert_pt_to_safetensors(model_path, safe_path)