test: fixed typo

app.py

@@ -155,805 +155,6 @@ def get_top_images(slider_value: int, toggle_btn: bool) -> List[Image.Image]:
         
     return _CACHE['top_images'][cache_key]
 
-# Initialize data
-data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root)
-
-
-# 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
-
-
-# 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
-    )
-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']
-
-def _load_image_file(image_file: str) -> Dict:
-    """Helper function to load a single image file."""
-    try:
-        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:
-        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]:
@@ -1413,9 +614,6 @@ def preload_all_model_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...")
@@ -1425,11 +623,7 @@ def precompute_activations():
             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"""
@@ -1444,13 +638,6 @@ def precompute_segmasks():
                 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(),
@@ -1672,6 +859,18 @@ if __name__ == "__main__":
     import threading
     start_memory_monitor()
 
+    
+    # Add to initialization
+    preload_all_model_data()
+    
+    _CACHE['precomputed_activations'] = {}
+    precompute_activations()
+    precompute_segmasks()
+    
+    data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root)
+    default_image_name = "christmas-imagenet"
+
+
     # Launch the app with memory-optimized settings
     demo.queue(max_size=min(20, int(total_ram_gb)))  # Scale queue with RAM
     demo.launch(