Update app.py

app.py

@@ -56,6 +56,7 @@ def resize(img, target_res=224, resize=True, to_pil=True, edge=False, sampling_f
 # ─── Feature extraction ──────────────────────────────────────────
 @spaces.GPU
 def get_processed_features_dino(num_patches, img,use_dummy):
+    torch.cuda.empty_cache()
     batch = extractor_vit.preprocess_pil(img)
     features_dino = extractor_vit.extract_descriptors(batch.to(extractor_vit.device), layer=11, facet='token') \
                                     .permute(0,1,3,2) \
@@ -68,7 +69,7 @@ def get_processed_features_dino(num_patches, img,use_dummy):
             desc = aggre_net(features_dino)
         norms = torch.linalg.norm(desc, dim=1, keepdim=True)
         desc = desc / (norms + 1e-8)
-    return desc  # shape [1, C, num_patches, num_patches]
+    return desc.cpu()  # shape [1, C, num_patches, num_patches]
 
 # ─── Similarity computation ───────────────────────────────────────
 def get_sim(
@@ -106,7 +107,6 @@ def draw_point(img_arr: np.ndarray, x: int, y: int, size: int, color=(255,0,0))
     return np.array(pil)
 
 # ─── Feature‐updating callback ───────────────────────────────────
-@spaces.GPU
 def update_features(
     img: Image,
     num_patches,
@@ -117,7 +117,6 @@ def update_features(
       1) the same PIL image (so it can be displayed)
       2) its DINO descriptor tensor, stored in a gr.State
     """
-    torch.cuda.empty_cache()
     if img is None:
         return None, None, None
     img = resize(img, target_res=target_res, resize=True, to_pil=True)
@@ -245,7 +244,5 @@ with gr.Blocks() as demo:
         outputs=[src,tgt]
     )
 
-    # Add examples=[["What are you doing?"], ["Where should we time travel to?"]], (https://huggingface.co/learn/llm-course/chapter9/4?fw=pt)
-
 if __name__ == "__main__":
     demo.launch()