Update app.py

app.py

@@ -26,10 +26,10 @@ if os.environ.get("SPACES_ZERO_GPU") is not None:
 else:
     class spaces:
         @staticmethod
-        def GPU(func):
-            def wrapper(*args, **kwargs):
-                return func(*args, **kwargs)
-            return wrapper
+        def GPU(func=None, duration=None):
+            def wrapper(fn):
+                return fn
+            return wrapper if func is None else wrapper(func)
 
 # Download necessary NLTK data
 def setup_nltk():
@@ -82,23 +82,35 @@ def check_gpu_availability():
 check_gpu_availability()
 
 # GPU-Safe MinDalle Model Loading
-def initialize_min_dalle_with_gpu():
+def initialize_min_dalle():
     """Load the MinDalle model with GPU support."""
-    @spaces.GPU()
-    def load_model():
-        print("Loading MinDalle model...")
+    if torch.cuda.is_available():
+        @spaces.GPU(duration=60 * 3)
+        def load_model():
+            print("Loading MinDalle model on GPU...")
+            return MinDalle(
+                is_mega=True,
+                models_root='pretrained',
+                is_reusable=False,
+                is_verbose=True,
+                dtype=torch.float16,
+                device='cuda'
+            )
+        return load_model()
+    else:
+        print("Loading MinDalle model on CPU...")
         return MinDalle(
             is_mega=True,
             models_root='pretrained',
             is_reusable=False,
             is_verbose=True,
-            dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device='cuda' if torch.cuda.is_available() else 'cpu'
+            dtype=torch.float32,
+            device='cpu'
         )
-    return load_model()
 
 # Initialize MinDalle Model
-min_dalle_model = initialize_min_dalle_with_gpu()
+min_dalle_model = initialize_min_dalle()
+
 
 
 def generate_image_with_min_dalle(