Assuming GPUZero is always available.

app.py

@@ -10,29 +10,28 @@ from pathlib import Path
 from PIL import Image
 
 from plots import get_pre_define_colors
-# from utils.load_model import load_xclip
+from utils.load_model import load_xclip
 from utils.predict import xclip_pred
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-# def initialize_model():
-#     global XCLIP, OWLVIT_PRECESSOR
-#     if XCLIP is None or OWLVIT_PRECESSOR is None:
-#         XCLIP, OWLVIT_PRECESSOR = load_xclip(DEVICE)
 
 #! Huggingface does not allow load model to main process, so we need to load the model when needed, it may not help in improve the speed of the app.
 try:
     import spaces
     XCLIP, OWLVIT_PRECESSOR = None, None
+    DEVICE = 'cuda'
 except:
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
     print(f"Not at Huggingface demo, load model to main process.")
     XCLIP, OWLVIT_PRECESSOR = load_xclip(DEVICE)
-    
+
+print(f"Device: {DEVICE}")
+
 XCLIP_DESC_PATH = "data/jsons/bs_cub_desc.json"
 XCLIP_DESC = json.load(open(XCLIP_DESC_PATH, "r"))
 IMAGES_FOLDER = "data/images"
 # XCLIP_RESULTS = json.load(open("data/jsons/xclip_org.json", "r"))
 IMAGE2GT = json.load(open("data/jsons/image2gt.json", 'r'))
-CUB_DESC_EMBEDS = torch.load('data/text_embeddings/cub_200_desc.pt').to(DEVICE)
+CUB_DESC_EMBEDS = torch.load('data/text_embeddings/cub_200_desc.pt')
 CUB_IDX2NAME = json.load(open('data/jsons/cub_desc_idx2name.json', 'r'))
 CUB_IDX2NAME = {int(k): v for k, v in CUB_IDX2NAME.items()}
 

utils/predict.py

@@ -91,7 +91,7 @@ def xclip_pred(new_desc: dict,
             modified_class_idx = 200
         else:
             n_classes = 200
-            query_embeds = cub_embeds
+            query_embeds = cub_embeds.to(device)
             idx2name = cub_idx2name
             modified_class_idx = None