Update app.py

app.py

@@ -18,24 +18,23 @@ prompt = st.text_input("Enter a prompt for image generation:", "a robot on a sta
 if st.button("Generate Image"):
     if api_key:
         with st.spinner("Generating image... Please wait!"):
-            # Call the HF API
-            image_bytes = client.text_to_image(prompt, model="black-forest-labs/FLUX.1-dev")
-            image = Image.open(BytesIO(image_bytes))
-            
-            # Show the generated image
+            # Generate image from Hugging Face model
+            image = client.text_to_image(prompt, model="black-forest-labs/FLUX.1-dev")
+
+            # Convert PIL Image to BytesIO
+            img_buffer = BytesIO()
+            image.save(img_buffer, format="PNG")
+            img_bytes = img_buffer.getvalue()  # Convert to bytes
+
+            # Display image
             st.image(image, caption="Generated Image", use_column_width=True)
-            
-            # Save the image
-            image.save("generated_image.png")
-            
-            # Provide a download button
-            with open("generated_image.png", "rb") as file:
-                st.download_button(
-                    label="Download Image",
-                    data=file,
-                    file_name="generated_image.png",
-                    mime="image/png"
-                )
+
+            # Provide download button
+            st.download_button(
+                label="Download Image",
+                data=img_bytes,
+                file_name="generated_image.png",
+                mime="image/png"
+            )
     else:
         st.error("API key not found. Please set the `HF_API_KEY` in secrets.")
-