cannot load to main process, load model directly at inference.

app.py

@@ -10,19 +10,25 @@ 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"
-XCLIP, OWLVIT_PRECESSOR = None, None
-def initialize_model():
-    global XCLIP, OWLVIT_PRECESSOR
-    if XCLIP is None or OWLVIT_PRECESSOR is None:
-        XCLIP, OWLVIT_PRECESSOR = load_xclip(DEVICE)
+# 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
+except:
+    print(f"Not at Huggingface demo, load model to main process.")
+    XCLIP, OWLVIT_PRECESSOR = load_xclip(DEVICE)
     
 XCLIP_DESC_PATH = "data/jsons/bs_cub_desc.json"
 XCLIP_DESC = json.load(open(XCLIP_DESC_PATH, "r"))
-PREPROCESS = lambda x: OWLVIT_PRECESSOR(images=x, return_tensors='pt')
 IMAGES_FOLDER = "data/images"
 # XCLIP_RESULTS = json.load(open("data/jsons/xclip_org.json", "r"))
 IMAGE2GT = json.load(open("data/jsons/image2gt.json", 'r'))
@@ -198,7 +204,6 @@ PART_IMAGES_DICT = {file_name: load_part_images(file_name) for file_name in IMAG
 
 # --- Gradio Functions ---
 def update_selected_image(event: gr.SelectData):
-    initialize_model()
     image_height = 400
     index = event.index
 

utils/predict.py

@@ -1,6 +1,7 @@
 try:
     import spaces
     gpu_decorator = spaces.GPU
+    from .load_model import load_xclip
 except ImportError:
     # Define a no-operation decorator as fallback
     def gpu_decorator(func):
@@ -52,6 +53,9 @@ 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)
+    
     # reorder the new description and the mask
     if new_class is not None:
         new_desc_ = {k: new_desc[k] for k in ORG_PART_ORDER}