Update app.py

app.py

@@ -27,7 +27,20 @@ if st.button("Generate Image"):
 
         # Display Image
         image = Image.open(io.BytesIO(image_bytes))
-        st.image(image, caption="Generated Image", use_column_width=True)
+        st.image(image, caption="Generated Image", use_container_width=True)
+
+        # Convert Image to Bytes for Download
+        img_buffer = io.BytesIO()
+        image.save(img_buffer, format="PNG")
+        img_buffer.seek(0)
+
+        # Download Button
+        st.download_button(
+            label="Download Image 📥",
+            data=img_buffer,
+            file_name="generated_image.png",
+            mime="image/png"
+        )
     else:
         st.warning("Please enter a prompt before generating an image.")