Update app.py

app.py

@@ -2,6 +2,7 @@ from ultralytics import YOLO
 import cv2
 import gradio as gr
 import numpy as np
+import spaces
 import os
 import torch
 import utils
@@ -13,21 +14,20 @@ from point_cloud_generator import display_pcd
 
 
 
-import spaces  # Required for ZeroGPU
 
-# Ensure CUDA is NOT initialized before ZeroGPU assigns a device
-torch.backends.cudnn.enabled = False  # Prevents CUDA errors on first GPU allocation
 
 
+device = torch.device("cpu")  # Start in CPU mode
+
 def initialize_gpu():
     """Ensure ZeroGPU assigns a GPU before initializing CUDA"""
     global device
     try:
-        with spaces.GPU():  # Ensures GPU is allocated
+        with spaces.GPU():  # Ensures ZeroGPU assigns a GPU
+            torch.cuda.empty_cache()  # Prevent leftover memory issues
             if torch.cuda.is_available():
                 device = torch.device("cuda")
                 print(f"✅ GPU initialized: {torch.cuda.get_device_name(0)}")
-                torch.cuda.empty_cache()  # Clear memory
             else:
                 print("❌ No GPU detected after ZeroGPU allocation.")
                 device = torch.device("cpu")
@@ -35,8 +35,6 @@ def initialize_gpu():
         print(f"🚨 GPU initialization failed: {e}")
         device = torch.device("cpu")
 
-# ✅ Don't call CUDA here!
-device = torch.device("cpu")  # Default to CPU
 
 # Run GPU initialization before using CUDA
 initialize_gpu()