Update app.py

app.py

@@ -7,7 +7,9 @@ from ultralytics.utils.torch_utils import smart_inference_mode
 from ultralytics.models.yolo.yoloe.predict_vp import YOLOEVPSegPredictor
 from gradio_image_prompter import ImagePrompter
 from huggingface_hub import hf_hub_download
+import spaces
 
+@spaces.GPU
 def init_model(model_id, is_pf=False):
     if not is_pf:
         path = hf_hub_download(repo_id="jameslahm/yoloe", filename=f"{model_id}-seg.pt")
@@ -16,10 +18,9 @@ def init_model(model_id, is_pf=False):
         path = hf_hub_download(repo_id="jameslahm/yoloe", filename=f"{model_id}-seg-pf.pt")
         model = YOLOE(path)
     model.eval()
-    model.to("cuda" if torch.cuda.is_available() else "cpu")
     return model
 
-
+@spaces.GPU
 @smart_inference_mode()
 def yoloe_inference(image, prompts, target_image, model_id, image_size, conf_thresh, iou_thresh, prompt_type):
     model = init_model(model_id)