update cache

app.py

@@ -464,20 +464,75 @@ def image_clip_feature(
     return outputs
 
 
+
+import hashlib
+import pickle
+import sys
+from collections import OrderedDict
+
+# Cache dictionary with limited size
+class LimitedSizeCache(OrderedDict):
+    def __init__(self, max_size_bytes):
+        self.max_size_bytes = max_size_bytes
+        self.current_size_bytes = 0
+        super().__init__()
+
+    def __setitem__(self, key, value):
+        item_size = self.get_item_size(value)
+        # Evict items until there is enough space
+        while self.current_size_bytes + item_size > self.max_size_bytes:
+            self.popitem(last=False)
+        super().__setitem__(key, value)
+        self.current_size_bytes += item_size
+
+    def __delitem__(self, key):
+        value = self[key]
+        super().__delitem__(key)
+        self.current_size_bytes -= self.get_item_size(value)
+
+    def get_item_size(self, value):
+        """Estimate the size of the value in bytes."""
+        return sys.getsizeof(value)
+
+# Initialize the cache with a 4GB limit
+cache = LimitedSizeCache(max_size_bytes=4 * 1024 * 1024 * 1024)  # 4GB
+
+def compute_hash(*args, **kwargs):
+    """Compute a unique hash based on the function arguments."""
+    hasher = hashlib.sha256()
+    pickled_args = pickle.dumps((args, kwargs))
+    hasher.update(pickled_args)
+    return hasher.hexdigest()
+
+
 def extract_features(images, model_name="sam", node_type="block", layer=-1):
+    # Compute the cache key
+    cache_key = compute_hash(images, model_name, node_type, layer)
+
+    # Check if the result is already in the cache
+    if cache_key in cache:
+        print("Cache hit!")
+        return cache[cache_key]
+    
+    
+    # Compute the result if not in cache
     if model_name == "SAM(sam_vit_b)":
         if not use_cuda:
             gr.warning("GPU not detected. Running SAM on CPU, ~30s/image.")
-        return image_sam_feature(images, node_type=node_type, layer=layer)
+        result = image_sam_feature(images, node_type=node_type, layer=layer)
     elif model_name == 'MobileSAM':
-        return image_mobilesam_feature(images, node_type=node_type, layer=layer)
+        result = image_mobilesam_feature(images, node_type=node_type, layer=layer)
     elif model_name == "DiNO(dinov2_vitb14_reg)":
-        return image_dino_feature(images, node_type=node_type, layer=layer)
+        result = image_dino_feature(images, node_type=node_type, layer=layer)
     elif model_name == "CLIP(openai/clip-vit-base-patch16)":
-        return image_clip_feature(images, node_type=node_type, layer=layer)
+        result = image_clip_feature(images, node_type=node_type, layer=layer)
     else:
         raise ValueError(f"Model {model_name} not supported.")
 
+    # Store the result in the cache
+    cache[cache_key] = result
+    
+    return result
 
 def compute_ncut(
     features,
@@ -584,7 +639,7 @@ demo = gr.Interface(
         gr.Dropdown(["attn", "mlp", "block"], label="Node type", value="block", elem_id="node_type", info="attn: attention output, mlp: mlp output, block: sum of residual stream"),
         gr.Slider(0, 11, step=1, label="Layer", value=11, elem_id="layer", info="which layer of the image backbone features"),
         gr.Slider(1, 1000, step=1, label="Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more object parts, decrease for whole object'),
-        gr.Slider(0.01, 1, step=0.01, label="Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for more aggressive cleaning on the affinity matrix"),
+        gr.Slider(0.01, 1, step=0.01, label="Affinity focal gamma", value=0.3, elem_id="affinity_focal_gamma", info="decrease for more aggressive cleaning on the affinity matrix"),
     ],
     gr.Gallery(value=default_outputs, label="NCUT Embedding", show_label=False, elem_id="ncut", columns=[3], rows=[1], object_fit="contain", height="auto"),
     additional_inputs=[