Update app.py

app.py

@@ -68,10 +68,11 @@ 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,
@@ -96,9 +97,7 @@ 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]
-    cos_map = cos_map.cpu()
-    torch.cuda.empty_cache()
-    return cos_map.numpy()         
+    return cos_map.cpu().numpy()         
 
 # ─── Drawing helper ───────────────────────────────────────────────
 def draw_point(img_arr: np.ndarray, x: int, y: int, size: int, color=(255,0,0)) -> np.ndarray:
@@ -109,7 +108,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(duration=20)
 def update_features(
     img: Image,
     num_patches,
@@ -124,7 +122,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, Image.fromarray(np.array(img))
+    return img, feat.cpu(), Image.fromarray(np.array(img))
 
 # ─── Click handler ───────────────────────────────────────────────
 def on_select(
@@ -177,7 +175,7 @@ def reload_img(
 
 
 # ─── Configuration ───────────────────────────────────────────────
-num_patches = 45
+num_patches = 60
 target_res = num_patches * 14
 ckpt_file = "./ckpts/dino_spair_0300.pth"