Update app.py

app.py

@@ -13,23 +13,31 @@ from point_cloud_generator import display_pcd
 
 
 
-import torch
-import subprocess
-import os
+import spaces  # Required for ZeroGPU
 
-try:
-    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
 
-    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
+def initialize_gpu():
+    """Ensure that ZeroGPU assigns a GPU before using CUDA"""
+    global device
+    try:
+        with spaces.GPU():  # Ensures GPU allocation
+            torch.cuda.init()
+            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")
+    except Exception as e:
+        print(f"🚨 GPU initialization failed: {e}")
+        device = torch.device("cpu")
+
+# Run GPU initialization before using CUDA
+initialize_gpu()
 
-except Exception as e:
-    print(f"🚨 Error checking CUDA: {e}")
-    device = torch.device("cpu")
 
 
 
@@ -39,15 +47,17 @@ 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
+@spaces.GPU  # Ensures GPU is allocated before running
 def process_image(image):
-    image = utils.resize(image)
-    image_segmentation, objects_data = img_seg.predict(image)
-    depthmap, depth_colormap = depth_estimator.make_prediction(image)
-    dist_image = utils.draw_depth_info(image, depthmap, objects_data)
-    objs_pcd = utils.generate_obj_pcd(depthmap, objects_data)
-    plot_fig = display_pcd(objs_pcd)
-    return image_segmentation, depth_colormap, dist_image, plot_fig
+    with spaces.GPU():  # Explicitly allocate a GPU
+        image = utils.resize(image)
+        image_segmentation, objects_data = img_seg.predict(image)
+        depthmap, depth_colormap = depth_estimator.make_prediction(image)
+        dist_image = utils.draw_depth_info(image, depthmap, objects_data)
+        objs_pcd = utils.generate_obj_pcd(depthmap, objects_data)
+        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):
@@ -56,19 +66,20 @@ def test_process_img(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
+@spaces.GPU
 def process_video(vid_path=None):
-    vid_cap = cv2.VideoCapture(vid_path)
-    while vid_cap.isOpened():
-        ret, frame = vid_cap.read()        
-        if ret:
-            print("making predictions ....")
-            frame = utils.resize(frame)
-            image_segmentation, objects_data = img_seg.predict(frame)
-            depthmap, depth_colormap = depth_estimator.make_prediction(frame)
-            dist_image = utils.draw_depth_info(frame, depthmap, objects_data)
-            yield cv2.cvtColor(image_segmentation, cv2.COLOR_BGR2RGB), depth_colormap, cv2.cvtColor(dist_image, cv2.COLOR_BGR2RGB)
-    
+    with spaces.GPU():
+        vid_cap = cv2.VideoCapture(vid_path)
+        while vid_cap.isOpened():
+            ret, frame = vid_cap.read()
+            if ret:
+                print("making predictions ....")
+                frame = utils.resize(frame)
+                image_segmentation, objects_data = img_seg.predict(frame)
+                depthmap, depth_colormap = depth_estimator.make_prediction(frame)
+                dist_image = utils.draw_depth_info(frame, depthmap, objects_data)
+                yield cv2.cvtColor(image_segmentation, cv2.COLOR_BGR2RGB), depth_colormap, cv2.cvtColor(dist_image, cv2.COLOR_BGR2RGB)
+
     return None