Update app.py

app.py

@@ -54,7 +54,7 @@ def resize(img, target_res=224, resize=True, to_pil=True, edge=False, sampling_f
     return canvas
 
 # ─── Feature extraction ──────────────────────────────────────────
-@spaces.GPU
+@spaces.GPU(duration=20)
 def get_processed_features_dino(num_patches, img,use_dummy):
     batch = extractor_vit.preprocess_pil(img)
     features_dino = extractor_vit.extract_descriptors(batch.to(extractor_vit.device), layer=11, facet='token') \
@@ -68,11 +68,10 @@ 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)
-    desc = desc.cpu()
-    torch.cuda.empty_cache()
     return desc  # shape [1, C, num_patches, num_patches]
 
 # ─── Similarity computation ───────────────────────────────────────
+spaces.GPU(duration=20)
 def get_sim(
     coord: tuple[int,int],
     feat1: torch.Tensor,
@@ -97,7 +96,9 @@ def get_sim(
     # Cosine similarity along channel‐dim
     cos = nn.CosineSimilarity(dim=1)
     cos_map = cos(src_vec, trg_ft)[0]           # [img_size, img_size]
-    return cos_map.cpu().numpy()         
+    cos_map = cos_map.cpu()
+    torch.cuda.empty_cache()
+    return cos_map.numpy()         
 
 # ─── Drawing helper ───────────────────────────────────────────────
 def draw_point(img_arr: np.ndarray, x: int, y: int, size: int, color=(255,0,0)) -> np.ndarray:
@@ -108,6 +109,7 @@ 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(duration=20)
 def update_features(
     img: Image,
     num_patches,
@@ -122,7 +124,7 @@ def update_features(
         return None, None, None
     img = resize(img, target_res=target_res, resize=True, to_pil=True)
     feat = get_processed_features_dino(num_patches, img=img,use_dummy=use_dummy)
-    return img, feat.cpu(), Image.fromarray(np.array(img))
+    return img, feat, Image.fromarray(np.array(img))
 
 # ─── Click handler ───────────────────────────────────────────────
 def on_select(