load the model to cpu first

app.py

@@ -15,14 +15,9 @@ from utils.predict import xclip_pred
 
 
 #! 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)
+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}")
 

utils/predict.py

@@ -53,8 +53,12 @@ def xclip_pred(new_desc: dict,
                cub_embeds: torch.Tensor = None,
                cub_idx2name: dict = None,
                descriptors: dict = None):
-    if model is None or owlvit_processor is None:
-        model, owlvit_processor = load_xclip(device=device)
+    # check if in huggingface space
+    try:
+        model.to('cuda')
+        device = 'cuda'
+    except:
+        device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
     
     # reorder the new description and the mask
     if new_class is not None:
@@ -78,7 +82,9 @@ def xclip_pred(new_desc: dict,
         n_classes = len(getprompt.name2idx)
         descs, class_idxs, class_mapping, org_desc_mapper, class_list = getprompt('chatgpt-no-template', max_len=12, pad=True)
         query_embeds = encode_descs_xclip(owlvit_processor, model, descs, device)
+        
     else:
+        cub_embeds = cub_embeds.to(device)
         if new_class is not None:
             if new_class in list(cub_idx2name.values()):
                 new_class = f"{new_class}_custom"
@@ -87,11 +93,11 @@ def xclip_pred(new_desc: dict,
             n_classes = 201
             query_tokens = owlvit_processor(text=list(new_desc_.values()), padding="max_length", truncation=True, return_tensors="pt").to(device)
             new_class_embed = model.owlvit.get_text_features(**query_tokens)
-            query_embeds = torch.cat([cub_embeds, new_class_embed], dim=0).to(device)
+            query_embeds = torch.cat([cub_embeds, new_class_embed], dim=0)
             modified_class_idx = 200
         else:
             n_classes = 200
-            query_embeds = cub_embeds.to(device)
+            query_embeds = cub_embeds
             idx2name = cub_idx2name
             modified_class_idx = None