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,12 +15,15 @@ generate_button = st.sidebar.button("Generate Image")
 # Load the pipeline when the app starts
 @st.cache_resource
 def load_pipeline():
-    pipe = DiffusionPipeline.from_pretrained(
+    # Use a smaller model or a more efficient pipeline
+    pipe = StableDiffusionPipeline.from_pretrained(
         "runwayml/stable-diffusion-v1-5", 
-        torch_dtype=torch.float16
+        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"
-    return pipe.to(device)
+    pipe = pipe.to(device)
+    return pipe
 
 pipe = load_pipeline()
 
@@ -28,19 +31,23 @@ pipe = load_pipeline()
 if generate_button:
     st.write(f"### Prompt: {prompt}")
     with st.spinner("Generating image... Please wait."):
-        # 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"
-            )
+        try:
+            # Generate the image
+            with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
+                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}")