Update app.py

app.py

@@ -7,8 +7,16 @@ from PIL import Image
 @st.cache_resource
 def load_model():
     model_id = "runwayml/stable-diffusion-v1-5"
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-    pipe = pipe.to("cuda")
+    
+    # Load the model
+    pipe = StableDiffusionPipeline.from_pretrained(model_id)
+    
+    # Check if CUDA is available and use it, otherwise fall back to CPU
+    if torch.cuda.is_available():
+        pipe = pipe.to("cuda")
+    else:
+        pipe = pipe.to("cpu")
+        
     return pipe
 
 pipe = load_model()