Update app.py

app.py

@@ -19,6 +19,17 @@ import matplotlib.pyplot as plt
 import gc  # Import the garbage collector
 from audio import *
 import os
+# Define a fallback for environments without GPU
+if os.environ.get("SPACES_ZERO_GPU") is not None:
+    import spaces
+else:
+    class spaces:
+        @staticmethod
+        def GPU(func):
+            def wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+            return wrapper
+
 
 # Download necessary NLTK data
 try:
@@ -46,30 +57,39 @@ def log_gpu_memory():
         print("CUDA is not available. Cannot log GPU memory.")
 
 # --------- MinDalle Image Generation Functions ---------
-
 # Load MinDalle model once
-def load_min_dalle_model(models_root: str = 'pretrained', fp16: bool = True):
+# Dynamically determine device and precision
+def load_min_dalle_model(models_root: str = 'pretrained'):
     """
-    Load the MinDalle model.
+    Load the MinDalle model, automatically selecting device and precision.
 
     Args:
         models_root: Path to the directory containing MinDalle models.
-        fp16: Whether to use float16 for faster generation (requires CUDA).
 
     Returns:
         An instance of the MinDalle model.
     """
     print("DEBUG: Loading MinDalle model...")
+    
+    if torch.cuda.is_available():
+        device = 'cuda'
+        dtype = torch.float16
+        print("DEBUG: Using GPU with float16 precision.")
+    else:
+        device = 'cpu'
+        dtype = torch.float32
+        print("DEBUG: Using CPU with float32 precision.")
+
     return MinDalle(
         is_mega=True,
         models_root=models_root,
-        is_reusable=False,  # Set is_reusable to False
+        is_reusable=False,
         is_verbose=True,
-        dtype=torch.float16 if fp16 else torch.float32,
+        dtype=dtype,
         device=device
     )
 
-# Initialize the MinDalle model
+# Initialize the MinDalle model (will now automatically use GPU if available)
 min_dalle_model = load_min_dalle_model()
 
 def generate_image_with_min_dalle(