import onnx
import argparse

def modify_batch_size(input_model_path, output_model_path, batch_size):
    # Load the ONNX model
    model = onnx.load(input_model_path)

    # Modify the batch size of the first input tensor
    input_tensor = model.graph.input[0]
    input_tensor.type.tensor_type.shape.dim[0].dim_value = batch_size

    # Save the modified model
    onnx.save(model, output_model_path)
    print(f"Modified model saved to {output_model_path}")

if __name__ == "__main__":
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Modify the batch size of the first input tensor of an ONNX model.")
    parser.add_argument("input_model_path", type=str, help="Path to the input ONNX model file.")
    parser.add_argument("output_model_path", type=str, help="Path to save the modified ONNX model file.")
    parser.add_argument("--batch_size", type=int, help="Desired batch size for the first input tensor.")

    args = parser.parse_args()

    # Ensure that the batch_size argument has been provided
    if args.batch_size is None:
        raise ValueError("Please specify the new batch size for the input tensor using --batch_size.")

    # Call the function to modify the model's batch size
    modify_batch_size(args.input_model_path, args.output_model_path, args.batch_size)