Update app.py

app.py

@@ -1,11 +1,12 @@
+import spaces
 import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
 import torch
+import torch.nn.functional as F
 from torchvision.transforms import Compose
 import tempfile
-from PIL import Image
-import numpy as np
-import cv2
-import torch.nn.functional as F
 
 from depth_anything.dpt import DPT_DINOv2
 from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
@@ -45,6 +46,11 @@ transform = Compose([
         PrepareForNet(),
 ])
 
+@spaces.GPU
+def predict_depth(model, image):
+    with torch.no_grad():
+        return model(image)
+
 with gr.Blocks(css=css) as demo:
     gr.Markdown(title)
     gr.Markdown(description)
@@ -63,9 +69,8 @@ with gr.Blocks(css=css) as demo:
         image = transform({'image': image})['image']
         image = torch.from_numpy(image).unsqueeze(0).to(DEVICE)
         
-        with torch.no_grad():
-            depth = model(image)
-            depth = F.interpolate(depth[None], (h, w), mode='bilinear', align_corners=False)[0, 0]
+        depth = predict_depth(model, image)
+        depth = F.interpolate(depth[None], (h, w), mode='bilinear', align_corners=False)[0, 0]
         
         raw_depth = Image.fromarray(depth.cpu().numpy().astype('uint16'))
         tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)