Update app.py

app.py

@@ -1,5 +1,6 @@
 
 import os
+import spaces
 import gradio as gr
 import numpy as np
 import torch
@@ -53,6 +54,7 @@ def resize(img, target_res=224, resize=True, to_pil=True, edge=False, sampling_f
     return canvas
 
 # ─── Feature extraction ──────────────────────────────────────────
+@spaces.GPU
 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') \
@@ -104,6 +106,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
 def update_features(
     img: Image,
     num_patches,
@@ -171,12 +174,13 @@ def reload_img(
 
 
 # ─── Configuration ───────────────────────────────────────────────
-num_patches = 30
+num_patches = 60
 target_res = num_patches * 14
-ckpt_file = "ckpts/dino_spair_0300.pth"
+ckpt_file = "./ckpts/dino_spair_0300.pth"
 
 # ─── Model setup ─────────────────────────────────────────────────
-device = 'cpu' #'cuda' if torch.cuda.is_available() else 'cpu'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"Using device: {device}")
 aggre_net = AggregationNetwork(feature_dims=[768], projection_dim=768, device=device)
 aggre_net.load_pretrained_weights(torch.load(ckpt_file, map_location=device))
 aggre_net_dummy  = DummyAggregationNetwork()
@@ -198,7 +202,6 @@ with gr.Blocks() as demo:
     Welcome to the DIY-SC demo!
     Upload two images and select a keypoint in the source image. This demo will compute and visualize the feature similarity map and a corresponding point in the target image.
     You can choose between the DIY-SC (DINOv2) or the DINOv2 feature extractor.
-    The demo is currently only running on CPUs. For this reason the run time is longer and a lower resolution is used, which results in a performance drop. We will update this soon.
     """)
 
     # Image upload / display components
@@ -241,5 +244,6 @@ with gr.Blocks() as demo:
         outputs=[src,tgt]
     )
 
+
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()