Update app.py

app.py

@@ -1,30 +1,21 @@
 # app.py
 
 import streamlit as st
+import base64
 from main import generate_and_upscale_image
-import os
 
-def main():
+def get_image_download_link(img_bytes, filename):
+    b64 = base64.b64encode(img_bytes).decode()
+    href = f'<a href="data:image/png;base64,{b64}" download="{filename}">Click to download</a>'
+    return href
 
-    st.title("Image Generator and Upscaler")
+st.title("Generate and Upscale Image")
 
-    text_prompt = st.text_input("Enter your text prompt")
+text_prompt = st.text_input("Enter a text prompt")
+clipdrop_api_key = st.text_input("Enter your Clipdrop API key")
+stability_api_key = st.text_input("Enter your Stability API key")
+replicate_api_token = st.text_input("Enter your Replicate API token")
 
-    clipdrop_api_key = st.text_input("Enter your ClipDrop API Key", type="password")
-    stability_api_key = st.text_input("Enter your Stability API Key", type="password")
-    replicate_api_token = st.text_input("Enter your Replicate API Token", type="password")
-
-    download_path = st.text_input("Enter the path where the image will be saved")
-    if not os.path.exists(download_path):
-        st.error("The provided download path does not exist. Please provide a valid path.")
-        return
-
-    if st.button("Generate and Upscale"):
-        try:
-            generate_and_upscale_image(text_prompt, clipdrop_api_key, stability_api_key, replicate_api_token, download_path)
-            st.success("The image has been successfully generated, upscaled, and saved to the specified path.")
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
-
-if __name__ == "__main__":
-    main()
+if st.button("Generate and Upscale Image"):
+    final_image_bytes = generate_and_upscale_image(text_prompt, clipdrop_api_key, stability_api_key, replicate_api_token)
+    st.markdown(get_image_download_link(final_image_bytes, "gfpgan_upscaled_image.png"), unsafe_allow_html=True)