add spaces

app.py

@@ -11,6 +11,7 @@ import subprocess
 import sys
 import subprocess
 import importlib
+import spaces
 
 try:
     import uniception
@@ -38,10 +39,6 @@ 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"""
     global model, USE_REFINEMENT_MODEL
@@ -59,17 +56,13 @@ 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:
         print(f"Error loading model: {e}")
         return False
 
-
+@spaces.GPU
 def process_images(source_image, target_image, model_type_choice):
     """
     Process two uploaded images and return visualizations
@@ -86,6 +79,9 @@ def process_images(source_image, target_image, model_type_choice):
     if model is None:
         return None, None, None, "Model not loaded. Please restart the application."
 
+    model = model.to("cuda" if torch.cuda.is_available() else "cpu")
+    use_gpu = torch.cuda.is_available()
+
     try:
         # Convert PIL images to numpy arrays
         source_np = np.array(source_image)