Update app.py

app.py

@@ -6,10 +6,15 @@ import streamlit as st
 # Load the Stable Diffusion pipeline
 @st.cache(allow_output_mutation=True)
 def load_pipeline():
-    pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
-    pipeline = pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
+    pipeline = StableDiffusionPipeline.from_pretrained(
+        "CompVis/stable-diffusion-v1-4", 
+        torch_dtype=torch.float32  # Change to float32 for CPU
+    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipeline.to(device)
     return pipeline
 
+
 def main():
     st.title("Stable Diffusion Image Generator")
     st.write("Generate images from text prompts using Stable Diffusion 2.1")
@@ -21,17 +26,19 @@ def main():
     prompt = st.text_input("Enter your text prompt", "")
 
     # Generate button
-    if st.button("Generate"):
-        if not prompt:
-            st.warning("Please enter a prompt first.")
-            return
-
-        st.write("Generating your image...")
-        with torch.no_grad():
-            image = pipeline(prompt).images[0]
-
-        st.write("Generated Image:")
-        st.image(image, use_column_width=True)
+if st.button("Generate"):
+    if not prompt:
+        st.warning("Please enter a prompt first.")
+        return
+
+    st.write("Generating your image...")
+    with torch.no_grad():
+        result = pipeline(prompt)  # Returns a dictionary
+        image = result.images[0]  # Extract the first generated image
+
+    st.write("Generated Image:")
+    st.image(image, use_column_width=True)
+   
 
 if __name__ == "__main__":
     main()