Update app.py

app.py

@@ -68,8 +68,8 @@ 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().detach()
-    del batch, features_dino
+    desc = desc.cpu()
+    # del batch, features_dino
     gc.collect()
     torch.cuda.empty_cache()
     return desc  # shape [1, C, num_patches, num_patches]
@@ -114,6 +114,7 @@ def update_features(
     num_patches,
     use_dummy
 ):
+    gc.collect()
     torch.cuda.empty_cache()
     """
     Given a PIL image, returns:
@@ -138,6 +139,8 @@ def on_select(
     or_src_img: Image,
     sel: gr.SelectData
 ):
+    gc.collect()
+    torch.cuda.empty_cache()
     # Convert to numpy arrays
     src_arr = np.array(or_src_img)
     tgt_arr = np.array(or_tgt_img)
@@ -194,6 +197,9 @@ extractor_vit.model.eval()
 
 upsampler = nn.Upsample(size=(target_res, target_res), mode='bilinear', align_corners=False)
 
+gc.collect()
+torch.cuda.empty_cache()
+
 # ─── Build Gradio UI ──────────────────────────────────────────────
 with gr.Blocks() as demo:
     # Hidden states to hold features