Update app.py

app.py

@@ -3,14 +3,16 @@ from diffusers import AutoPipelineForText2Image
 import torch
 
 def load_model():
-    # Load the pipeline with your specific configuration
-    # Ensure the model and any custom weights are accessible
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+    
     pipeline = AutoPipelineForText2Image.from_pretrained(
         "stabilityai/stable-diffusion-xl-base-1.0",
-        torch_dtype=torch.float16,
-        variant="fp16"
-    ).to("cuda")
-    pipeline.load_lora_weights(weight_name="EnvyFloorplansXL01.safetensors")
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,  # Use float32 if on CPU
+        variant="fp16" if device == "cuda" else "no"
+    ).to(device)
+    
+    pipeline.load_lora_weights("/content/", weight_name="EnvyFloorplansXL01.safetensors")
     return pipeline
 
 pipeline = load_model()