Upload app.py

app.py

@@ -27,6 +27,8 @@ from CLIP_Explainability.vit_cam import (
 
 from pytorch_grad_cam.grad_cam import GradCAM
 
+RUN_LITE = True  # Load vision model for CAM viz explainability for M-CLIP only
+
 MAX_IMG_WIDTH = 500
 MAX_IMG_HEIGHT = 800
 
@@ -172,9 +174,10 @@ def init():
         ja_model_name = "hakuhodo-tech/japanese-clip-vit-h-14-bert-wider"
         ja_model_path = "./models/ViT-H-14-laion2B-s32B-b79K.bin"
 
-        st.session_state.ja_image_model, st.session_state.ja_image_preprocess = load(
-            ja_model_path, device=device, jit=False
-        )
+        if not RUN_LITE:
+            st.session_state.ja_image_model, st.session_state.ja_image_preprocess = (
+                load(ja_model_path, device=device, jit=False)
+            )
 
         st.session_state.ja_model = AutoModel.from_pretrained(
             ja_model_name, trust_remote_code=True
@@ -183,9 +186,10 @@ def init():
             ja_model_name, trust_remote_code=True
         )
 
-        st.session_state.rn_image_model, st.session_state.rn_image_preprocess = (
-            clip.load("RN50x4", device=device)
-        )
+        if not RUN_LITE:
+            st.session_state.rn_image_model, st.session_state.rn_image_preprocess = (
+                clip.load("RN50x4", device=device)
+            )
 
         st.session_state.rn_model = legacy_multilingual_clip.load_model(
             "M-BERT-Base-69"
@@ -701,11 +705,12 @@ for image_id in batch:
                 <div>""",
             unsafe_allow_html=True,
         )
-        st.button(
-            "Explain this",
-            on_click=image_modal,
-            args=[image_id],
-            use_container_width=True,
-            key=image_id,
-        )
+        if not RUN_LITE or st.session_state.active_model == "M-CLIP (multilingual ViT)":
+            st.button(
+                "Explain this",
+                on_click=image_modal,
+                args=[image_id],
+                use_container_width=True,
+                key=image_id,
+            )
     col = (col + 1) % row_size