import torch
from transformers import AutoImageProcessor, Swinv2ForImageClassification
import onnx
import onnxruntime as ort

# Load the model and processor
image_processor = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
model = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")

# Set the model to evaluation mode
model.eval()

# Create dummy input for tracing
dummy_input = torch.randn(1, 3, 256, 256)  # Batch size of 1, 3 color channels, 256x256 image

# Export the model to ONNX
onnx_model_path = "model.onnx"
# torch.onnx.export(
#     model,
#     dummy_input,
#     onnx_model_path,
#     input_names=["pixel_values"],
#     output_names=["logits"],
#     opset_version=11,
#     dynamic_axes={
#         "pixel_values": {0: "batch_size"},
#         "logits": {0: "batch_size"}
#     }
# )

# Verify the ONNX model
# onnx_model = onnx.load(onnx_model_path)
# onnx.checker.check_model(onnx_model)
print("The ONNX model is valid.")