Update app.py

app.py

@@ -1,30 +1,32 @@
-from diffusers import StableDiffusionPipeline  # Import the StableDiffusionPipeline class from diffusers
+import streamlit as st
+from diffusers import StableDiffusionPipeline
 import torch
+from PIL import Image
 
-# Define the model ID for Stable Diffusion
+# Load the Stable Diffusion model
 model_id = "CompVis/stable-diffusion-v1-4"
-
-# Load the pre-trained model
 pipe = StableDiffusionPipeline.from_pretrained(model_id)
 
 # Move the model to GPU if available, otherwise use CPU
 pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Function to generate images based on text prompts
-def generate_image(prompt, output_path="generated_image.png"):
-    # Generate the image
-    image = pipe(prompt).images[0]
-    
-    # Save the generated image
-    image.save(output_path)
-    print(f"Image saved at: {output_path}")
+# Streamlit app layout
+st.title("Text-to-Image Generation using Stable Diffusion")
+
+# Text input for prompt
+prompt = st.text_input("Enter a prompt to generate an image:", value="A futuristic city skyline at sunset")
 
-# Example prompts for image generation
-prompt1 = "A futuristic city skyline at sunset"
-generate_image(prompt1, "image1.png")
+# Button to trigger image generation
+if st.button("Generate Image"):
+    with st.spinner("Generating image..."):
+        # Generate image based on prompt
+        image = pipe(prompt).images[0]
+        
+        # Display the image
+        st.image(image, caption="Generated Image", use_column_width=True)
 
-prompt2 = "A beautiful landscape from mountains during sunrise"
-generate_image(prompt2, "image2.png")
+        # Option to save the generated image
+        output_path = "generated_image.png"
+        image.save(output_path)
+        st.success(f"Image saved as {output_path}")
 
-prompt3 = "A beautiful landscape from mountains during sunset"
-generate_image(prompt3, "image3.png")