Update app.py

app.py

@@ -19,6 +19,7 @@ import matplotlib.pyplot as plt
 import gc  # Import the garbage collector
 from audio import *
 import os
+multiprocessing.set_start_method("spawn")
 # Define a fallback for environments without GPU
 if os.environ.get("SPACES_ZERO_GPU") is not None:
     import spaces
@@ -29,8 +30,6 @@ else:
             def wrapper(*args, **kwargs):
                 return func(*args, **kwargs)
             return wrapper
-
-
 # Download necessary NLTK data
 try:
     nltk.data.find('tokenizers/punkt')
@@ -57,41 +56,31 @@ def log_gpu_memory():
         print("CUDA is not available. Cannot log GPU memory.")
 
 # --------- MinDalle Image Generation Functions ---------
-# Load MinDalle model once
-# Dynamically determine device and precision
-@spaces.GPU(duration=60 * 3)
-def load_min_dalle_model(models_root: str = 'pretrained'):
-    """
-    Load the MinDalle model, automatically selecting device and precision.
-
-    Args:
-        models_root: Path to the directory containing MinDalle models.
-
-    Returns:
-        An instance of the MinDalle model.
-    """
-    print("DEBUG: Loading MinDalle model...")
-    
+# Check for GPU availability
+def check_gpu_availability():
     if torch.cuda.is_available():
-        device = 'cuda'
-        dtype = torch.float16
-        print("DEBUG: Using GPU with float16 precision.")
+        print(f"CUDA devices: {torch.cuda.device_count()}")
+        print(f"Current device: {torch.cuda.current_device()}")
+        print(torch.cuda.get_device_properties(torch.cuda.current_device()))
     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,
-        is_verbose=True,
-        dtype=dtype,
-        device=device
-    )
+        print("CUDA is not available. Running on CPU.")
+check_gpu_availability()
+# GPU-safe model loading
+def initialize_min_dalle_with_gpu():
+    @spaces.GPU(duration=60 * 3)
+    def load_model():
+        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'
+        )
+    return load_model()
 
-# Initialize the MinDalle model (will now automatically use GPU if available)
-min_dalle_model = load_min_dalle_model()
+# Initialize MinDalle model
+min_dalle_model = initialize_min_dalle_with_gpu()
 
 def generate_image_with_min_dalle(
     model: MinDalle,