Update app.py

app.py

@@ -3,11 +3,13 @@ import torch
 import numpy as np
 from PIL import Image
 import matplotlib.pyplot as plt
+import spaces
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-depth_anything = pipeline(task = "depth-estimation", model="nielsr/depth-anything-small", device=0)
+depth_anything = pipeline(task = "depth-estimation", model="nielsr/depth-anything-small", device=device)
 checkpoint = "BAAI/seggpt-vit-large"
 image_processor = SegGptImageProcessor.from_pretrained(checkpoint)
-model = SegGptForImageSegmentation.from_pretrained(checkpoint)
+model = SegGptForImageSegmentation.from_pretrained(checkpoint).to(device)
 
 def infer_seggpt(image_input, image_prompt, mask_prompt):
   num_labels = 100
@@ -17,7 +19,7 @@ def infer_seggpt(image_input, image_prompt, mask_prompt):
         prompt_masks=mask_prompt,
         return_tensors="pt",
         num_labels=num_labels
-    )
+    ).to(device)
   with torch.no_grad():
       outputs = model(**inputs)
 
@@ -38,6 +40,7 @@ def infer_seggpt(image_input, image_prompt, mask_prompt):
   plt.savefig("masks.png", bbox_inches='tight', pad_inches=0)
   return "masks.png"
 
+@spaces.GPU
 def infer(image_input, image_prompt, mask_prompt):
   sg_masks = []
   mask_prompt = depth_anything(image_prompt)["depth"].convert("RGB")