Update app.py

app.py

@@ -1,5 +1,17 @@
-import gradio as gr
+import os
+
+# If running on Hugging Face Spaces, import the GPU decorator
+if "SPACE_ID" in os.environ:
+    from spaces.zero.decorator import GPU
+else:
+    # Define a dummy GPU decorator when running locally
+    def GPU(func):
+        return func
+
+# Now you can safely import torch and other libraries
 import torch
+import gradio as gr
+
 import numpy as np
 from PIL import Image
 import os
@@ -20,15 +32,6 @@ args = parser.parse_args()
 
 
 
-# try:
-#     import spaces
-#     GPU = spaces.GPU
-# except ImportError:
-#     # Define a no-op decorator if running locally
-#     def GPU(func):
-#         return func
-import spaces
-
 # Create model directories if they don't exist
 os.makedirs("models", exist_ok=True)
 os.makedirs("stimuli", exist_ok=True)
@@ -36,7 +39,7 @@ os.makedirs("stimuli", exist_ok=True)
 # Initialize model
 model = GenerativeInferenceModel()
 
-@spaces.GPU 
+@GPU 
 def run_inference(image, model_type, illusion_type, eps_value, num_iterations):
     # Convert eps to float
     eps = float(eps_value)