Update app.py

app.py

@@ -41,9 +41,9 @@ def generate_image(api_key, prompt, aspect_ratio, mode, model, seed, output_form
         files=files,
         data=data
     )
+    
     return response
 
-# Streamlit Ui
 def main():
     st.image("static/SD3_webui_logo_image.png", use_column_width=True)
     
@@ -53,6 +53,10 @@ def main():
         st.session_state.prompt = ""
     if 'negative_prompt' not in st.session_state:
         st.session_state.negative_prompt = ""
+    if 'image_data' not in st.session_state:
+        st.session_state.image_data = None
+    if 'image_path' not in st.session_state:
+        st.session_state.image_path = None
     
     api_key = st.text_input("Enter your API key:", type="password", value=st.session_state.api_key)
     models = st.selectbox("Select model:", ["sd3", "sd3-turbo"], index=0)
@@ -74,18 +78,18 @@ def main():
         with st.spinner('Generating Image...'):
             result = generate_image(api_key, prompt, aspect_ratios, mode, models, seed, output_formats, strength, negative_prompt, uploaded_file)
             if result.status_code == 200:
-                image = Image.open(io.BytesIO(result.content))
-                st.image(image, use_column_width=True)
-                current_time = time.strftime("%H.%M.%S").lower()
-                file_path = os.path.join(f"gen_{models}_{seed}_{current_time}.{output_formats}")
-                image.save(file_path)
-                st.success(f'Image saved to {file_path}')
-                # Download generated image for Huggingface Spaces add-on
-                with open(file_path, "rb") as file:
-                    st.download_button(label=":green[Download Image]", data=file, file_name=file_path, mime=f"image/{output_formats}", use_container_width=True)
+                st.session_state.image_data = Image.open(io.BytesIO(result.content))
+                st.session_state.image_path = os.path.join(f"gen_{models}_{seed}_{time.strftime('%H.%M.%S').lower()}.{output_formats}")
+                st.session_state.image_data.save(st.session_state.image_path)
+                st.success(f'Image saved to {st.session_state.image_path}')
             else:
                 st.error('Failed to generate image: ' + str(result.json()))
 
+    if st.session_state.image_data is not None:
+        st.image(st.session_state.image_data, use_column_width=True)
+        with open(st.session_state.image_path, "rb") as file:
+            st.download_button(label=":green[Download Image]", data=file, file_name=st.session_state.image_path, mime=f"image/{output_formats}", use_container_width=True)
+
     with st.sidebar:
         st.link_button("Stability.Ai _ API Documentation", "https://platform.stability.ai/docs/api-reference", use_container_width=True)
         with st.expander("Image Generation Costs", expanded=True):
@@ -98,6 +102,6 @@ def main():
             Credits cost $10 per 1,000 credits, which is enough credits for roughly 154 SD3 images or 250 SD3 Turbo images.
             """)
         st.link_button("Account_Page", "https://platform.stability.ai/account/credits", use_container_width=True)
-                
+
 if __name__ == "__main__":
     main()