test: add memory monitor

app.py

@@ -609,148 +609,1078 @@ def load_all_data(image_root, pkl_root):
 default_image_name = "christmas-imagenet"
 
 
-def safe_operation(func, *args, timeout=10, default=None):
-    """Execute function with timeout and return default value on failure."""
-    try:
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            future = executor.submit(func, *args)
-            return future.result(timeout=timeout)
-    except concurrent.futures.TimeoutError:
-        print(f"Operation timed out: {func.__name__}")
-        return default
-    except Exception as e:
-        print(f"Operation failed: {func.__name__}, Error: {e}")
-        return default
-
-def create_interface():
-    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; }
-        """,
-    ) as demo:
-        # State management for preventing duplicate operations
-        state = gr.State({
-            'last_update': 0,
-            'processing': False
-        })
-
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("## Select input image and patch on the image")
-                image_selector = gr.Dropdown(
-                    choices=list(_CACHE['data_dict'].keys()),
-                    value=default_image_name,
-                    label="Select Image",
-                )
-                image_display = gr.Image(
-                    value=_CACHE['data_dict'][default_image_name]["image"],
-                    type="pil",
-                    interactive=True,
-                )
-
-                def safe_update_image(img_name):
-                    return safe_operation(
-                        lambda: _CACHE['data_dict'][img_name]["image"],
-                        timeout=5,
-                        default=Image.new('RGB', (IMAGE_SIZE, IMAGE_SIZE), 'gray')
-                    )
-
-                image_selector.change(
-                    fn=safe_update_image,
-                    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)",
-                )
-
-                def safe_plot_activation(evt, selected_image, model_name):
-                    try:
-                        return safe_operation(
-                            plot_activation_distribution,
-                            evt, selected_image, model_name,
-                            timeout=10,
-                            default=go.Figure()  # Return empty figure on timeout
-                        )
-                    except Exception as e:
-                        print(f"Error in plot_activation: {e}")
-                        return go.Figure()
-
-                neuron_plot = gr.Plot(
-                    label="Neuron Activation",
-                    value=safe_plot_activation(None, default_image_name, model_options[0]),
-                    show_label=False,
-                )
-
-                def debounced_update(*args, minimum_interval=1.0):
-                    """Prevent updates that are too close together"""
-                    current_time = time.time()
-                    if current_time - state.value['last_update'] < minimum_interval:
-                        raise gr.Error("Please wait a moment before updating again")
-                    state.value['last_update'] = current_time
-                    return safe_plot_activation(*args)
-
-                # Add error handling and debouncing to all event handlers
-                image_selector.change(
-                    fn=debounced_update,
-                    inputs=[image_selector, model_selector],
-                    outputs=neuron_plot,
-                )
-                image_display.select(
-                    fn=debounced_update,
-                    inputs=[image_selector, model_selector],
-                    outputs=neuron_plot,
-                )
-                model_selector.change(
-                    fn=debounced_update,
-                    inputs=[image_selector, model_selector],
-                    outputs=neuron_plot,
-                )
-
-        # Rest of your interface code with similar error handling...
-        
-        return demo
+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])
 
-demo = create_interface()
+            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],
+            )
+
+            model_selector.change(
+                fn=update_radio_options,
+                inputs=[image_selector, model_selector],
+                outputs=[radio_choices],
+            )
+
+            image_selector.select(
+                fn=update_radio_options,
+                inputs=[image_selector, model_selector],
+                outputs=[radio_choices],
+            )
+
+        radio_choices.change(
+            fn=update_all,
+            inputs=[image_selector, radio_choices, toggle_btn, model_selector],
+            outputs=[
+                seg_mask_display,
+                seg_mask_display_maple,
+                top_image_1,
+                top_image_2,
+                top_image_3,
+                act_value_1,
+                act_value_2,
+                act_value_3,
+                markdown_display,
+                markdown_display_2,
+            ],
+        )
+
+        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,
+            ],
+        )
+
+    # Launch the app
+    # demo.queue()
+    # demo.launch()
 
+    
 if __name__ == "__main__":
-    # Configure logging
-    import logging
-    logging.basicConfig(
-        level=logging.INFO,
-        format='%(asctime)s - %(levelname)s - %(message)s'
+    demo.queue()  # Enable queuing for better handling of concurrent users
+    demo.launch(
+        server_name="0.0.0.0",  # Allow external access
+        server_port=7860,
+        share=False,  # Set to True if you want to create a public URL
+        show_error=True,
+        # Optimize concurrency
+        max_threads=8,  # Adjust based on your CPU cores
     )
+import gzip
+import os
+import pickle
+from glob import glob
+from time import sleep
+
+from functools import lru_cache
+import concurrent.futures
+from typing import Dict, Tuple, List
+
+import gradio as gr
+import numpy as np
+import plotly.graph_objects as go
+import torch
+from PIL import Image, ImageDraw
+from plotly.subplots import make_subplots
+
+IMAGE_SIZE = 400
+DATASET_LIST = ["imagenet", "oxford_flowers", "ucf101", "caltech101", "dtd", "eurosat"]
+GRID_NUM = 14
+pkl_root = "./data/out"
+preloaded_data = {}
+
+
+# Global cache for data
+_CACHE = {
+    'data_dict': {},
+    'sae_data_dict': {},
+    'model_data': {},
+    'segmasks': {},
+    'top_images': {}
+}
+
+def load_all_data(image_root: str, pkl_root: str) -> Tuple[Dict, Dict]:
+    """Load all data with optimized parallel processing."""
+    # Load images in parallel
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        image_files = glob(f"{image_root}/*")
+        future_to_file = {
+            executor.submit(_load_image_file, image_file): image_file 
+            for image_file in image_files
+        }
+        
+        for future in concurrent.futures.as_completed(future_to_file):
+            image_file = future_to_file[future]
+            image_name = os.path.basename(image_file).split(".")[0]
+            result = future.result()
+            if result is not None:
+                _CACHE['data_dict'][image_name] = result
+
+    # Load SAE data
+    with open("./data/sae_data/mean_acts.pkl", "rb") as f:
+        _CACHE['sae_data_dict']["mean_acts"] = pickle.load(f)
+
+    # Load mean act values in parallel
+    datasets = ["imagenet", "imagenet-sketch", "caltech101"]
+    _CACHE['sae_data_dict']["mean_act_values"] = {}
+    
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        future_to_dataset = {
+            executor.submit(_load_mean_act_values, dataset): dataset 
+            for dataset in datasets
+        }
+        
+        for future in concurrent.futures.as_completed(future_to_dataset):
+            dataset = future_to_dataset[future]
+            result = future.result()
+            if result is not None:
+                _CACHE['sae_data_dict']["mean_act_values"][dataset] = result
+
+    return _CACHE['data_dict'], _CACHE['sae_data_dict']
 
-    # Initialize the interface with error handling
+def _load_image_file(image_file: str) -> Dict:
+    """Helper function to load a single image file."""
     try:
-        demo = create_interface()
-        demo.queue()  # Limit concurrent operations
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            share=False,
-            show_error=True,
-            max_threads=4,  # Reduced for better stability
-            prevent_thread_lock=True,
-            # Add health check endpoint
-            root_path="/healthz",
-            # Add automatic cleanup
-            _teardown=lambda: (
-                clear_caches(),
-                logging.info("Caches cleared during teardown")
-            ),
-            # Configure timeouts
-            api_open_timeout=60,
-            api_call_timeout=300,
-        )
+        image = Image.open(image_file).resize((IMAGE_SIZE, IMAGE_SIZE))
+        return {
+            "image": image,
+            "image_path": image_file,
+        }
+    except Exception as e:
+        print(f"Error loading {image_file}: {e}")
+        return None
+
+def _load_mean_act_values(dataset: str) -> np.ndarray:
+    """Helper function to load mean act values for a dataset."""
+    try:
+        with gzip.open(f"./data/sae_data/mean_act_values_{dataset}.pkl.gz", "rb") as f:
+            return pickle.load(f)
     except Exception as e:
-        logging.error(f"Failed to launch interface: {e}")
-        raise
+        print(f"Error loading mean act values for {dataset}: {e}")
+        return None
+
+@lru_cache(maxsize=1024)
+def get_data(image_name: str, model_name: str) -> np.ndarray:
+    """Cached function to get model data."""
+    cache_key = f"{model_name}_{image_name}"
+    if cache_key not in _CACHE['model_data']:
+        data_dir = f"{pkl_root}/{model_name}/{image_name}.pkl.gz"
+        with gzip.open(data_dir, "rb") as f:
+            _CACHE['model_data'][cache_key] = pickle.load(f)
+    return _CACHE['model_data'][cache_key]
+
+@lru_cache(maxsize=1024)
+def get_activation_distribution(image_name: str, model_type: str) -> np.ndarray:
+    """Cached function to get activation distribution."""
+    activation = get_data(image_name, model_type)[0]
+    noisy_features_indices = (
+        (_CACHE['sae_data_dict']["mean_acts"]["imagenet"] > 0.1).nonzero()[0].tolist()
+    )
+    activation[:, noisy_features_indices] = 0
+    return activation
+
+@lru_cache(maxsize=1024)
+def get_segmask(selected_image: str, slider_value: int, model_type: str) -> np.ndarray:
+    """Cached function to get segmentation mask."""
+    cache_key = f"{selected_image}_{slider_value}_{model_type}"
+    if cache_key not in _CACHE['segmasks']:
+        image = _CACHE['data_dict'][selected_image]["image"]
+        sae_act = get_data(selected_image, model_type)[0]
+        temp = sae_act[:, slider_value]
+        
+        mask = torch.Tensor(temp[1:].reshape(14, 14)).view(1, 1, 14, 14)
+        mask = torch.nn.functional.interpolate(mask, (image.height, image.width))[0][0].numpy()
+        mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-10)
+
+        base_opacity = 30
+        image_array = np.array(image)[..., :3]
+        rgba_overlay = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
+        rgba_overlay[..., :3] = image_array[..., :3]
+
+        darkened_image = (image_array[..., :3] * (base_opacity / 255)).astype(np.uint8)
+        rgba_overlay[mask == 0, :3] = darkened_image[mask == 0]
+        rgba_overlay[..., 3] = 255
+
+        _CACHE['segmasks'][cache_key] = rgba_overlay
+        
+    return _CACHE['segmasks'][cache_key]
+
+@lru_cache(maxsize=1024)
+def get_top_images(slider_value: int, toggle_btn: bool) -> List[Image.Image]:
+    """Cached function to get top images."""
+    cache_key = f"{slider_value}_{toggle_btn}"
+    if cache_key not in _CACHE['top_images']:
+        dataset_path = "./data/top_images_masked" if toggle_btn else "./data/top_images"
+        paths = [
+            os.path.join(dataset_path, dataset, f"{slider_value}.jpg")
+            for dataset in ["imagenet", "imagenet-sketch", "caltech101"]
+        ]
+        
+        _CACHE['top_images'][cache_key] = [
+            Image.open(path) if os.path.exists(path) else Image.new("RGB", (256, 256), (255, 255, 255))
+            for path in paths
+        ]
+        
+    return _CACHE['top_images'][cache_key]
+
+
+# def preload_activation(image_name):
+#     for model in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
+#         image_file = f"{pkl_root}/{model}/{image_name}.pkl.gz"
+#         with gzip.open(image_file, "rb") as f:
+#             preloaded_data[model] = pickle.load(f)
+
+
+# def get_activation_distribution(image_name: str, model_type: str):
+#     activation = get_data(image_name, model_type)[0]
+
+#     noisy_features_indices = (
+#         (sae_data_dict["mean_acts"]["imagenet"] > 0.1).nonzero()[0].tolist()
+#     )
+#     activation[:, noisy_features_indices] = 0
+
+#     return activation
+
+
+def get_grid_loc(evt, image):
+    # Get click coordinates
+    x, y = evt._data["index"][0], evt._data["index"][1]
+
+    cell_width = image.width // GRID_NUM
+    cell_height = image.height // GRID_NUM
+
+    grid_x = x // cell_width
+    grid_y = y // cell_height
+    return grid_x, grid_y, cell_width, cell_height
+
+
+def highlight_grid(evt: gr.EventData, image_name):
+    image = data_dict[image_name]["image"]
+    grid_x, grid_y, cell_width, cell_height = get_grid_loc(evt, image)
+
+    highlighted_image = image.copy()
+    draw = ImageDraw.Draw(highlighted_image)
+    box = [
+        grid_x * cell_width,
+        grid_y * cell_height,
+        (grid_x + 1) * cell_width,
+        (grid_y + 1) * cell_height,
+    ]
+    draw.rectangle(box, outline="red", width=3)
+
+    return highlighted_image
+
+
+def load_image(img_name):
+    return Image.open(data_dict[img_name]["image_path"]).resize(
+        (IMAGE_SIZE, IMAGE_SIZE)
+    )
+
+
+def plot_activations(
+    all_activation,
+    tile_activations=None,
+    grid_x=None,
+    grid_y=None,
+    top_k=5,
+    colors=("blue", "cyan"),
+    model_name="CLIP",
+):
+    fig = go.Figure()
+
+    def _add_scatter_with_annotation(fig, activations, model_name, color, label):
+        fig.add_trace(
+            go.Scatter(
+                x=np.arange(len(activations)),
+                y=activations,
+                mode="lines",
+                name=label,
+                line=dict(color=color, dash="solid"),
+                showlegend=True,
+            )
+        )
+        top_neurons = np.argsort(activations)[::-1][:top_k]
+        for idx in top_neurons:
+            fig.add_annotation(
+                x=idx,
+                y=activations[idx],
+                text=str(idx),
+                showarrow=True,
+                arrowhead=2,
+                ax=0,
+                ay=-15,
+                arrowcolor=color,
+                opacity=0.7,
+            )
+        return fig
+
+    label = f"{model_name.split('-')[-0]} Image-level"
+    fig = _add_scatter_with_annotation(
+        fig, all_activation, model_name, colors[0], label
+    )
+    if tile_activations is not None:
+        label = f"{model_name.split('-')[-0]} Tile ({grid_x}, {grid_y})"
+        fig = _add_scatter_with_annotation(
+            fig, tile_activations, model_name, colors[1], label
+        )
+
+    fig.update_layout(
+        title="Activation Distribution",
+        xaxis_title="SAE latent index",
+        yaxis_title="Activation Value",
+        template="plotly_white",
+    )
+    fig.update_layout(
+        legend=dict(orientation="h", yanchor="middle", y=0.5, xanchor="center", x=0.5)
+    )
+
+    return fig
+
+
+def get_activations(evt: gr.EventData, selected_image: str, model_name: str, colors):
+    activation = get_activation_distribution(selected_image, model_name)
+    all_activation = activation.mean(0)
+
+    tile_activations = None
+    grid_x = None
+    grid_y = None
+
+    if evt is not None:
+        if evt._data is not None:
+            image = data_dict[selected_image]["image"]
+            grid_x, grid_y, cell_width, cell_height = get_grid_loc(evt, image)
+            token_idx = grid_y * GRID_NUM + grid_x + 1
+            tile_activations = activation[token_idx]
+
+    fig = plot_activations(
+        all_activation,
+        tile_activations,
+        grid_x,
+        grid_y,
+        top_k=5,
+        model_name=model_name,
+        colors=colors,
+    )
+    return fig
+
+
+def plot_activation_distribution(
+    evt: gr.EventData, selected_image: str, model_name: str
+):
+    fig = make_subplots(
+        rows=2,
+        cols=1,
+        shared_xaxes=True,
+        subplot_titles=["CLIP Activation", f"{model_name} Activation"],
+    )
+
+    fig_clip = get_activations(
+        evt, selected_image, "CLIP", colors=("#00b4d8", "#90e0ef")
+    )
+    fig_maple = get_activations(
+        evt, selected_image, model_name, colors=("#ff5a5f", "#ffcad4")
+    )
+
+    def _attach_fig(fig, sub_fig, row, col, yref):
+        for trace in sub_fig.data:
+            fig.add_trace(trace, row=row, col=col)
+
+        for annotation in sub_fig.layout.annotations:
+            annotation.update(yref=yref)
+            fig.add_annotation(annotation)
+        return fig
+
+    fig = _attach_fig(fig, fig_clip, row=1, col=1, yref="y1")
+    fig = _attach_fig(fig, fig_maple, row=2, col=1, yref="y2")
+
+    fig.update_xaxes(title_text="SAE Latent Index", row=2, col=1)
+    fig.update_xaxes(title_text="SAE Latent Index", row=1, col=1)
+    fig.update_yaxes(title_text="Activation Value", row=1, col=1)
+    fig.update_yaxes(title_text="Activation Value", row=2, col=1)
+    fig.update_layout(
+        # height=500,
+        # title="Activation Distributions",
+        template="plotly_white",
+        showlegend=True,
+        legend=dict(orientation="h", yanchor="bottom", y=-0.2, xanchor="center", x=0.5),
+        margin=dict(l=20, r=20, t=40, b=20),
+    )
+
+    return fig
+
+
+# def get_segmask(selected_image, slider_value, model_type):
+#     image = data_dict[selected_image]["image"]
+#     sae_act = get_data(selected_image, model_type)[0]
+#     temp = sae_act[:, slider_value]
+#     try:
+#         mask = torch.Tensor(temp[1:,].reshape(14, 14)).view(1, 1, 14, 14)
+#     except Exception as e:
+#         print(sae_act.shape, slider_value)
+#     mask = torch.nn.functional.interpolate(mask, (image.height, image.width))[0][
+#         0
+#     ].numpy()
+#     mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-10)
+
+#     base_opacity = 30
+#     image_array = np.array(image)[..., :3]
+#     rgba_overlay = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
+#     rgba_overlay[..., :3] = image_array[..., :3]
+
+#     darkened_image = (image_array[..., :3] * (base_opacity / 255)).astype(np.uint8)
+#     rgba_overlay[mask == 0, :3] = darkened_image[mask == 0]
+#     rgba_overlay[..., 3] = 255  # Fully opaque
+
+#     return rgba_overlay
+
+
+# def get_top_images(slider_value, toggle_btn):
+#     def _get_images(dataset_path):
+#         top_image_paths = [
+#             os.path.join(dataset_path, "imagenet", f"{slider_value}.jpg"),
+#             os.path.join(dataset_path, "imagenet-sketch", f"{slider_value}.jpg"),
+#             os.path.join(dataset_path, "caltech101", f"{slider_value}.jpg"),
+#         ]
+#         top_images = [
+#             (
+#                 Image.open(path)
+#                 if os.path.exists(path)
+#                 else Image.new("RGB", (256, 256), (255, 255, 255))
+#             )
+#             for path in top_image_paths
+#         ]
+#         return top_images
+
+#     if toggle_btn:
+#         top_images = _get_images("./data/top_images_masked")
+#     else:
+#         top_images = _get_images("./data/top_images")
+#     return top_images
+
+
+def show_activation_heatmap(selected_image, slider_value, model_type, toggle_btn=False):
+    slider_value = int(slider_value.split("-")[-1])
+    rgba_overlay = get_segmask(selected_image, slider_value, model_type)
+    top_images = get_top_images(slider_value, toggle_btn)
+
+    act_values = []
+    for dataset in ["imagenet", "imagenet-sketch", "caltech101"]:
+        act_value = sae_data_dict["mean_act_values"][dataset][slider_value, :5]
+        act_value = [str(round(value, 3)) for value in act_value]
+        act_value = " | ".join(act_value)
+        out = f"#### Activation values: {act_value}"
+        act_values.append(out)
+
+    return rgba_overlay, top_images, act_values
+
+
+def show_activation_heatmap_clip(selected_image, slider_value, toggle_btn):
+    rgba_overlay, top_images, act_values = show_activation_heatmap(
+        selected_image, slider_value, "CLIP", toggle_btn
+    )
+    sleep(0.1)
+    return (
+        rgba_overlay,
+        top_images[0],
+        top_images[1],
+        top_images[2],
+        act_values[0],
+        act_values[1],
+        act_values[2],
+    )
+
+
+def show_activation_heatmap_maple(selected_image, slider_value, model_name):
+    slider_value = int(slider_value.split("-")[-1])
+    rgba_overlay = get_segmask(selected_image, slider_value, model_name)
+    sleep(0.1)
+    return rgba_overlay
+
+
+def get_init_radio_options(selected_image, model_name):
+    clip_neuron_dict = {}
+    maple_neuron_dict = {}
+
+    def _get_top_actvation(selected_image, model_name, neuron_dict, top_k=5):
+        activations = get_activation_distribution(selected_image, model_name).mean(0)
+        top_neurons = list(np.argsort(activations)[::-1][:top_k])
+        for top_neuron in top_neurons:
+            neuron_dict[top_neuron] = activations[top_neuron]
+        sorted_dict = dict(
+            sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True)
+        )
+        return sorted_dict
+
+    clip_neuron_dict = _get_top_actvation(selected_image, "CLIP", clip_neuron_dict)
+    maple_neuron_dict = _get_top_actvation(
+        selected_image, model_name, maple_neuron_dict
+    )
+
+    radio_choices = get_radio_names(clip_neuron_dict, maple_neuron_dict)
+
+    return radio_choices
+
+
+def get_radio_names(clip_neuron_dict, maple_neuron_dict):
+    clip_keys = list(clip_neuron_dict.keys())
+    maple_keys = list(maple_neuron_dict.keys())
+
+    common_keys = list(set(clip_keys).intersection(set(maple_keys)))
+    clip_only_keys = list(set(clip_keys) - (set(maple_keys)))
+    maple_only_keys = list(set(maple_keys) - (set(clip_keys)))
+
+    common_keys.sort(
+        key=lambda x: max(clip_neuron_dict[x], maple_neuron_dict[x]), reverse=True
+    )
+    clip_only_keys.sort(reverse=True)
+    maple_only_keys.sort(reverse=True)
+
+    out = []
+    out.extend([f"common-{i}" for i in common_keys[:5]])
+    out.extend([f"CLIP-{i}" for i in clip_only_keys[:5]])
+    out.extend([f"MaPLE-{i}" for i in maple_only_keys[:5]])
+
+    return out
+
+
+def update_radio_options(evt: gr.EventData, selected_image, model_name):
+    def _sort_and_save_top_k(activations, neuron_dict, top_k=5):
+        top_neurons = list(np.argsort(activations)[::-1][:top_k])
+        for top_neuron in top_neurons:
+            neuron_dict[top_neuron] = activations[top_neuron]
+
+    def _get_top_actvation(evt, selected_image, model_name, neuron_dict):
+        all_activation = get_activation_distribution(selected_image, model_name)
+        image_activation = all_activation.mean(0)
+        _sort_and_save_top_k(image_activation, neuron_dict)
+
+        if evt is not None:
+            if evt._data is not None and isinstance(evt._data["index"], list):
+                image = data_dict[selected_image]["image"]
+                grid_x, grid_y, cell_width, cell_height = get_grid_loc(evt, image)
+                token_idx = grid_y * GRID_NUM + grid_x + 1
+                tile_activations = all_activation[token_idx]
+                _sort_and_save_top_k(tile_activations, neuron_dict)
+
+        sorted_dict = dict(
+            sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True)
+        )
+        return sorted_dict
+
+    clip_neuron_dict = {}
+    maple_neuron_dict = {}
+    clip_neuron_dict = _get_top_actvation(evt, selected_image, "CLIP", clip_neuron_dict)
+    maple_neuron_dict = _get_top_actvation(
+        evt, selected_image, model_name, maple_neuron_dict
+    )
+
+    clip_keys = list(clip_neuron_dict.keys())
+    maple_keys = list(maple_neuron_dict.keys())
+
+    common_keys = list(set(clip_keys).intersection(set(maple_keys)))
+    clip_only_keys = list(set(clip_keys) - (set(maple_keys)))
+    maple_only_keys = list(set(maple_keys) - (set(clip_keys)))
+
+    common_keys.sort(
+        key=lambda x: max(clip_neuron_dict[x], maple_neuron_dict[x]), reverse=True
+    )
+    clip_only_keys.sort(reverse=True)
+    maple_only_keys.sort(reverse=True)
+
+    out = []
+    out.extend([f"common-{i}" for i in common_keys[:5]])
+    out.extend([f"CLIP-{i}" for i in clip_only_keys[:5]])
+    out.extend([f"MaPLE-{i}" for i in maple_only_keys[:5]])
+
+    radio_choices = gr.Radio(
+        choices=out, label="Top activating SAE latent", value=out[0]
+    )
+    sleep(0.1)
+    return radio_choices
+
+
+def update_markdown(option_value):
+    latent_idx = int(option_value.split("-")[-1])
+    out_1 = f"## Segmentation mask for the selected SAE latent - {latent_idx}"
+    out_2 = f"## Top reference images for the selected SAE latent - {latent_idx}"
+    return out_1, out_2
+
+
+def get_data(image_name, model_name):
+    pkl_root = "./data/out"
+    data_dir = f"{pkl_root}/{model_name}/{image_name}.pkl.gz"
+    with gzip.open(data_dir, "rb") as f:
+        data = pickle.load(f)
+        out = data
+
+    return out
+
+
+def update_all(selected_image, slider_value, toggle_btn, model_name):
+    (
+        seg_mask_display,
+        top_image_1,
+        top_image_2,
+        top_image_3,
+        act_value_1,
+        act_value_2,
+        act_value_3,
+    ) = show_activation_heatmap_clip(selected_image, slider_value, toggle_btn)
+    seg_mask_display_maple = show_activation_heatmap_maple(
+        selected_image, slider_value, model_name
+    )
+    markdown_display, markdown_display_2 = update_markdown(slider_value)
+
+    return (
+        seg_mask_display,
+        seg_mask_display_maple,
+        top_image_1,
+        top_image_2,
+        top_image_3,
+        act_value_1,
+        act_value_2,
+        act_value_3,
+        markdown_display,
+        markdown_display_2,
+    )
+
+
+def load_all_data(image_root, pkl_root):
+    image_files = glob(f"{image_root}/*")
+    data_dict = {}
+    for image_file in image_files:
+        image_name = os.path.basename(image_file).split(".")[0]
+        if image_file not in data_dict:
+            data_dict[image_name] = {
+                "image": Image.open(image_file).resize((IMAGE_SIZE, IMAGE_SIZE)),
+                "image_path": image_file,
+            }
+
+    sae_data_dict = {}
+    with open("./data/sae_data/mean_acts.pkl", "rb") as f:
+        data = pickle.load(f)
+        sae_data_dict["mean_acts"] = data
+
+    sae_data_dict["mean_act_values"] = {}
+    for dataset in ["imagenet", "imagenet-sketch", "caltech101"]:
+        with gzip.open(f"./data/sae_data/mean_act_values_{dataset}.pkl.gz", "rb") as f:
+            data = pickle.load(f)
+            sae_data_dict["mean_act_values"][dataset] = data
+
+    return data_dict, sae_data_dict
+
+
+def preload_all_model_data():
+    """Preload all model data into memory at startup"""
+    print("Preloading model data...")
+    for image_name in data_dict.keys():
+        for model_name in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
+            try:
+                data = get_data(image_name, model_name)
+                cache_key = f"{model_name}_{image_name}"
+                _CACHE['model_data'][cache_key] = data
+            except Exception as e:
+                print(f"Error preloading {cache_key}: {e}")
+
+# Add to initialization
+preload_all_model_data()
+
+def precompute_activations():
+    """Precompute and cache common activation patterns"""
+    print("Precomputing activations...")
+    for image_name in data_dict.keys():
+        for model_name in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
+            activation = get_activation_distribution(image_name, model_name)
+            cache_key = f"activation_{model_name}_{image_name}"
+            _CACHE['precomputed_activations'][cache_key] = activation.mean(0)
+
+# Add to _CACHE initialization
+_CACHE['precomputed_activations'] = {}
+
+# Add to initialization
+precompute_activations()
+
+def precompute_segmasks():
+    """Precompute common segmentation masks"""
+    print("Precomputing segmentation masks...")
+    for image_name in data_dict.keys():
+        for model_type in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
+            for slider_value in range(0, 100):  # Adjust range as needed
+                try:
+                    mask = get_segmask(image_name, slider_value, model_type)
+                    cache_key = f"{image_name}_{slider_value}_{model_type}"
+                    _CACHE['segmasks'][cache_key] = mask
+                except Exception as e:
+                    print(f"Error precomputing mask {cache_key}: {e}")
+
+# Add to initialization
+precompute_segmasks()
+
+
+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],
+            )
+
+            model_selector.change(
+                fn=update_radio_options,
+                inputs=[image_selector, model_selector],
+                outputs=[radio_choices],
+            )
+
+            image_selector.select(
+                fn=update_radio_options,
+                inputs=[image_selector, model_selector],
+                outputs=[radio_choices],
+            )
+
+        radio_choices.change(
+            fn=update_all,
+            inputs=[image_selector, radio_choices, toggle_btn, model_selector],
+            outputs=[
+                seg_mask_display,
+                seg_mask_display_maple,
+                top_image_1,
+                top_image_2,
+                top_image_3,
+                act_value_1,
+                act_value_2,
+                act_value_3,
+                markdown_display,
+                markdown_display_2,
+            ],
+        )
+
+        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,
+            ],
+        )
+
+    # Launch the app
+    # demo.queue()
+    # demo.launch()
+
+    
+# if __name__ == "__main__":
+#     demo.queue()  # Enable queuing for better handling of concurrent users
+#     demo.launch(
+#         server_name="0.0.0.0",  # Allow external access
+#         server_port=7860,
+#         share=False,  # Set to True if you want to create a public URL
+#         show_error=True,
+#         # Optimize concurrency
+#         max_threads=8,  # Adjust based on your CPU cores
+#     )
+
+if __name__ == "__main__":
+    import psutil
+    
+    # Get system memory info
+    mem = psutil.virtual_memory()
+    total_ram_gb = mem.total / (1024**3)
+    
+    # Configure cache sizes based on available RAM
+    cache_size = int(total_ram_gb * 100)  # Rough estimate: 100 entries per GB
+    
+    # Memory monitoring function
+    def monitor_memory_usage():
+        """Monitor and log memory usage"""
+        process = psutil.Process()
+        mem_info = process.memory_info()
+        print(f"""
+        Memory Usage:
+        - RSS: {mem_info.rss / (1024**2):.2f} MB
+        - VMS: {mem_info.vms / (1024**2):.2f} MB
+        - Cache Size: {len(_CACHE['model_data'])} entries
+        """)
+
+    # Start periodic monitoring
+    def start_memory_monitor():
+        threading.Timer(300.0, start_memory_monitor).start()  # Every 5 minutes
+        monitor_memory_usage()
+
+    # Start the monitoring
+    import threading
+    start_memory_monitor()
+
+    # Launch the app with memory-optimized settings
+    demo.queue(max_size=min(20, int(total_ram_gb)))  # Scale queue with RAM
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True,
+        max_threads=min(16, psutil.cpu_count()),  # Scale threads with CPU
+        websocket_ping_timeout=60,
+        preventive_refresh=True,
+        memory_limit_mb=int(total_ram_gb * 1024 * 0.8)  # Use up to 80% of RAM
+    )