Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces  # Import spaces first
 import cv2
 import torch
 import numpy as np
@@ -11,6 +12,7 @@ model_configs = {
     'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}
 }
 
+# Initialize model globally
 def initialize_model():
     encoder = 'vitl'
     max_depth = 1
@@ -33,38 +35,49 @@ def initialize_model():
         my_state_dict[new_key] = state_dict[key]
     
     model.load_state_dict(my_state_dict)
-    model.eval()
     return model
 
-# Initialize model globally
 MODEL = initialize_model()
 
+@spaces.GPU
 def process_image(input_image):
     """
     Process the input image and return depth maps
     """
-    # Convert from RGB to BGR (since original code uses cv2.imread which loads in BGR)
     if input_image is None:
         return None, None
-        
-    input_image = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB_BGR)
     
-    # Get depth map
-    depth = MODEL.infer_image(input_image)
+    # Move model to GPU for processing
+    MODEL.to('cuda')
+    MODEL.eval()
     
-    # Normalize depth for visualization (0-255)
-    depth_normalized = ((depth - depth.min()) / (depth.max() - depth.min()) * 255).astype(np.uint8)
+    # Convert from RGB to BGR
+    input_image = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
+    
+    with torch.no_grad():
+        # Get depth map
+        depth = MODEL.infer_image(input_image)
+        
+        # Normalize depth for visualization (0-255)
+        depth_normalized = ((depth - depth.min()) / (depth.max() - depth.min()) * 255).astype(np.uint8)
+        
+        # Apply colormap for better visualization
+        depth_colormap = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_INFERNO)
+        depth_colormap = cv2.cvtColor(depth_colormap, cv2.COLOR_BGR2RGB)
     
-    # Apply colormap for better visualization
-    depth_colormap = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_INFERNO)
-    depth_colormap = cv2.cvtColor(depth_colormap, cv2.COLOR_BGR2RGB)  # Convert back to RGB for Gradio
+    # Move model back to CPU after processing
+    MODEL.to('cpu')
     
     return depth_normalized, depth_colormap
 
-# Create Gradio interface
+@spaces.GPU
 def gradio_interface(input_img):
-    depth_raw, depth_colored = process_image(input_img)
-    return [input_img, depth_raw, depth_colored]
+    try:
+        depth_raw, depth_colored = process_image(input_img)
+        return [input_img, depth_raw, depth_colored]
+    except Exception as e:
+        print(f"Error processing image: {str(e)}")
+        return [input_img, None, None]
 
 # Define interface
 iface = gr.Interface(
@@ -77,7 +90,7 @@ iface = gr.Interface(
     ],
     title="Depth Estimation",
     description="Upload an image to generate its depth map.",
-    examples=["image.jpg"]  # Add example images here
+    examples=["image.jpg"]
 )
 
 # Launch the app