add CUDA inference

app.py

@@ -37,6 +37,9 @@ from uniflowmatch.utils.viz import warp_image_with_flow
 model = None
 USE_REFINEMENT_MODEL = False
 
+use_gpu = torch.cuda.is_available()
+if use_gpu:
+    print("Using GPU for processing.")
 
 def initialize_model(use_refinement: bool = False):
     """Initialize the model - call this once at startup"""
@@ -55,6 +58,10 @@ def initialize_model(use_refinement: bool = False):
         if hasattr(model, "eval"):
             model.eval()
 
+        if use_gpu:
+            print("Moving model to GPU...")
+            model = model.to("cuda")
+
         print("Model loaded successfully!")
         return True
     except Exception as e:
@@ -99,8 +106,8 @@ def process_images(source_image, target_image, model_type_choice):
         # === Predict Correspondences ===
         with torch.no_grad():
             result = model.predict_correspondences_batched(
-                source_image=torch.from_numpy(source_rgb),
-                target_image=torch.from_numpy(target_rgb),
+                source_image=torch.from_numpy(source_rgb).to("cuda" if use_gpu else "cpu"),
+                target_image=torch.from_numpy(target_rgb).to("cuda" if use_gpu else "cpu"),
             )
 
             # Extract results based on your model's output structure