Update

app.py

@@ -5,6 +5,7 @@ from sae_auto_interp.features.features import upsample_mask
 import torch
 from transformers import AutoTokenizer
 from PIL import Image
+import spaces
 
 CITATION_BUTTON_TEXT = """
 @misc{zhang2024largemultimodalmodelsinterpret,
@@ -46,7 +47,7 @@ greedy_file_path = "assets/greedy.jpg"
 railway_file_path = "assets/railway.jpg"
 happy_file_path = "assets/happy.jpg"
 
-
+@spaces.GPU
 def generate_activations(image):
     prompt = "<image>"
     inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
@@ -116,6 +117,7 @@ def visualize_activations(image, feature_num):
 
     return activation_images
 
+@spaces.GPU
 def clamp_features_max(
     sae: Sae, feature: int, hooked_module: torch.nn.Module, k: float = 10
 ):
@@ -142,6 +144,7 @@ def clamp_features_max(
 
     return handles
 
+@spaces.GPU
 def generate_with_clamp(feature_idx, feature_strength, text, image, chat_history):
     if not isinstance(feature_idx, int):
         feature_idx = int(feature_idx)

requirements.txt

@@ -3,4 +3,4 @@ gradio
 sae_auto_interp @ git+https://github.com/EvolvingLMMs-Lab/multimodal-sae
 fastapi==0.112.2
 gradio==4.44.1
-httpx==0.23.3
+httpx==0.24.1