Global lock to avoid concurrent caching

global_session.py

@@ -0,0 +1,17 @@
+# Hacky way to provide global session.
+import sys
+
+GLOBAL_CONTAINER = sys
+
+
+class GlobalState(object):
+
+    def __new__(cls, key="default"):
+        if not hasattr(GLOBAL_CONTAINER, '_global_states'):
+            GLOBAL_CONTAINER._global_states = {}
+            print("Global state container created")
+
+        if not GLOBAL_CONTAINER._global_states.get(key):
+            GLOBAL_CONTAINER._global_states[key] = super(GlobalState, cls).__new__(cls)
+
+        return GLOBAL_CONTAINER._global_states[key]

utils.py

@@ -5,7 +5,8 @@ from transformers import AutoTokenizer, CLIPProcessor, ViTFeatureExtractor
 
 from config import MODEL_LIST
 from koclip import FlaxHybridCLIP
-
+from global_session import GlobalState
+from threading import Lock
 
 @st.cache(allow_output_mutation=True)
 def load_index(img_file):
@@ -24,8 +25,20 @@ def load_index(img_file):
     return filenames, index
 
 
-@st.cache(allow_output_mutation=True)
 def load_model(model_name="koclip/koclip-base"):
+    state = GlobalState(model_name)
+    if not hasattr(state, '_lock'):
+        state._lock = Lock()
+    print(f"Locking loading of model : {model_name} to avoid concurrent caching.")
+
+    with state._lock:
+        cached_model = load_model_cached(model_name)
+
+    print(f"Unlocking loading of model : {model_name} to avoid concurrent caching.")
+    return cached_model
+
+@st.cache(allow_output_mutation=True)
+def load_model_cached(model_name):
     assert model_name in {f"koclip/{model}" for model in MODEL_LIST}
     model = FlaxHybridCLIP.from_pretrained(model_name)
     processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")