Upload 2 files

app.py

@@ -1,11 +1,39 @@
+# for Zero GPU Spaces compatibility
+import spaces
+@spaces.GPU
+def dummy_gpu():
+    pass
+
 import gradio as gr
 import numpy as np
 import cv2
 import torch
 import onnxruntime as ort
+from optimum.onnxruntime import ORTModel
 from ultralytics import YOLO
 import os
 from typing import Tuple, List
+import subprocess
+
+def install_cuda_toolkit():
+    print("Installing CUDA Toolkit.") 
+    #CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
+    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
+    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
+    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
+    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
+    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
+
+    os.environ["CUDA_HOME"] = "/usr/local/cuda"
+    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
+    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
+        os.environ["CUDA_HOME"],
+        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
+    )
+    # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
+    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
+
+install_cuda_toolkit()
 
 # Configuration - UPDATE THESE VALUES
 MODEL_PT_PATH = "model.pt"    # Your trained PyTorch model
@@ -14,9 +42,14 @@ INPUT_SIZE = 640              # Must match training size
 CLASS_NAMES = ["class0", "class1"]  # Your actual class names
 CONF_THRESHOLD = 0.5          # Confidence threshold
 IOU_THRESHOLD = 0.45          # NMS IoU threshold
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+session_options = ort.SessionOptions()
+session_options.log_severity_level = 0
 
 def convert_pt_to_onnx():
     """Convert PyTorch model to ONNX format if not exists"""
+    print(f'Converting model on {"cuda" if torch.cuda.is_available() else "cpu"}')
     if not os.path.exists(MODEL_ONNX_PATH):
         print("Converting PyTorch model to ONNX...")
         try:
@@ -43,15 +76,19 @@ def convert_pt_to_onnx():
 
 def load_onnx_model() -> ort.InferenceSession:
     """Initialize ONNX runtime session"""
-    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+    print(f'Loading model on {"cuda" if torch.cuda.is_available() else "cpu"}')
+    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if DEVICE != "cpu" else ['CPUExecutionProvider']
     try:
-        return ort.InferenceSession(MODEL_ONNX_PATH, providers=providers)
+        #return ort.InferenceSession(MODEL_ONNX_PATH, providers=providers, session_options=session_options, export=True)
+        return ORTModel.load_model(MODEL_ONNX_PATH, provider='CUDAExecutionProvider' if DEVICE != "cpu" else 'CPUExecutionProvider', session_options=session_options)
     except Exception as e:
         raise RuntimeError(f"Failed to load ONNX model: {str(e)}")
 
 # Initialize model
 convert_pt_to_onnx()
 ort_session = load_onnx_model()
+print("Available Providers: ", ort_session._providers)
+#assert "CUDAExecutionProvider" in ort_session._providers
 
 def letterbox_image(image: np.ndarray) -> Tuple[np.ndarray, float, Tuple[int, int]]:
     """

requirements.txt

@@ -1,7 +1,10 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+torch
 gradio
 numpy
-onnxruntime-gpu
+onnxslim
+#onnxruntime-gpu==1.19.2
+optimum[onnxruntime-gpu]
 opencv-python
 Pillow
-torch
 ultralytics