Update app.py

app.py

@@ -4,7 +4,8 @@ import numpy as np
 import cv2
 from PIL import Image
 import tempfile # For creating temporary video files
-import os # <--- ADDED: Import the 'os' module
+import os # Import the 'os' module
+import accelerate # Import accelerate for better memory management (recommended)
 
 # Marigold specific imports
 from diffusers import MarigoldDepthPipeline, DDIMScheduler
@@ -27,13 +28,16 @@ try:
     pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
     pipe = pipe.to(device=device, dtype=dtype)
     
-    # Enable xformers for memory-efficient attention if available (optional)
-    try:
-        import xformers
-        pipe.enable_xformers_memory_efficient_attention()
-        print("xFormers enabled for Marigold pipeline.")
-    except ImportError:
-        print("xFormers not found, running without memory-efficient attention.")
+    # Enable xformers for memory-efficient attention ONLY IF CUDA is available
+    if torch.cuda.is_available():
+        try:
+            import xformers
+            pipe.enable_xformers_memory_efficient_attention()
+            print("xFormers enabled for Marigold pipeline.")
+        except ImportError:
+            print("xFormers not found, running without memory-efficient attention (on GPU).")
+    else:
+        print("Running on CPU or MPS. xFormers memory-efficient attention is not applicable.")
 
     print(f"MarigoldDepthPipeline loaded successfully from {CHECKPOINT} on {device}.")
 except Exception as e: