Update app.py

app.py

@@ -1,20 +1,25 @@
+import streamlit as st
 from diffusers import DiffusionPipeline
 import torch
+from PIL import Image
 
-# Load the pipeline
-pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float16)
+# Streamlit app title
+st.title("Stable Diffusion Image Generator")
 
+# Load the diffusion pipeline
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float32)
+pipe.to("cpu")  # Use CPU for inference
 
-pipe.to("cpu")
+# Get user input (prompt)
+prompt = st.text_input("What do you want to see?", "A beautiful landscape")
 
-# Get the prompt from the user
-prompt = input("What do you want to see? ")
+# Button to generate image
+if st.button('Generate Image'):
+    with st.spinner('Generating...'):
+        # Generate image
+        image = pipe(prompt).images[0]
+        
+        # Display image in Streamlit
+        st.image(image, caption="Generated Image", use_column_width=True)
 
-# Generate the image
-image = pipe(prompt).images[0]
-
-# Display the image
-image.show()
-
-# Print a message
-print("How cool is this!")
+    st.success("Image generated successfully!")