Don't run activations viz with ResNet model, to save memory

app.py

@@ -27,7 +27,7 @@ from CLIP_Explainability.vit_cam import (
 
 from pytorch_grad_cam.grad_cam import GradCAM
 
-RUN_LITE = False  # Load vision model for CAM viz explainability for M-CLIP only
+RUN_LITE = True  # Load models for CAM viz for M-CLIP and J-CLIP only
 
 MAX_IMG_WIDTH = 500
 MAX_IMG_HEIGHT = 800
@@ -110,6 +110,10 @@ def clip_search(search_query):
 
 
 def string_search():
+    st.session_state.disable_uploader = (
+        RUN_LITE and st.session_state.active_model == "Legacy (multilingual ResNet)"
+    )
+
     if "search_field_value" in st.session_state:
         clip_search(st.session_state.search_field_value)
 
@@ -179,10 +183,9 @@ 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"
 
-        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_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
@@ -216,6 +219,9 @@ def init():
     st.session_state.search_image_ids = []
     st.session_state.search_image_scores = {}
     st.session_state.text_table_df = None
+    st.session_state.disable_uploader = (
+        RUN_LITE and st.session_state.active_model == "Legacy (multilingual ResNet)"
+    )
 
     with st.spinner("Loading models and data, please wait..."):
         load_image_features()
@@ -430,7 +436,7 @@ def visualize_gradcam(image):
 
     header_cols = st.columns([80, 20], vertical_alignment="bottom")
     with header_cols[0]:
-        st.title("Image + query details")
+        st.title("Image + query activation gradients")
     with header_cols[1]:
         if st.button("Close"):
             st.rerun()
@@ -457,6 +463,8 @@ def visualize_gradcam(image):
             st.session_state.search_field_value
         )
 
+    st.image(image)
+
     with st.spinner("Calculating..."):
         # info_text = st.text("Calculating activation regions...")
 
@@ -743,6 +751,7 @@ with controls[3]:
         key="uploaded_image",
         label_visibility="collapsed",
         on_change=vis_uploaded_image,
+        disabled=st.session_state.disable_uploader,
     )
 
 
@@ -777,7 +786,9 @@ for image_id in batch:
                 <div>""",
             unsafe_allow_html=True,
         )
-        if not RUN_LITE or st.session_state.active_model == "M-CLIP (multilingual ViT)":
+        if not (
+            RUN_LITE and st.session_state.active_model == "Legacy (multilingual ResNet)"
+        ):
             st.button(
                 "Explain this",
                 on_click=vis_known_image,
@@ -785,4 +796,6 @@ for image_id in batch:
                 use_container_width=True,
                 key=image_id,
             )
+        else:
+            st.empty()
     col = (col + 1) % row_size