Update predict.py

predict.py

@@ -13,17 +13,18 @@ from detectron2.utils.visualizer import ColorMode, Visualizer
 from color_palette import ade_palette
 from transformers import Mask2FormerImageProcessor, Mask2FormerForUniversalSegmentation
 
+@spaces.GPU
 def load_model_and_processor(model_ckpt: str):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = Mask2FormerForUniversalSegmentation.from_pretrained(model_ckpt).to(torch.device(device))
     model.eval()
     image_preprocessor = Mask2FormerImageProcessor.from_pretrained(model_ckpt)
     return model, image_preprocessor
-
+@spaces.GPU
 def load_default_ckpt():
     default_ckpt = "facebook/mask2former-swin-tiny-coco-panoptic"
     return default_ckpt
-
+@spaces.GPU
 def draw_panoptic_segmentation(predicted_segmentation_map, seg_info, image):
     metadata = MetadataCatalog.get("coco_2017_val_panoptic")
     for res in seg_info:
@@ -41,7 +42,7 @@ def draw_panoptic_segmentation(predicted_segmentation_map, seg_info, image):
     return output_img, labels
 
 
-
+@spaces.GPU
 def predict_masks(input_img_path: str):
     
     #load model and image processor
@@ -66,7 +67,7 @@ def predict_masks(input_img_path: str):
 
 
 
-
+@spaces.GPU
 def get_mask_for_label(results, label):
     import numpy as np
     from PIL import Image