app.py

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]:
#         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}")

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)



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}")


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()

    
    # 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(
        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
    )