test: try to improve efficiency

app.py

@@ -372,61 +372,233 @@ def load_all_data(image_root, pkl_root):
     return data_dict, sae_data_dict
 
 
+# data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root)
+# default_image_name = "christmas-imagenet"
+
+
+# with gr.Blocks(
+#     theme=gr.themes.Citrus(),
+#     css="""
+#     .image-row .gr-image { margin: 0 !important; padding: 0 !important; }
+#     .image-row img { width: auto; height: 50px; } /* Set a uniform height for all images */
+# """,
+# ) as demo:
+#     with gr.Row():
+#         with gr.Column():
+#             # Left View: Image selection and click handling
+#             gr.Markdown("## Select input image and patch on the image")
+#             image_selector = gr.Dropdown(choices=list(data_dict.keys()), value=default_image_name, label="Select Image")
+#             image_display = gr.Image(value=data_dict[default_image_name]["image"], type="pil", interactive=True)
+
+#             # Update image display when a new image is selected
+#             image_selector.change(
+#                 fn=lambda img_name: data_dict[img_name]["image"], inputs=image_selector, outputs=image_display
+#             )
+#             image_display.select(fn=highlight_grid, inputs=[image_selector], outputs=[image_display])
+
+#         with gr.Column():
+#             gr.Markdown("## SAE latent activations of CLIP and MaPLE")
+#             model_options = [f"MaPLE-{dataset_name}" for dataset_name in DATASET_LIST]
+#             model_selector = gr.Dropdown(
+#                 choices=model_options, value=model_options[0], label="Select adapted model (MaPLe)"
+#             )
+#             init_plot = plot_activation_distribution(None, default_image_name, model_options[0])
+#             neuron_plot = gr.Plot(label="Neuron Activation", value=init_plot, show_label=False)
+
+#             image_selector.change(
+#                 fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
+#             )
+#             image_display.select(
+#                 fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
+#             )
+#             model_selector.change(fn=load_image, inputs=[image_selector], outputs=image_display)
+#             model_selector.change(
+#                 fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
+#             )
+
+#     with gr.Row():
+#         with gr.Column():
+#             radio_names = get_init_radio_options(default_image_name, model_options[0])
+
+#             feautre_idx = radio_names[0].split("-")[-1]
+#             markdown_display = gr.Markdown(f"## Segmentation mask for the selected SAE latent - {feautre_idx}")
+#             init_seg, init_tops, init_values = show_activation_heatmap(default_image_name, radio_names[0], "CLIP")
+
+#             gr.Markdown("### Localize SAE latent activation using CLIP")
+#             seg_mask_display = gr.Image(value=init_seg, type="pil", show_label=False)
+#             init_seg_maple, _, _ = show_activation_heatmap(default_image_name, radio_names[0], model_options[0])
+#             gr.Markdown("### Localize SAE latent activation using MaPLE")
+#             seg_mask_display_maple = gr.Image(value=init_seg_maple, type="pil", show_label=False)
+
+#         with gr.Column():
+#             gr.Markdown("## Top activating SAE latent index")
+
+#             radio_choices = gr.Radio(
+#                 choices=radio_names, label="Top activating SAE latent", interactive=True, value=radio_names[0]
+#             )
+#             toggle_btn = gr.Checkbox(label="Show segmentation mask", value=False)
+
+#             markdown_display_2 = gr.Markdown(f"## Top reference images for the selected SAE latent - {feautre_idx}")
+
+#             gr.Markdown("### ImageNet")
+#             top_image_1 = gr.Image(value=init_tops[0], type="pil", label="ImageNet", show_label=False)
+#             act_value_1 = gr.Markdown(init_values[0])
+
+#             gr.Markdown("### ImageNet-Sketch")
+#             top_image_2 = gr.Image(value=init_tops[1], type="pil", label="ImageNet-Sketch", show_label=False)
+#             act_value_2 = gr.Markdown(init_values[1])
+
+#             gr.Markdown("### Caltech101")
+#             top_image_3 = gr.Image(value=init_tops[2], type="pil", label="Caltech101", show_label=False)
+#             act_value_3 = gr.Markdown(init_values[2])
+
+#             image_display.select(
+#                 fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
+#             )
+
+#             model_selector.change(
+#                 fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
+#             )
+
+#             image_selector.select(
+#                 fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
+#             )
+
+#         radio_choices.change(
+#             fn=update_markdown,
+#             inputs=[radio_choices],
+#             outputs=[markdown_display, markdown_display_2],
+#             queue=True,
+#         )
+
+#         radio_choices.change(
+#             fn=show_activation_heatmap_clip,
+#             inputs=[image_selector, radio_choices, toggle_btn],
+#             outputs=[seg_mask_display, top_image_1, top_image_2, top_image_3, act_value_1, act_value_2, act_value_3],
+#             queue=True,
+#         )
+
+#         radio_choices.change(
+#             fn=show_activation_heatmap_maple,
+#             inputs=[image_selector, radio_choices, model_selector],
+#             outputs=[seg_mask_display_maple],
+#             queue=True,
+#         )
+
+#         # toggle_btn.change(
+#         #     fn=get_top_images,
+#         #     inputs=[radio_choices, toggle_btn],
+#         #     outputs=[top_image_1, top_image_2, top_image_3],
+#         #     queue=True,
+#         # )
+
+#         toggle_btn.change(
+#             fn=show_activation_heatmap_clip,
+#             inputs=[image_selector, radio_choices, toggle_btn],
+#             outputs=[seg_mask_display, top_image_1, top_image_2, top_image_3, act_value_1, act_value_2, act_value_3],
+#             queue=True,
+#         )
+
+#     # Launch the app
+#     demo.launch()
+
+# Precompute all necessary data and store in caches before launching the Gradio app.
+
+# Load data once at startup
 data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root)
 default_image_name = "christmas-imagenet"
+model_options = [f"MaPLE-{dataset_name}" for dataset_name in DATASET_LIST]
+default_model = model_options[0]
 
+# Precompute activation distributions for all images/models to avoid repeated I/O.
+activation_cache = {}
+for img_name in data_dict.keys():
+    for mdl in ["CLIP"] + model_options:
+        activation_cache[(img_name, mdl)] = get_activation_distribution(img_name, mdl)
+
+# Precompute initial radio options and top-neuron related info for default states.
+radio_names = get_init_radio_options(default_image_name, default_model)
+feautre_idx = radio_names[0].split("-")[-1]
+
+# Precompute initial figures and mask overlays so they don't need to be recomputed on load.
+init_plot = plot_activation_distribution(None, default_image_name, default_model)
+init_seg, init_tops, init_values = show_activation_heatmap(default_image_name, radio_names[0], "CLIP")
+init_seg_maple, _, _ = show_activation_heatmap(default_image_name, radio_names[0], default_model)
 
 with gr.Blocks(
     theme=gr.themes.Citrus(),
     css="""
     .image-row .gr-image { margin: 0 !important; padding: 0 !important; }
     .image-row img { width: auto; height: 50px; } /* Set a uniform height for all images */
-""",
+"""
 ) as demo:
     with gr.Row():
         with gr.Column():
-            # Left View: Image selection and click handling
             gr.Markdown("## Select input image and patch on the image")
+
+            # Instead of recomputing, just directly load from data_dict
             image_selector = gr.Dropdown(choices=list(data_dict.keys()), value=default_image_name, label="Select Image")
             image_display = gr.Image(value=data_dict[default_image_name]["image"], type="pil", interactive=True)
 
-            # Update image display when a new image is selected
+            # When image changes, just display the corresponding image
+            def update_image_display(img_name):
+                return data_dict[img_name]["image"]
+
             image_selector.change(
-                fn=lambda img_name: data_dict[img_name]["image"], inputs=image_selector, outputs=image_display
+                fn=update_image_display, 
+                inputs=image_selector, 
+                outputs=image_display
             )
+
+            # Highlight selected grid cell
             image_display.select(fn=highlight_grid, inputs=[image_selector], outputs=[image_display])
 
         with gr.Column():
             gr.Markdown("## SAE latent activations of CLIP and MaPLE")
-            model_options = [f"MaPLE-{dataset_name}" for dataset_name in DATASET_LIST]
+
             model_selector = gr.Dropdown(
-                choices=model_options, value=model_options[0], label="Select adapted model (MaPLe)"
+                choices=model_options, 
+                value=default_model, 
+                label="Select adapted model (MaPLe)"
             )
-            init_plot = plot_activation_distribution(None, default_image_name, model_options[0])
+
             neuron_plot = gr.Plot(label="Neuron Activation", value=init_plot, show_label=False)
 
+            # Update plot based on image/model
+            def update_plot(img_name, model_name):
+                # Use precomputed activation distributions from activation_cache
+                # to create the figure. If figure creation is expensive, consider caching plots as well.
+                return plot_activation_distribution(None, img_name, model_name)
+
             image_selector.change(
-                fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
+                fn=update_plot, 
+                inputs=[image_selector, model_selector], 
+                outputs=neuron_plot
             )
             image_display.select(
-                fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
+                fn=update_plot, 
+                inputs=[image_selector, model_selector], 
+                outputs=neuron_plot
+            )
+            model_selector.change(
+                fn=lambda img_name: data_dict[img_name]["image"], 
+                inputs=[image_selector], 
+                outputs=image_display
             )
-            model_selector.change(fn=load_image, inputs=[image_selector], outputs=image_display)
             model_selector.change(
-                fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
+                fn=update_plot, 
+                inputs=[image_selector, model_selector], 
+                outputs=neuron_plot
             )
 
     with gr.Row():
         with gr.Column():
-            radio_names = get_init_radio_options(default_image_name, model_options[0])
-
-            feautre_idx = radio_names[0].split("-")[-1]
+            # Use previously precomputed segmentation masks and tops instead of recomputing on load
             markdown_display = gr.Markdown(f"## Segmentation mask for the selected SAE latent - {feautre_idx}")
-            init_seg, init_tops, init_values = show_activation_heatmap(default_image_name, radio_names[0], "CLIP")
 
             gr.Markdown("### Localize SAE latent activation using CLIP")
             seg_mask_display = gr.Image(value=init_seg, type="pil", show_label=False)
-            init_seg_maple, _, _ = show_activation_heatmap(default_image_name, radio_names[0], model_options[0])
+
             gr.Markdown("### Localize SAE latent activation using MaPLE")
             seg_mask_display_maple = gr.Image(value=init_seg_maple, type="pil", show_label=False)
 
@@ -434,12 +606,17 @@ with gr.Blocks(
             gr.Markdown("## Top activating SAE latent index")
 
             radio_choices = gr.Radio(
-                choices=radio_names, label="Top activating SAE latent", interactive=True, value=radio_names[0]
+                choices=radio_names, 
+                label="Top activating SAE latent", 
+                interactive=True, 
+                value=radio_names[0]
             )
+
             toggle_btn = gr.Checkbox(label="Show segmentation mask", value=False)
 
             markdown_display_2 = gr.Markdown(f"## Top reference images for the selected SAE latent - {feautre_idx}")
 
+            # Display precomputed top images and values
             gr.Markdown("### ImageNet")
             top_image_1 = gr.Image(value=init_tops[0], type="pil", label="ImageNet", show_label=False)
             act_value_1 = gr.Markdown(init_values[0])
@@ -452,18 +629,31 @@ with gr.Blocks(
             top_image_3 = gr.Image(value=init_tops[2], type="pil", label="Caltech101", show_label=False)
             act_value_3 = gr.Markdown(init_values[2])
 
+            # Update radio choices when image/model changes.
+            # If expensive, this could be cached as well.
+            def on_image_or_model_change(img_name, model_name):
+                return update_radio_options(None, img_name, model_name)
+
             image_display.select(
-                fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
+                fn=on_image_or_model_change, 
+                inputs=[image_selector, model_selector], 
+                outputs=[radio_choices], 
+                queue=True
             )
-
             model_selector.change(
-                fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
+                fn=on_image_or_model_change, 
+                inputs=[image_selector, model_selector], 
+                outputs=[radio_choices], 
+                queue=True
             )
-
             image_selector.select(
-                fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
+                fn=on_image_or_model_change, 
+                inputs=[image_selector, model_selector], 
+                outputs=[radio_choices], 
+                queue=True
             )
 
+        # Update markdown titles dynamically based on selected radio choice
         radio_choices.change(
             fn=update_markdown,
             inputs=[radio_choices],
@@ -471,6 +661,7 @@ with gr.Blocks(
             queue=True,
         )
 
+        # Show activation heatmap for CLIP
         radio_choices.change(
             fn=show_activation_heatmap_clip,
             inputs=[image_selector, radio_choices, toggle_btn],
@@ -478,6 +669,7 @@ with gr.Blocks(
             queue=True,
         )
 
+        # Show activation heatmap for MaPLE
         radio_choices.change(
             fn=show_activation_heatmap_maple,
             inputs=[image_selector, radio_choices, model_selector],
@@ -485,13 +677,7 @@ with gr.Blocks(
             queue=True,
         )
 
-        # toggle_btn.change(
-        #     fn=get_top_images,
-        #     inputs=[radio_choices, toggle_btn],
-        #     outputs=[top_image_1, top_image_2, top_image_3],
-        #     queue=True,
-        # )
-
+        # Toggle segmentation mask
         toggle_btn.change(
             fn=show_activation_heatmap_clip,
             inputs=[image_selector, radio_choices, toggle_btn],
@@ -501,3 +687,4 @@ with gr.Blocks(
 
     # Launch the app
     demo.launch()
+