add gpu support

app.py

@@ -1,3 +1,4 @@
+import spaces
 import io
 import os
 debug = False

utils/load_model.py

@@ -1,10 +1,12 @@
 
 
+import spaces
 import torch
 from transformers import OwlViTProcessor, OwlViTForObjectDetection
 
 from .model import OwlViTForClassification
 
+@spaces.GPU
 def load_xclip(device: str = "cuda:0", 
                n_classes: int = 183, 
                use_teacher_logits: bool = False, 

utils/predict.py

@@ -1,3 +1,4 @@
+import spaces
 import PIL
 import torch
 
@@ -29,7 +30,7 @@ def encode_descs_xclip(owlvit_det_processor: callable, model: callable, descs: l
 #     text_embeds = torch.cat(text_embeds, dim=0)
 #     text_embeds = torch.nn.functional.normalize(text_embeds, dim=-1)
 #     return text_embeds.to(device)
-
+@spaces.GPU
 def xclip_pred(new_desc: dict, 
                new_part_mask: dict, 
                new_class: str,