Update app.py

app.py

@@ -10,13 +10,14 @@ from PIL import Image
 import moviepy.editor as mp
 from infer import lotus, load_models
 import logging
+import spaces  # Import spaces for @spaces.GPU decorator
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Set device to use the L40s GPU explicitly
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+# Set device to use GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load models once
 task_name = 'depth'
@@ -37,8 +38,6 @@ def preprocess_video(video_path, target_fps=24):
 def process_frame(frame, seed=0):
     """Process a single frame and return depth map."""
     try:
-        torch.cuda.empty_cache()  # Clear GPU cache
-
         # Convert frame to PIL Image
         image = Image.fromarray(frame).convert('RGB')
 
@@ -51,7 +50,7 @@ def process_frame(frame, seed=0):
         logger.error(f"Error processing frame: {e}")
         return None
 
-# Process video frames and generate depth maps
+@spaces.GPU  # Decorate the function to use Hugging Face Spaces GPUs
 def process_video(video_path, fps=0, seed=0):
     """Process video frames individually and generate depth maps."""
     # Create a persistent temporary directory