Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionPipeline
 import torch
 from PIL import Image
 
@@ -15,7 +15,7 @@ generate_button = st.sidebar.button("Generate Image")
 # Cache the Stable Diffusion model to load it only once
 @st.cache_resource
 def load_model():
-    model = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+    model = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float16)
     model = model.to("cuda" if torch.cuda.is_available() else "cpu")
     return model
 
@@ -28,10 +28,10 @@ if generate_button:
         st.error("Prompt cannot be empty. Please enter a valid text prompt.")
     else:
         st.write(f"### Prompt: {prompt}")
-        with st.spinner("Generating image... Please wait."):
+        with st.spinner("Generating image... Please wait. This may take a few moments depending on system performance."):
             try:
                 # Generate the image
-                image = pipe(prompt).images[0]
+                image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
 
                 # Display the image
                 st.image(image, caption="Generated Image", use_column_width=True)