Update app.py

app.py

@@ -5,6 +5,7 @@ import os
 import utils
 import plotly.graph_objects as go
 import spaces
+import torch
 
 from image_segmenter import ImageSegmenter
 from monocular_depth_estimator import MonocularDepthEstimator
@@ -24,7 +25,21 @@ def initialize_models():
     if depth_estimator is None:
         depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256")
 
-@spaces.GPU
+def safe_gpu_decorator(func):
+    """Custom decorator to handle GPU operations safely"""
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except RuntimeError as e:
+            if "cudaGetDeviceCount" in str(e):
+                print("GPU initialization failed, falling back to CPU")
+                # Set environment variable to force CPU
+                os.environ['CUDA_VISIBLE_DEVICES'] = ''
+                return func(*args, **kwargs)
+            raise
+    return wrapper
+
+@safe_gpu_decorator
 def process_image(image):
     try:
         print("Starting image processing")
@@ -43,7 +58,7 @@ def process_image(image):
         print(traceback.format_exc())
         raise
 
-@spaces.GPU
+@safe_gpu_decorator
 def test_process_img(image):
     initialize_models()
     image = utils.resize(image)
@@ -51,7 +66,7 @@ def test_process_img(image):
     depthmap, depth_colormap = depth_estimator.make_prediction(image)
     return image_segmentation, objects_data, depthmap, depth_colormap
 
-@spaces.GPU
+@safe_gpu_decorator
 def process_video(vid_path=None):
     try:
         initialize_models()
@@ -66,6 +81,7 @@ def process_video(vid_path=None):
                 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.release()
         return None
     except Exception as e:
         print(f"Error in process_video: {str(e)}")
@@ -83,7 +99,7 @@ def update_confidence_threshold(thres_val):
     initialize_models()
     img_seg.confidence_threshold = thres_val/100
 
-@spaces.GPU
+@safe_gpu_decorator
 def model_selector(model_type):
     global img_seg, depth_estimator
     
@@ -104,6 +120,23 @@ def cancel():
     CANCEL_PROCESSING = True
 
 if __name__ == "__main__":
+    # Try to initialize CUDA early to catch any issues
+    try:
+        if torch.cuda.is_available():
+            print("CUDA is available. Using GPU.")
+            # Test CUDA initialization
+            torch.cuda.init()
+            device = torch.device("cuda")
+        else:
+            print("CUDA is not available. Using CPU.")
+            os.environ['CUDA_VISIBLE_DEVICES'] = ''
+            device = torch.device("cpu")
+    except RuntimeError as e:
+        print(f"CUDA initialization failed: {e}")
+        print("Falling back to CPU mode")
+        os.environ['CUDA_VISIBLE_DEVICES'] = ''
+        device = torch.device("cpu")
+
     with gr.Blocks() as my_app:
         # title
         gr.Markdown("<h1><center>Simultaneous Segmentation and Depth Estimation</center></h1>")
@@ -134,7 +167,6 @@ if __name__ == "__main__":
                         dist_img_output = gr.Image(height=300, label="Distance")
                         pcd_img_output = gr.Plot(label="Point Cloud")
             
-            gr.Markdown("## Sample Images")
             gr.Examples(
                 examples=[os.path.join(os.path.dirname(__file__), "assets/images/baggage_claim.jpg"),
                           os.path.join(os.path.dirname(__file__), "assets/images/kitchen_2.png"),
@@ -172,7 +204,6 @@ if __name__ == "__main__":
                     with gr.Row():
                         dist_vid_output = gr.Image(height=300, label="Distance")
             
-            gr.Markdown("## Sample Videos")
             gr.Examples(
                 examples=[os.path.join(os.path.dirname(__file__), "assets/videos/input_video.mp4"),
                           os.path.join(os.path.dirname(__file__), "assets/videos/driving.mp4"),