Update app.py

app.py

@@ -18,12 +18,12 @@ 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
 
+
 def initialize_gpu():
-    """Ensure that ZeroGPU assigns a GPU before using CUDA"""
+    """Ensure ZeroGPU assigns a GPU before initializing CUDA"""
     global device
     try:
-        with spaces.GPU():  # Ensures GPU allocation
-            torch.cuda.init()
+        with spaces.GPU():  # Ensures GPU is allocated
             if torch.cuda.is_available():
                 device = torch.device("cuda")
                 print(f"✅ GPU initialized: {torch.cuda.get_device_name(0)}")
@@ -35,6 +35,9 @@ 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()
 
@@ -47,16 +50,16 @@ 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 before running
+@spaces.GPU  # Ensures ZeroGPU assigns a GPU
 def process_image(image):
-    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
+    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
@@ -68,21 +71,20 @@ def test_process_img(image):
 
 @spaces.GPU
 def process_video(vid_path=None):
-    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)
-
+    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
 
 
+
 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