perf: improve delay

app.py

@@ -2,7 +2,10 @@ import gzip
 import os
 import pickle
 from glob import glob
-from time import sleep
+from functools import lru_cache
+import concurrent.futures
+import threading
+import time
 
 import gradio as gr
 import numpy as np
@@ -11,47 +14,259 @@ 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()
 
-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)
+# 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,
+            ],
+        )
 
-def get_activation_distribution(image_name: str, model_type: str):
-    activation = get_data(image_name, model_type)[0]
+    # 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: gr.EventData, image_name):
+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 = [
@@ -61,16 +276,14 @@ def highlight_grid(evt: gr.EventData, image_name):
         (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)
-    )
-
+    """Load image by name"""
+    return data_dict[img_name]["image"]
 
+# Optimized plotting with less annotations
 def plot_activations(
     all_activation,
     tile_activations=None,
@@ -80,19 +293,28 @@ def plot_activations(
     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=np.arange(len(activations)),
-                y=activations,
+                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(
@@ -107,45 +329,46 @@ def plot_activations(
                 opacity=0.7,
             )
         return fig
-
-    label = f"{model_name.split('-')[-0]} Image-level"
+    
+    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('-')[-0]} Tile ({grid_x}, {grid_y})"
+        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),
     )
-    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):
+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:
-        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
+    
+    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,
@@ -155,124 +378,291 @@ def get_activations(evt: gr.EventData, selected_image: str, model_name: str, col
         model_name=model_name,
         colors=colors,
     )
+    
     return fig
 
-
-def plot_activation_distribution(
-    evt: gr.EventData, selected_image: str, model_name: str
-):
+# 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(
-        # 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
 
-
+# Cache segmentation masks
+@lru_cache(maxsize=32)
 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]
+    """Generate segmentation mask with caching"""
     try:
-        mask = torch.Tensor(temp[1:,].reshape(14, 14)).view(1, 1, 14, 14)
+        # 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(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
-
+        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 = [
-            (
-                Image.open(path)
-                if os.path.exists(path)
-                else Image.new("RGB", (256, 256), (255, 255, 255))
-            )
-            for path in top_image_paths
-        ]
+        
+        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):
-    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
-
+    """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
     )
-    sleep(0.1)
+    
     return (
         rgba_overlay,
         top_images[0],
@@ -283,18 +673,19 @@ def show_activation_heatmap_clip(selected_image, slider_value, toggle_btn):
         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)
+    """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])
@@ -304,127 +695,138 @@ def get_init_radio_options(selected_image, model_name):
             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
-    )
-
+    
+    # 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)))
-
+    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[x], maple_neuron_dict[x]), reverse=True
+        key=lambda x: max(clip_neuron_dict.get(x, 0), maple_neuron_dict.get(x, 0)), 
+        reverse=True
     )
-    clip_only_keys.sort(reverse=True)
-    maple_only_keys.sort(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: 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):
+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)
-        _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
+        
+        # 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]
-                _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]
+                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
     )
-    sleep(0.1)
-    return radio_choices
-
+    
+    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 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
-    )
+    """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,
@@ -438,42 +840,17 @@ def update_all(selected_image, slider_value, toggle_btn, model_name):
         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
-
-
+# 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():
@@ -485,21 +862,36 @@ with gr.Blocks(
                 label="Select Image",
             )
             image_display = gr.Image(
-                value=data_dict[default_image_name]["image"],
+                value=load_image(default_image_name),
                 type="pil",
                 interactive=True,
             )
-
-            # Update image display when a new image is selected
+            
+            # Update image display when a new image is selected (with debounce)
             image_selector.change(
-                fn=lambda img_name: data_dict[img_name]["image"],
+                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]
+                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]
@@ -508,139 +900,108 @@ with gr.Blocks(
                 value=model_options[0],
                 label="Select adapted model (MaPLe)",
             )
-            init_plot = plot_activation_distribution(
-                None, default_image_name, model_options[0]
-            )
+            
+            # Initialize with a placeholder plot to avoid delays
             neuron_plot = gr.Plot(
-                label="Neuron Activation", value=init_plot, show_label=False
+                label="Neuron Activation", 
+                show_label=False
             )
-
-            image_selector.change(
-                fn=plot_activation_distribution,
+            
+            # 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=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,
+                fn=update_plot,
                 inputs=[image_selector, model_selector],
                 outputs=neuron_plot,
             )
 
     with gr.Row():
         with gr.Column():
-            radio_names = get_init_radio_options(default_image_name, model_options[0])
-
-            feautre_idx = radio_names[0].split("-")[-1]
-            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"
-            )
-
+            # 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(value=init_seg, type="pil", show_label=False)
-            init_seg_maple, _, _ = show_activation_heatmap(
-                default_image_name, radio_names[0], model_options[0]
-            )
+            seg_mask_display = gr.Image(type="pil", show_label=False)
+            
             gr.Markdown("### Localize SAE latent activation using MaPLE")
-            seg_mask_display_maple = gr.Image(
-                value=init_seg_maple, type="pil", show_label=False
-            )
-
+            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(
-                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}"
+            
+            # 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(
-                value=init_tops[0], type="pil", label="ImageNet", show_label=False
-            )
-            act_value_1 = gr.Markdown(init_values[0])
-
+            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(
-                value=init_tops[1],
-                type="pil",
-                label="ImageNet-Sketch",
-                show_label=False,
-            )
-            act_value_2 = gr.Markdown(init_values[1])
-
+            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(
-                value=init_tops[2], type="pil", label="Caltech101", show_label=False
-            )
-            act_value_3 = gr.Markdown(init_values[2])
-
+            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],
+                outputs=radio_choices,
             )
-
+            
+            # Update radio options on model change
             model_selector.change(
                 fn=update_radio_options,
                 inputs=[image_selector, model_selector],
-                outputs=[radio_choices],
+                outputs=radio_choices,
             )
-
-            image_selector.select(
+            
+            # Update radio options on image selection
+            image_selector.change(
                 fn=update_radio_options,
                 inputs=[image_selector, model_selector],
-                outputs=[radio_choices],
+                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()
+        
+        # Initialize