Update app.py

app.py

@@ -10,6 +10,7 @@ import matplotlib.pyplot as plt
 from matplotlib.colors import ListedColormap
 from utils import add_feature_on_area, replace_with_feature
 import threading
+import spaces
 
 code_to_block = {
     "down.2.1": "unet.down_blocks.2.attentions.1",
@@ -19,6 +20,7 @@ code_to_block = {
 }
 lock = threading.Lock()
 
+@spaces.GPU
 def process_cache(cache, saes_dict):
 
     top_features_dict = {}
@@ -65,6 +67,7 @@ def plot_image_heatmap(cache, block_select, radio):
 
 
 def create_prompt_part(pipe, saes_dict, demo):
+    @spaces.GPU
     def image_gen(prompt):
         lock.acquire()
         try:
@@ -137,6 +140,7 @@ def downsample_mask(image, factor):
     return downsampled
 
 def create_intervene_part(pipe: HookedStableDiffusionXLPipeline, saes_dict, means_dict, demo):
+    @spaces.GPU
     def image_gen(prompt, num_steps):
         lock.acquire()
         try:
@@ -151,6 +155,7 @@ def create_intervene_part(pipe: HookedStableDiffusionXLPipeline, saes_dict, mean
             lock.release()
         return images.images[0]
 
+    @spaces.GPU
     def image_mod(prompt, block_str, brush_index, strength, num_steps, input_image):
         block = block_str.split(" ")[0]
 
@@ -184,6 +189,7 @@ def create_intervene_part(pipe: HookedStableDiffusionXLPipeline, saes_dict, mean
             lock.release()
         return image
 
+    @spaces.GPU
     def feature_icon(block_str, brush_index):
         block = block_str.split(" ")[0]
         if block in ["mid.0", "up.0.0"]: