import sys
import os
from onnx2torch import convert
import torch

def onnx_to_torch_converter(dir_name):
    if not os.path.exists(dir_name):
        print(f"Directory {dir_name} does not exist!")
        return

    onnx_model_path = os.path.join(dir_name, "onnx", "model.onnx")

    if not os.path.exists(onnx_model_path):
        print(f"ONNX model at {onnx_model_path} does not exist!")
        return

    torch_model = convert(onnx_model_path)

    torch_model_save_path = os.path.join(dir_name, "pytorch_model.bin")
    torch.save(torch_model.state_dict(), torch_model_save_path)
    print(f"PyTorch model saved at {torch_model_save_path}")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python onnx2torch.py [directory_path]")
    else:
        dir_name = sys.argv[1]
        onnx_to_torch_converter(dir_name)