load model inside gr.Block

app.py

@@ -1,16 +1,10 @@
 import os
-import gradio as gr
-print(f"Gradio version {gr.__version__}")
-# if gr.__version__ != '4.28.2':
-#     os.system("pip uninstall -y gradio")
-#     os.system("pip install gradio==4.28.2")
-#     print(f"Gradio version: {gr.__version__}")
-
 import io
 
 import torch
 import json
 import base64
+import gradio as gr
 import numpy as np
 from pathlib import Path
 from PIL import Image
@@ -20,7 +14,12 @@ 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 = load_xclip(DEVICE)
+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)
+    
 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')
@@ -383,6 +382,7 @@ custom_css = """
 
 # Define the Gradio interface
 with gr.Blocks(theme=gr.themes.Soft(), css=custom_css, title="PEEB") as demo:
+    initialize_model()
     current_image = gr.State("")
     current_predicted_class = gr.State("")
     gt_class = gr.State("")