Update app.py

app.py

@@ -7,6 +7,9 @@ import torch
 
 # Function to generate and display image
 def generate_and_display_image(prompt_text):
+    # Set device to CPU if GPU is not available
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
     # Create the diffusion pipeline
     pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
     pipe = pipe.to(device)
@@ -37,7 +40,6 @@ def display_image(image_path):
 
 # Model parameters
 model_id = "CompVis/stable-diffusion-v1-4"
-device = "cuda"
 
 # Streamlit text input for prompt
 prompt_text = st.text_input("Enter the prompt:")