turn off xformers for cpu

app.py

@@ -10,7 +10,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Load the model in FP16
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float16)
 
-# Move the pipeline to GPU
+# Move the pipeline to GPU if available
 pipe = pipe.to(device)
 
 # Convert text encoders to full precision
@@ -18,12 +18,21 @@ pipe.text_encoder = pipe.text_encoder.to(torch.float32)
 if hasattr(pipe, 'text_encoder_2'):
     pipe.text_encoder_2 = pipe.text_encoder_2.to(torch.float32)
 
-# Enable memory efficient attention if available
-if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
-    pipe.enable_xformers_memory_efficient_attention()
+# Enable memory efficient attention if available and on CUDA
+if device == "cuda" and hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
+    try:
+        pipe.enable_xformers_memory_efficient_attention()
+        print("xformers memory efficient attention enabled")
+    except Exception as e:
+        print(f"Could not enable memory efficient attention: {e}")
 
-# Compile the UNet for potential speedups
-pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# Compile the UNet for potential speedups if on CUDA
+if device == "cuda":
+    try:
+        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+        print("UNet compiled for potential speedups")
+    except Exception as e:
+        print(f"Could not compile UNet: {e}")
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048