Update app.py

app.py

@@ -3,53 +3,48 @@ from diffusers import StableDiffusionPipeline
 import torch
 from PIL import Image
 
-# Title and description of the app
+# Streamlit App Title
 st.title("🖼️ Stable Diffusion Image Generator")
-st.write("Generate images from text using the Stable Diffusion v1.5 model!")
+st.write("Generate images from text prompts using Stable Diffusion!")
 
-# Sidebar for user inputs
-st.sidebar.title("Input Options")
+# Sidebar for prompt input and button
+st.sidebar.title("Image Generation Settings")
 prompt = st.sidebar.text_input("Enter your prompt", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k")
 generate_button = st.sidebar.button("Generate Image")
 
-# Load the pipeline when the app starts
+# Cache the Stable Diffusion model to load it only once
 @st.cache_resource
-def load_pipeline():
-    # Use a smaller model or a more efficient pipeline
-    pipe = StableDiffusionPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5", 
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        revision="fp16" if torch.cuda.is_available() else None
-    )
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipe = pipe.to(device)
-    return pipe
+def load_model():
+    model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+    model = model.to("cuda" if torch.cuda.is_available() else "cpu")
+    return model
 
-pipe = load_pipeline()
+# Load the model
+pipe = load_model()
 
-# Generate image when button is clicked
+# Generate the image when the button is clicked
 if generate_button:
-    st.write(f"### Prompt: {prompt}")
-    with st.spinner("Generating image... Please wait."):
-        try:
-            # Generate the image
-            with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+    if prompt.strip() == "":
+        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."):
+            try:
+                # Generate the image
                 image = pipe(prompt).images[0]
-            
-            # Display the generated image
-            st.image(image, caption="Generated Image", use_column_width=True)
-            
-            # Option to download the image
-            img_path = "generated_image.png"
-            image.save(img_path)
-            with open(img_path, "rb") as img_file:
-                st.download_button(
-                    label="Download Image",
-                    data=img_file,
-                    file_name="generated_image.png",
-                    mime="image/png"
-                )
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
 
-            ## print
+                # Display the image
+                st.image(image, caption="Generated Image", use_column_width=True)
+
+                # Provide download option
+                img_path = "generated_image.png"
+                image.save(img_path)
+                with open(img_path, "rb") as img_file:
+                    st.download_button(
+                        label="Download Image",
+                        data=img_file,
+                        file_name="generated_image.png",
+                        mime="image/png"
+                    )
+            except Exception as e:
+                st.error(f"An error occurred: {e}")