Update app.py

app.py

@@ -1,16 +1,13 @@
 import streamlit as st
 import requests
 from PIL import Image
-import datetime
 import io
 import os
+import time
 
-output_dir = "outputs"
-os.makedirs(output_dir, exist_ok=True)
+st.set_page_config(initial_sidebar_state="collapsed", layout="centered")
 
-st.set_page_config(initial_sidebar_state="collapsed", page_icon="static/sai_logo.ico", layout="centered")
-
-def generate_image(api_key, prompt, aspect_ratio, mode, model, seed, output_format, strength=None, negative_prompt=None, image_file=None): #, num_outputs=1):
+def generate_image(api_key, prompt, aspect_ratio, mode, model, seed, output_format, strength=None, negative_prompt=None, image_file=None):
     headers = {
         "authorization": f"Bearer {api_key}",
         "accept": "image/*"
@@ -19,11 +16,10 @@ def generate_image(api_key, prompt, aspect_ratio, mode, model, seed, output_form
         "prompt": prompt,
         "mode": mode,
         "model": model,
-#       "num_outputs": num_outputs,
         "output_format": output_format
     }
     
-    files = {"none": ''}  # Default to no file if none provided
+    files = {"none": ''}
 
     if mode == "text-to-image":
         files["aspect_ratio"] = (None, aspect_ratio)
@@ -31,10 +27,10 @@ def generate_image(api_key, prompt, aspect_ratio, mode, model, seed, output_form
         if seed != 0:
             data['seed'] = seed
     if negative_prompt:
-            data['negative_prompt'] = negative_prompt
+        data['negative_prompt'] = negative_prompt
     if image_file is not None:
-            image_bytes = image_file.getvalue()
-            files = {'image': ('image.jpg', image_bytes, 'image/jpeg')}
+        image_bytes = image_file.getvalue()
+        files = {'image': ('image.jpg', image_bytes, 'image/jpeg')}
     if strength is not None:
         files["strength"] = (None, str(strength))
     if seed is not None:
@@ -47,6 +43,7 @@ def generate_image(api_key, prompt, aspect_ratio, mode, model, seed, output_form
     )
     return response
 
+# Streamlit Ui
 def main():
     st.image("static/SD3_webui_logo_image.png", use_column_width=True)
     
@@ -56,8 +53,6 @@ def main():
         st.session_state.prompt = ""
     if 'negative_prompt' not in st.session_state:
         st.session_state.negative_prompt = ""
-    if 'seed' not in st.session_state:
-        st.session_state.seed = 0
     
     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)
@@ -65,45 +60,44 @@ def main():
     prompt = st.text_area("Enter positive prompt:", value=st.session_state.prompt)
     negative_prompt = st.text_input("Enter negative prompt: (optional)", value=st.session_state.negative_prompt)
     aspect_ratios = st.selectbox("Select the aspect ratio:", ["1:1", "2:3", "3:2", "4:5", "5:4", "16:9", "21:9", "9:16", "9:21"], index=0) if mode == "text-to-image" else None
-    seed = st.slider("Set seed: (randomize = 0)", min_value=0, max_value=4294967294, value=st.session_state.seed, step=1)
+    seed = st.number_input("Set seed: (randomize = 0)", min_value=0, max_value=4294967294, value=0, step=1)
     strength = st.slider("Denoising strength:", min_value=0.0, max_value=1.0, value=0.5, step=0.01) if mode == "image-to-image" else None
-#   num_outputs = st.number_input("Batch Count:", min_value=1, max_value=20, value=1, step=1)
     output_formats = st.selectbox("Select the output format:", ["jpeg", "png"], index=1)
     
     st.session_state.api_key = api_key
     st.session_state.prompt = prompt
     st.session_state.negative_prompt = negative_prompt
-    st.session_state.seed = seed
     
     uploaded_file = st.file_uploader("Upload for image-to-image generation:", type=["png", "jpg", "jpeg"]) if mode == "image-to-image" else None
 
-    if st.button('Generate Image', use_container_width=True):
+    if st.button(':blue[Generate Image]', use_container_width=True):
         with st.spinner('Generating Image...'):
-            result = generate_image(api_key, prompt, aspect_ratios, mode, models, seed, output_formats, strength, negative_prompt, uploaded_file)
+            result = generate_image(api_key, prompt, aspect_ratios, mode, models, seed, output_formats, strength, negative_prompt, uploaded_file, download)
             if result.status_code == 200:
                 image = Image.open(io.BytesIO(result.content))
                 st.image(image, use_column_width=True)
-                current_time = datetime.datetime.now().strftime("%Y%m%d_%H.%M.%S").lower()
-                file_path = os.path.join(output_dir, f"gen_{models}_{seed}_{current_time}.{output_formats}")
+                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:
+                    download = st.download_button(label=":green[Download Image]", data=file, file_name=file_path, mime=f"image/{output_formats}", use_container_width=True)
             else:
                 st.error('Failed to generate image: ' + str(result.json()))
-                
+
     with st.sidebar:
-        st.link_button("Stability.Ai _ API Documentation", "https://platform.stability.ai/docs/api-reference")
-        with st.expander(("Image Generation Costs"), expanded=True):
-            st.markdown((
-            """
+        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):
+            st.markdown("""
             - SD3 - 6.5 credits per image or $0.065
             - SD3 Turbo - 4.0 credits per image or $0.04
             
             Additional credits can be purchased via the account_page button below.
             
-            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")
+            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()