Create utils/onnx.py

utils/onnx.py

@@ -0,0 +1,34 @@
+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.")