Update app.py

app.py

@@ -17,25 +17,19 @@ import torch
 import subprocess
 import os
 
-# Check if CUDA is available in PyTorch
-print(f"🚀 PyTorch CUDA Available: {torch.cuda.is_available()}")
-
-# Check available CUDA devices
-if torch.cuda.is_available():
-    print(f"✅ GPU detected: {torch.cuda.get_device_name(0)}")
-else:
-    print("❌ No GPU detected by PyTorch!")
-
-# Check `nvidia-smi` output (if NVIDIA GPU is assigned)
 try:
-    gpu_check = subprocess.run(["nvidia-smi"], capture_output=True, text=True)
-    print(f"🖥️ nvidia-smi output:\n{gpu_check.stdout}")
-except FileNotFoundError:
-    print("❌ nvidia-smi not found! GPU might not be assigned.")
+    import spaces  # Required for ZeroGPU
 
-# Check if Hugging Face assigned a GPU
-hf_gpu_env = os.getenv("CUDA_VISIBLE_DEVICES")
-print(f"🔍 Hugging Face CUDA_VISIBLE_DEVICES: {hf_gpu_env}")
+    if torch.cuda.is_available():
+        print(f"✅ GPU detected: {torch.cuda.get_device_name(0)}")
+        device = torch.device("cuda")
+    else:
+        print("❌ No GPU detected initially. ZeroGPU should allocate one dynamically.")
+        device = torch.device("cpu")  # Default to CPU until ZeroGPU assigns GPU
+
+except Exception as e:
+    print(f"🚨 Error checking CUDA: {e}")
+    device = torch.device("cpu")
 
 
 
@@ -45,6 +39,7 @@ CANCEL_PROCESSING = False
 img_seg = ImageSegmenter(model_type="yolov8s-seg")
 depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256")
 
+@spaces.GPU  # Ensures GPU is allocated when this function runs
 def process_image(image):
     image = utils.resize(image)
     image_segmentation, objects_data = img_seg.predict(image)
@@ -54,12 +49,14 @@ def process_image(image):
     plot_fig = display_pcd(objs_pcd)
     return image_segmentation, depth_colormap, dist_image, plot_fig
 
+@spaces.GPU  # Requests GPU for depth estimation
 def test_process_img(image):
     image = utils.resize(image)
     image_segmentation, objects_data = img_seg.predict(image)
     depthmap, depth_colormap = depth_estimator.make_prediction(image)
     return image_segmentation, objects_data, depthmap, depth_colormap
 
+@spaces.GPU  # Ensures GPU is used for video processing
 def process_video(vid_path=None):
     vid_cap = cv2.VideoCapture(vid_path)
     while vid_cap.isOpened():
@@ -74,6 +71,7 @@ def process_video(vid_path=None):
     
     return None
 
+
 def update_segmentation_options(options):
     img_seg.is_show_bounding_boxes = True if 'Show Boundary Box' in options else False
     img_seg.is_show_segmentation = True if 'Show Segmentation Region' in options else False
@@ -82,7 +80,9 @@ def update_segmentation_options(options):
 def update_confidence_threshold(thres_val):
     img_seg.confidence_threshold = thres_val/100
 
+@spaces.GPU  # Ensures YOLO + MiDaS get GPU access
 def model_selector(model_type):
+    global img_seg, depth_estimator
 
     if "Small - Better performance and less accuracy" == model_type:
         midas_model, yolo_model = "midas_v21_small_256", "yolov8s-seg"