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': {}, 'precomputed_activations' = {} } 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}") data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root) default_image_name = "christmas-imagenet" with gr.Blocks( theme=gr.themes.Citrus(), css=""" .image-row .gr-image { margin: 0 !important; padding: 0 !important; } .image-row img { width: auto; height: 50px; } /* Set a uniform height for all images */ """, ) as demo: with gr.Row(): with gr.Column(): # Left View: Image selection and click handling gr.Markdown("## Select input image and patch on the image") image_selector = gr.Dropdown( choices=list(data_dict.keys()), value=default_image_name, label="Select Image", ) image_display = gr.Image( value=data_dict[default_image_name]["image"], type="pil", interactive=True, ) # Update image display when a new image is selected image_selector.change( fn=lambda img_name: data_dict[img_name]["image"], inputs=image_selector, outputs=image_display, ) image_display.select( fn=highlight_grid, inputs=[image_selector], outputs=[image_display] ) with gr.Column(): gr.Markdown("## SAE latent activations of CLIP and MaPLE") model_options = [f"MaPLE-{dataset_name}" for dataset_name in DATASET_LIST] model_selector = gr.Dropdown( choices=model_options, value=model_options[0], label="Select adapted model (MaPLe)", ) init_plot = plot_activation_distribution( None, default_image_name, model_options[0] ) neuron_plot = gr.Plot( label="Neuron Activation", value=init_plot, show_label=False ) image_selector.change( fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot, ) image_display.select( fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot, ) model_selector.change( fn=load_image, inputs=[image_selector], outputs=image_display ) model_selector.change( fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot, ) with gr.Row(): with gr.Column(): radio_names = get_init_radio_options(default_image_name, model_options[0]) feautre_idx = radio_names[0].split("-")[-1] markdown_display = gr.Markdown( f"## Segmentation mask for the selected SAE latent - {feautre_idx}" ) init_seg, init_tops, init_values = show_activation_heatmap( default_image_name, radio_names[0], "CLIP" ) gr.Markdown("### Localize SAE latent activation using CLIP") seg_mask_display = gr.Image(value=init_seg, type="pil", show_label=False) init_seg_maple, _, _ = show_activation_heatmap( default_image_name, radio_names[0], model_options[0] ) gr.Markdown("### Localize SAE latent activation using MaPLE") seg_mask_display_maple = gr.Image( value=init_seg_maple, type="pil", show_label=False ) with gr.Column(): gr.Markdown("## Top activating SAE latent index") radio_choices = gr.Radio( choices=radio_names, label="Top activating SAE latent", interactive=True, value=radio_names[0], ) toggle_btn = gr.Checkbox(label="Show segmentation mask", value=False) markdown_display_2 = gr.Markdown( f"## Top reference images for the selected SAE latent - {feautre_idx}" ) gr.Markdown("### ImageNet") top_image_1 = gr.Image( value=init_tops[0], type="pil", label="ImageNet", show_label=False ) act_value_1 = gr.Markdown(init_values[0]) gr.Markdown("### ImageNet-Sketch") top_image_2 = gr.Image( value=init_tops[1], type="pil", label="ImageNet-Sketch", show_label=False, ) act_value_2 = gr.Markdown(init_values[1]) gr.Markdown("### Caltech101") top_image_3 = gr.Image( value=init_tops[2], type="pil", label="Caltech101", show_label=False ) act_value_3 = gr.Markdown(init_values[2]) image_display.select( fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], ) model_selector.change( fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], ) image_selector.select( fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], ) radio_choices.change( fn=update_all, inputs=[image_selector, radio_choices, toggle_btn, model_selector], outputs=[ seg_mask_display, seg_mask_display_maple, top_image_1, top_image_2, top_image_3, act_value_1, act_value_2, act_value_3, markdown_display, markdown_display_2, ], ) toggle_btn.change( fn=show_activation_heatmap_clip, inputs=[image_selector, radio_choices, toggle_btn], outputs=[ seg_mask_display, top_image_1, top_image_2, top_image_3, act_value_1, act_value_2, act_value_3, ], ) # Launch the app # demo.queue() # demo.launch() # if __name__ == "__main__": # demo.queue() # Enable queuing for better handling of concurrent users # demo.launch( # server_name="0.0.0.0", # Allow external access # server_port=7860, # share=False, # Set to True if you want to create a public URL # show_error=True, # # Optimize concurrency # max_threads=8, # Adjust based on your CPU cores # ) if __name__ == "__main__": import psutil # Get system memory info mem = psutil.virtual_memory() total_ram_gb = mem.total / (1024**3) # Configure cache sizes based on available RAM cache_size = int(total_ram_gb * 100) # Rough estimate: 100 entries per GB # Precompute all data print("Starting precomputation...") preload_all_model_data() precompute_activations() precompute_segmasks() print("Precomputation complete!") # Memory monitoring function def monitor_memory_usage(): """Monitor and log memory usage""" process = psutil.Process() mem_info = process.memory_info() print(f""" Memory Usage: - RSS: {mem_info.rss / (1024**2):.2f} MB - VMS: {mem_info.vms / (1024**2):.2f} MB - Cache Size: {len(_CACHE['model_data'])} entries """) # Start periodic monitoring def start_memory_monitor(): threading.Timer(300.0, start_memory_monitor).start() # Every 5 minutes monitor_memory_usage() # Start the monitoring import threading start_memory_monitor() # Launch the app with memory-optimized settings demo.queue(max_size=min(20, int(total_ram_gb))) # Scale queue with RAM demo.launch( server_name="0.0.0.0", server_port=7860, share=False, show_error=True, max_threads=min(16, psutil.cpu_count()), # Scale threads with CPU websocket_ping_timeout=60, preventive_refresh=True, memory_limit_mb=int(total_ram_gb * 1024 * 0.8) # Use up to 80% of RAM )