Update app.py

app.py

@@ -12,8 +12,11 @@ import os
 # Run the script to get pretrained models
 subprocess.run(["bash", "get_pretrained_models.sh"])
 
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
 # Load model and preprocessing transform
 model, transform = depth_pro.create_model_and_transforms()
+model = model.to(device)
 model.eval()
 
 def resize_image(image_path, max_size=1024):
@@ -30,7 +33,7 @@ def resize_image(image_path, max_size=1024):
             img.save(temp_file, format="PNG")
             return temp_file.name
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=20)
 def predict_depth(input_image):
     temp_file = None
     try:
@@ -42,6 +45,7 @@ def predict_depth(input_image):
         image = result[0]
         f_px = result[-1]  # Assuming f_px is the last item in the returned tuple
         image = transform(image)
+        image = image.to(device)
 
         # Run inference
         prediction = model.infer(image, f_px=f_px)