Added CPU support

app.py

@@ -70,9 +70,15 @@ def set_example(selected_prompt):
             return example["prompt"], example["ckpt"], example["aspect"], example["mode"]
     return None, None, None, None  # Default values if not found
 
-# Ensure model and scheduler are initialized in GPU-enabled function
-if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
+# Check if CUDA is available (GPU support), and set the appropriate device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load the pipeline for the specified device
+# For GPU, use torch_dtype=torch.float16 for better performance
+if device == "cuda":
+    pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to(device)
+else:
+    pipe = StableDiffusionXLPipeline.from_pretrained(base).to(device)
 
 if SAFETY_CHECKER:
     from safety_checker import StableDiffusionSafetyChecker