Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import streamlit as st
 import matplotlib.pyplot as plt
 from PIL import Image
 from diffusers import StableDiffusionPipeline
@@ -19,29 +20,28 @@ def generate_and_display_image(prompt_text):
         image.save(image_path)
 
         # Show success message
-        print("Image generated successfully!")
+        st.write("Image generated successfully!")
 
         # Display the generated image
         display_image(image_path)
 
     except Exception as e:
         # Show error message if an error occurs
-        print(f"An error occurred: {str(e)}")
+        st.error(f"An error occurred: {str(e)}")
 
 # Function to display image
 def display_image(image_path):
     # Open the generated image
     generated_image = Image.open(image_path)
-    plt.imshow(generated_image)
-    plt.axis('off')
-    plt.show()
+    st.image(generated_image, use_column_width=True)
 
 # Model parameters
 model_id = "CompVis/stable-diffusion-v1-4"
 device = "cuda"
 
-# Text input for prompt
-prompt_text = input("Enter the prompt: ")
+# Streamlit text input for prompt
+prompt_text = st.text_input("Enter the prompt:")
 
-# Generate and display image
-generate_and_display_image(prompt_text)
+# Generate and display image if prompt is provided
+if prompt_text:
+    generate_and_display_image(prompt_text)