app.py

import gzip
import os
import pickle
from glob import glob
from functools import lru_cache
import concurrent.futures
import threading
import time

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

# Constants
IMAGE_SIZE = 400
DATASET_LIST = ["imagenet", "oxford_flowers", "ucf101", "caltech101", "dtd", "eurosat"]
GRID_NUM = 14
pkl_root = "./data/out"

# Global cache for preloaded data
preloaded_data = {}
data_dict = {}
sae_data_dict = {}
activation_cache = {}
segmask_cache = {}
top_images_cache = {}

# Thread lock for thread-safe operations
data_lock = threading.Lock()

# Load data more efficiently
def load_all_data(image_root, pkl_root):
    """Load all necessary data with optimized caching"""
    # Load image data
    image_files = glob(f"{image_root}/*")
    data_dict = {}
    
    # Use thread pool for parallel image loading
    def load_image_data(image_file):
        image_name = os.path.basename(image_file).split(".")[0]
        # Only load thumbnail for initial display, load full image on demand
        thumbnail = Image.open(image_file).resize((IMAGE_SIZE, IMAGE_SIZE))
        return image_name, {
            "image": thumbnail,
            "image_path": image_file,
        }
    
    # Load images in parallel
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
        results = executor.map(load_image_data, image_files)
        for image_name, data in results:
            data_dict[image_name] = data
    
    # Load SAE data with minimal processing
    sae_data_dict = {}
    
    # Load mean acts only once
    with open("./data/sae_data/mean_acts.pkl", "rb") as f:
        sae_data_dict["mean_acts"] = pickle.load(f)
    
    # Update all components when radio selection changes
        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,
            ],
            _js="""
            function(img, radio, toggle, model) {
                // Add a small delay to prevent rapid UI updates
                clearTimeout(window._radioTimeout);
                return new Promise((resolve) => {
                    window._radioTimeout = setTimeout(() => {
                        resolve([img, radio, toggle, model]);
                    }, 100);
                });
            }
            """
        )
        
        # Update components when toggle button changes
        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,
            ],
            _js="""
            function(img, radio, toggle) {
                // Add a small delay to prevent rapid UI updates
                clearTimeout(window._toggleTimeout);
                return new Promise((resolve) => {
                    window._toggleTimeout = setTimeout(() => {
                        resolve([img, radio, toggle]);
                    }, 100);
                });
            }
            """
        )

    # Initialize UI with default values
    default_options = get_init_radio_options(default_image_name, model_options[0])
    if default_options:
        default_option = default_options[0]
        
        # Set initial values to avoid blank UI at start
        gr.on(
            gr.Blocks.load,
            fn=lambda: update_all(
                default_image_name, 
                default_option,
                False,
                model_options[0]
            ),
            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,
            ],
        )

    # Add a status indicator to show processing state
    status_indicator = gr.Markdown("Status: Ready")
    
    # Add a refresh button to manually reload data if needed
    refresh_btn = gr.Button("Refresh Data")
    
    def reload_data():
        global data_dict, sae_data_dict
        
        # Update status
        yield "Status: Reloading data..."
        
        # Reload data
        try:
            data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root)
            yield "Status: Data reloaded successfully!"
        except Exception as e:
            yield f"Status: Error reloading data - {str(e)}"
    
    refresh_btn.click(
        fn=reload_data,
        inputs=[],
        outputs=[status_indicator],
        queue=False
    )
    
    # Launch app with optimized settings
    demo.queue(concurrency_count=3, max_size=10)  # Balanced concurrency for better performance
    
    # Add startup message
    print("Starting visualization application...")
    print(f"Loaded {len(data_dict)} images and {len(sae_data_dict)} datasets")
    
    # Launch with proper error handling
    demo.launch(
        share=False,        # Don't share publicly
        debug=False,        # Disable debug mode for production
        show_error=True,    # Show errors for debugging
        quiet=False,        # Show startup messages
        favicon_path=None,  # Default favicon
        server_port=None,   # Use default port
        server_name=None,   # Bind to all interfaces
        height=None,        # Use default height
        width=None,         # Use default width
        enable_queue=True,  # Enable queue for better performance
    ) dictionary for dataset values
    sae_data_dict["mean_act_values"] = {}
    
    # Load dataset values in parallel
    def load_dataset_values(dataset):
        with gzip.open(f"./data/sae_data/mean_act_values_{dataset}.pkl.gz", "rb") as f:
            return dataset, pickle.load(f)
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
        futures = [
            executor.submit(load_dataset_values, dataset)
            for dataset in ["imagenet", "imagenet-sketch", "caltech101"]
        ]
        for future in concurrent.futures.as_completed(futures):
            dataset, data = future.result()
            sae_data_dict["mean_act_values"][dataset] = data
    
    return data_dict, sae_data_dict

# Cache activation data with LRU cache
@lru_cache(maxsize=32)
def preload_activation(image_name, model_name):
    """Preload and cache activation data for a specific image and model"""
    image_file = f"{pkl_root}/{model_name}/{image_name}.pkl.gz"
    
    try:
        with gzip.open(image_file, "rb") as f:
            return pickle.load(f)
    except Exception as e:
        print(f"Error loading {image_file}: {e}")
        return None

# Get activation with caching
def get_data(image_name, model_type):
    """Get activation data with caching for better performance"""
    cache_key = f"{image_name}_{model_type}"
    
    with data_lock:
        if cache_key not in activation_cache:
            activation_cache[cache_key] = preload_activation(image_name, model_type)
    
    return activation_cache[cache_key]

def get_activation_distribution(image_name, model_type):
    """Get activation distribution with noise filtering"""
    activation = get_data(image_name, model_type)
    
    if activation is None:
        # Return empty tensor if data loading failed
        return torch.zeros((GRID_NUM * GRID_NUM + 1, 1000))
    
    activation = activation[0]
    
    # Filter out noisy features
    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 grid location from click event"""
    # 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, image_name):
    """Highlight grid cell on click"""
    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):
    """Load image by name"""
    return data_dict[img_name]["image"]

# Optimized plotting with less annotations
def plot_activations(
    all_activation,
    tile_activations=None,
    grid_x=None,
    grid_y=None,
    top_k=5,
    colors=("blue", "cyan"),
    model_name="CLIP",
):
    """Plot activations with optimized rendering"""
    fig = go.Figure()
    
    def _add_scatter_with_annotation(fig, activations, model_name, color, label):
        # Only plot non-zero values to reduce points
        non_zero_indices = np.where(np.abs(activations) > 1e-5)[0]
        if len(non_zero_indices) == 0:
            # If all values are near zero, use full array
            non_zero_indices = np.arange(len(activations))
        
        fig.add_trace(
            go.Scatter(
                x=non_zero_indices,
                y=activations[non_zero_indices],
                mode="lines",
                name=label,
                line=dict(color=color, dash="solid"),
                showlegend=True,
            )
        )
        
        # Only annotate the top_k activations
        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('-')[-1]} 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('-')[-1]} Tile ({grid_x}, {grid_y})"
        fig = _add_scatter_with_annotation(
            fig, tile_activations, model_name, colors[1], label
        )
    
    # Optimize layout with minimal settings
    fig.update_layout(
        title="Activation Distribution",
        xaxis_title="SAE latent index",
        yaxis_title="Activation Value",
        template="plotly_white",
        legend=dict(orientation="h", yanchor="middle", y=0.5, xanchor="center", x=0.5),
    )
    
    return fig

def get_activations(evt, selected_image, model_name, colors):
    """Get activations for plotting"""
    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 and 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
        # Ensure token_idx is within bounds
        if token_idx < activation.shape[0]:
            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

# Cache plot results
@lru_cache(maxsize=16)
def plot_activation_distribution(evt_data, selected_image, model_name):
    """Plot activation distribution with caching"""
    # Convert event data to hashable format for caching
    if evt_data is not None:
        evt = type('obj', (object,), {'_data': evt_data})
    else:
        evt = None
    
    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")
    
    # Optimize layout with minimal settings
    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(
        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

# Cache segmentation masks
@lru_cache(maxsize=32)
def get_segmask(selected_image, slider_value, model_type):
    """Generate segmentation mask with caching"""
    try:
        # Check if image exists
        if selected_image not in data_dict:
            print(f"Image {selected_image} not found in data dictionary")
            # Return blank mask with IMAGE_SIZE dimensions
            return np.zeros((IMAGE_SIZE, IMAGE_SIZE, 4), dtype=np.uint8)
        
        # Use cache if available
        cache_key = f"{selected_image}_{slider_value}_{model_type}"
        with data_lock:
            if cache_key in segmask_cache:
                return segmask_cache[cache_key]
        
        # Get image
        image = data_dict[selected_image]["image"]
        
        # Get activation data
        sae_act = get_data(selected_image, model_type)
        
        if sae_act is None:
            # Return blank mask if data loading failed
            return np.zeros((image.height, image.width, 4), dtype=np.uint8)
        
        # Handle array shape issues
        try:
            # Check array shape and dimensions
            if isinstance(sae_act, tuple) and len(sae_act) > 0:
                # First element of tuple
                act_data = sae_act[0]
            else:
                # Direct array
                act_data = sae_act
            
            # Check if slider_value is within bounds
            if slider_value >= act_data.shape[1]:
                print(f"Slider value {slider_value} out of bounds for activation shape {act_data.shape}")
                return np.zeros((image.height, image.width, 4), dtype=np.uint8)
            
            # Get activation for specific latent
            temp = act_data[:, slider_value]
            
            # Skip first token (CLS token) and reshape to grid
            if len(temp) > 1:  # Ensure we have enough tokens
                mask = torch.Tensor(temp[1:].reshape(GRID_NUM, GRID_NUM)).view(1, 1, GRID_NUM, GRID_NUM)
                
                # Upsample to image dimensions
                mask = torch.nn.functional.interpolate(mask, (image.height, image.width))[0][0].numpy()
                
                # Normalize mask values between 0 and 1
                mask_min, mask_max = mask.min(), mask.max()
                if mask_max > mask_min:  # Avoid division by zero
                    mask = (mask - mask_min) / (mask_max - mask_min)
                else:
                    mask = np.zeros_like(mask)
            else:
                # Not enough tokens
                print(f"Not enough tokens in activation data: {len(temp)}")
                return np.zeros((image.height, image.width, 4), dtype=np.uint8)
                
        except Exception as e:
            print(f"Error processing activation data: {e}")
            print(f"Shape info - sae_act: {type(sae_act)}, slider_value: {slider_value}")
            return np.zeros((image.height, image.width, 4), dtype=np.uint8)
        
        # Create RGBA overlay
        try:
            # Set base opacity for darkened areas
            base_opacity = 30
            
            # Convert image to numpy array
            image_array = np.array(image)
            
            # Handle grayscale images
            if len(image_array.shape) == 2:
                # Convert grayscale to RGB
                image_array = np.stack([image_array] * 3, axis=-1)
            elif image_array.shape[2] == 4:
                # Use only RGB channels
                image_array = image_array[..., :3]
            
            # Create overlay
            rgba_overlay = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
            rgba_overlay[..., :3] = image_array
            
            # Use vectorized operations for better performance
            darkened_image = (image_array * (base_opacity / 255)).astype(np.uint8)
            
            # Create mask for darkened areas
            mask_threshold = 0.1  # Adjust threshold if needed
            mask_zero = mask < mask_threshold
            
            # Apply darkening only to low-activation areas
            rgba_overlay[mask_zero, :3] = darkened_image[mask_zero]
            
            # Set alpha channel
            rgba_overlay[..., 3] = 255  # Fully opaque
            
            # Cache result for future use
            with data_lock:
                segmask_cache[cache_key] = rgba_overlay
            
            return rgba_overlay
            
        except Exception as e:
            print(f"Error creating overlay: {e}")
            return np.zeros((image.height, image.width, 4), dtype=np.uint8)
    
    except Exception as e:
        print(f"Unexpected error in get_segmask: {e}")
        # Return a blank image of standard size
        return np.zeros((IMAGE_SIZE, IMAGE_SIZE, 4), dtype=np.uint8)

# Cache top images
@lru_cache(maxsize=32)
def get_top_images(slider_value, toggle_btn):
    """Get top images with caching"""
    cache_key = f"{slider_value}_{toggle_btn}"
    
    if cache_key in top_images_cache:
        return top_images_cache[cache_key]
    
    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 = []
        for path in top_image_paths:
            if os.path.exists(path):
                top_images.append(Image.open(path))
            else:
                top_images.append(Image.new("RGB", (256, 256), (255, 255, 255)))
        
        return top_images
    
    if toggle_btn:
        top_images = _get_images("./data/top_images_masked")
    else:
        top_images = _get_images("./data/top_images")
    
    # Cache result
    top_images_cache[cache_key] = top_images
    
    return top_images

def show_activation_heatmap(selected_image, slider_value, model_type, toggle_btn=False):
    """Show activation heatmap with optimized processing"""
    try:
        # Parse slider value safely
        if not slider_value:
            # Fallback to the first option if no slider value
            radio_options = get_init_radio_options(selected_image, model_type)
            if not radio_options:
                # Create placeholder data if no options available
                return (
                    np.zeros((IMAGE_SIZE, IMAGE_SIZE, 4), dtype=np.uint8),
                    [Image.new("RGB", (256, 256), (255, 255, 255)) for _ in range(3)],
                    ["#### Activation values: No data available"] * 3
                )
            slider_value = radio_options[0]
        
        # Extract the integer value
        try:
            slider_value_int = int(slider_value.split("-")[-1])
        except (ValueError, IndexError):
            print(f"Error parsing slider value: {slider_value}")
            slider_value_int = 0
        
        # Process in parallel with thread pool and add timeout
        results = []
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            # Start both tasks
            segmask_future = executor.submit(get_segmask, selected_image, slider_value_int, model_type)
            top_images_future = executor.submit(get_top_images, slider_value_int, toggle_btn)
            
            # Get results with timeout to prevent hanging
            try:
                rgba_overlay = segmask_future.result(timeout=5)
            except (concurrent.futures.TimeoutError, Exception) as e:
                print(f"Error or timeout generating segmentation mask: {e}")
                rgba_overlay = np.zeros((IMAGE_SIZE, IMAGE_SIZE, 4), dtype=np.uint8)
            
            try:
                top_images = top_images_future.result(timeout=5)
            except (concurrent.futures.TimeoutError, Exception) as e:
                print(f"Error or timeout getting top images: {e}")
                top_images = [Image.new("RGB", (256, 256), (255, 255, 255)) for _ in range(3)]
        
        # Prepare activation values with error handling
        act_values = []
        for dataset in ["imagenet", "imagenet-sketch", "caltech101"]:
            try:
                if dataset in sae_data_dict["mean_act_values"]:
                    values = sae_data_dict["mean_act_values"][dataset]
                    if slider_value_int < values.shape[0]:
                        act_value = values[slider_value_int, :5]
                        act_value = [str(round(value, 3)) for value in act_value]
                        act_value = " | ".join(act_value)
                        out = f"#### Activation values: {act_value}"
                    else:
                        out = f"#### Activation values: Index out of range"
                else:
                    out = f"#### Activation values: Dataset not available"
            except Exception as e:
                print(f"Error getting activation values for {dataset}: {e}")
                out = f"#### Activation values: Error retrieving data"
            
            act_values.append(out)
        
        return rgba_overlay, top_images, act_values
    
    except Exception as e:
        print(f"Error in show_activation_heatmap: {e}")
        # Return placeholder data in case of error
        return (
            np.zeros((IMAGE_SIZE, IMAGE_SIZE, 4), dtype=np.uint8),
            [Image.new("RGB", (256, 256), (255, 255, 255)) for _ in range(3)],
            ["#### Activation values: Error occurred"] * 3
        )

def show_activation_heatmap_clip(selected_image, slider_value, toggle_btn):
    """Show CLIP activation heatmap"""
    rgba_overlay, top_images, act_values = show_activation_heatmap(
        selected_image, slider_value, "CLIP", toggle_btn
    )
    
    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):
    """Show MaPLE activation heatmap"""
    slider_value_int = int(slider_value.split("-")[-1])
    rgba_overlay = get_segmask(selected_image, slider_value_int, model_name)
    
    return rgba_overlay

# Optimize radio options generation
def get_init_radio_options(selected_image, model_name):
    """Get initial radio options with optimized processing"""
    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
    
    # Process in parallel
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        future_clip = executor.submit(_get_top_actvation, selected_image, "CLIP", {})
        future_maple = executor.submit(_get_top_actvation, selected_image, model_name, {})
        
        clip_neuron_dict = future_clip.result()
        maple_neuron_dict = future_maple.result()
    
    radio_choices = get_radio_names(clip_neuron_dict, maple_neuron_dict)
    
    return radio_choices

def get_radio_names(clip_neuron_dict, maple_neuron_dict):
    """Get radio button names based on neuron activations"""
    clip_keys = list(clip_neuron_dict.keys())
    maple_keys = list(maple_neuron_dict.keys())
    
    # Use set operations for better performance
    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))
    
    # Sort keys by activation values
    common_keys.sort(
        key=lambda x: max(clip_neuron_dict.get(x, 0), maple_neuron_dict.get(x, 0)), 
        reverse=True
    )
    clip_only_keys.sort(key=lambda x: clip_neuron_dict.get(x, 0), reverse=True)
    maple_only_keys.sort(key=lambda x: maple_neuron_dict.get(x, 0), reverse=True)
    
    # Limit number of choices to improve performance
    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, selected_image, model_name):
    """Update radio options based on user interaction"""
    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)
        
        # Get top activations from image-level
        top_neurons = list(np.argsort(image_activation)[::-1][:5])
        for top_neuron in top_neurons:
            neuron_dict[top_neuron] = image_activation[top_neuron]
        
        # Get top activations from tile-level if available
        if evt is not None and 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
            
            # Ensure token_idx is within bounds
            if token_idx < all_activation.shape[0]:
                tile_activations = all_activation[token_idx]
                top_tile_neurons = list(np.argsort(tile_activations)[::-1][:5])
                for top_neuron in top_tile_neurons:
                    neuron_dict[top_neuron] = max(
                        neuron_dict.get(top_neuron, 0), 
                        tile_activations[top_neuron]
                    )
        
        # Sort by activation value
        return dict(sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True))
    
    # Process in parallel
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        future_clip = executor.submit(_get_top_actvation, evt, selected_image, "CLIP")
        future_maple = executor.submit(_get_top_actvation, evt, selected_image, model_name)
        
        clip_neuron_dict = future_clip.result()
        maple_neuron_dict = future_maple.result()
    
    # Get radio choices
    radio_choices = get_radio_names(clip_neuron_dict, maple_neuron_dict)
    
    # Create radio component
    radio = gr.Radio(
        choices=radio_choices, 
        label="Top activating SAE latent", 
        value=radio_choices[0] if radio_choices else None
    )
    
    return radio

def update_markdown(option_value):
    """Update markdown text"""
    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 update_all(selected_image, slider_value, toggle_btn, model_name):
    """Update all UI components in optimized way"""
    # Use a thread pool to parallelize operations
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        # Start both tasks
        clip_future = executor.submit(
            show_activation_heatmap_clip, 
            selected_image, 
            slider_value, 
            toggle_btn
        )
        
        maple_future = executor.submit(
            show_activation_heatmap_maple,
            selected_image,
            slider_value,
            model_name
        )
        
        # Get results
        (
            seg_mask_display,
            top_image_1,
            top_image_2,
            top_image_3,
            act_value_1,
            act_value_2,
            act_value_3,
        ) = clip_future.result()
        
        seg_mask_display_maple = maple_future.result()
    
    # Update markdown
    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,
    )

# Initialize data - load at startup
data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root)
default_image_name = "christmas-imagenet"

# Define UI with lazy loading
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=load_image(default_image_name),
                type="pil",
                interactive=True,
            )
            
            # Update image display when a new image is selected (with debounce)
            image_selector.change(
                fn=load_image,
                inputs=image_selector,
                outputs=image_display,
                _js="""
                function(img_name) {
                    // Simple debounce
                    clearTimeout(window._imageSelectTimeout);
                    return new Promise((resolve) => {
                        window._imageSelectTimeout = setTimeout(() => {
                            resolve(img_name);
                        }, 100);
                    });
                }
                """
            )
            
            # Handle grid highlighting
            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)",
            )
            
            # Initialize with a placeholder plot to avoid delays
            neuron_plot = gr.Plot(
                label="Neuron Activation", 
                show_label=False
            )
            
            # Add event handlers with proper data flow
            def update_plot(evt, selected_image, model_name):
                if hasattr(evt, '_data') and evt._data is not None:
                    return plot_activation_distribution(
                        tuple(map(tuple, evt._data.get('index', []))), 
                        selected_image, 
                        model_name
                    )
                return plot_activation_distribution(None, selected_image, model_name)
            
            # Load initial plot after UI is rendered
            gr.on(
                [image_selector.change, model_selector.change],
                fn=lambda img, model: plot_activation_distribution(None, img, model),
                inputs=[image_selector, model_selector],
                outputs=neuron_plot,
            )
            
            # Update plot on image click
            image_display.select(
                fn=update_plot,
                inputs=[image_selector, model_selector],
                outputs=neuron_plot,
            )

    with gr.Row():
        with gr.Column():
            # Initialize radio options
            radio_names = gr.State(value=get_init_radio_options(default_image_name, model_options[0]))
            
            # Initialize markdown displays
            markdown_display = gr.Markdown(f"## Segmentation mask for the selected SAE latent")
            
            # Initialize segmentation displays
            gr.Markdown("### Localize SAE latent activation using CLIP")
            seg_mask_display = gr.Image(type="pil", show_label=False)
            
            gr.Markdown("### Localize SAE latent activation using MaPLE")
            seg_mask_display_maple = gr.Image(type="pil", show_label=False)
        
        with gr.Column():
            gr.Markdown("## Top activating SAE latent index")
            
            # Initialize radio component
            radio_choices = gr.Radio(
                label="Top activating SAE latent",
                interactive=True,
            )
            
            # Initialize as soon as UI loads
            gr.on(
                gr.Blocks.load,
                fn=lambda: gr.Radio.update(
                    choices=get_init_radio_options(default_image_name, model_options[0]),
                    value=get_init_radio_options(default_image_name, model_options[0])[0]
                ),
                outputs=radio_choices
            )
            
            toggle_btn = gr.Checkbox(label="Show segmentation mask", value=False)
            
            markdown_display_2 = gr.Markdown(f"## Top reference images for the selected SAE latent")
            
            # Initialize image displays
            gr.Markdown("### ImageNet")
            top_image_1 = gr.Image(type="pil", label="ImageNet", show_label=False)
            act_value_1 = gr.Markdown()
            
            gr.Markdown("### ImageNet-Sketch")
            top_image_2 = gr.Image(type="pil", label="ImageNet-Sketch", show_label=False)
            act_value_2 = gr.Markdown()
            
            gr.Markdown("### Caltech101")
            top_image_3 = gr.Image(type="pil", label="Caltech101", show_label=False)
            act_value_3 = gr.Markdown()
            
            # Update radio options on image interaction
            image_display.select(
                fn=update_radio_options,
                inputs=[image_selector, model_selector],
                outputs=radio_choices,
            )
            
            # Update radio options on model change
            model_selector.change(
                fn=update_radio_options,
                inputs=[image_selector, model_selector],
                outputs=radio_choices,
            )
            
            # Update radio options on image selection
            image_selector.change(
                fn=update_radio_options,
                inputs=[image_selector, model_selector],
                outputs=radio_choices,
            )
        
        # Initialize