fix spaces

app.py

@@ -11,6 +11,7 @@ import subprocess
 import sys
 import subprocess
 import importlib
+from functools import lru_cache
 import spaces
 
 try:
@@ -36,31 +37,28 @@ from uniflowmatch.models.ufm import (
 from uniflowmatch.utils.viz import warp_image_with_flow
 
 # Global model variable
-model = None
 USE_REFINEMENT_MODEL = False
 
+@lru_cache(maxsize=2)
 def initialize_model(use_refinement: bool = False):
     """Initialize the model - call this once at startup"""
-    global model, USE_REFINEMENT_MODEL
-    USE_REFINEMENT_MODEL = use_refinement
-
     try:
-        if USE_REFINEMENT_MODEL:
+        if use_refinement:
             print("Loading UFM Refinement model from infinity1096/UFM-Refine...")
-            model = UniFlowMatchClassificationRefinement.from_pretrained("infinity1096/UFM-Refine")
+            model_obj = UniFlowMatchClassificationRefinement.from_pretrained("infinity1096/UFM-Refine")
         else:
             print("Loading UFM Base model from infinity1096/UFM-Base...")
-            model = UniFlowMatchConfidence.from_pretrained("infinity1096/UFM-Base")
+            model_obj = UniFlowMatchConfidence.from_pretrained("infinity1096/UFM-Base")
 
         # Set model to evaluation mode
-        if hasattr(model, "eval"):
-            model.eval()
+        if hasattr(model_obj, "eval"):
+            model_obj.eval()
 
         print("Model loaded successfully!")
-        return True
+        return True, model_obj
     except Exception as e:
         print(f"Error loading model: {e}")
-        return False
+        return False, None
 
 @spaces.GPU
 def process_images(source_image, target_image, model_type_choice):
@@ -72,9 +70,8 @@ def process_images(source_image, target_image, model_type_choice):
 
     # Reinitialize model if type has changed
     current_refinement = model_type_choice == "Refinement Model"
-    if current_refinement != USE_REFINEMENT_MODEL:
-        print(f"Switching to {model_type_choice}...")
-        initialize_model(current_refinement)
+    print(f"Switching to {model_type_choice}...")
+    ret, model = initialize_model(current_refinement)
 
     if model is None:
         return None, None, None, "Model not loaded. Please restart the application."