Update app.py

app.py

@@ -1,28 +1,48 @@
+# app.py
+
 import streamlit as st
-from main import generate_image_from_text, upscale_image_esrgan, further_upscale_image
-
-def main():
-    st.title("Image Generation and Upscaling")
-    st.write("Enter a text prompt and an image will be generated and upscaled.")
-
-    prompt = st.text_input("Enter a textual prompt to generate an image...")
-    
-    if prompt:
-        progress_bar = st.progress(0)
-        
-        st.write("Generating image from text prompt...")
-        image_bytes = generate_image_from_text(prompt)
-        progress_bar.progress(33)
-
-        st.write("Upscaling image with ESRGAN...")
-        upscaled_image_bytes = upscale_image_esrgan(image_bytes)
-        progress_bar.progress(66)
-
-        st.write("Further upscaling image with GFPGAN...")
-        download_link = further_upscale_image(upscaled_image_bytes)
-        progress_bar.progress(100)
-        
-        st.markdown(download_link, unsafe_allow_html=True)
-
-if __name__ == "__main__":
-    main()
+from main import generate_and_upscale_image
+from PIL import Image
+import io
+
+# Display title and instructions
+st.title("Text to Image Upscaling")
+st.write("Please enter a text prompt, and we will generate and upscale an image based on it.")
+
+# User input for text prompt
+text_prompt = st.text_input('Enter your text prompt here:')
+
+# User input for API keys
+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")
+
+# When the button is clicked, the processing functions are called with the user's input
+if st.button('Generate and Upscale Image'):
+
+    st.write("Processing... Please wait, this may take a while.")
+
+    # Calls function from main.py and passes user input
+    upscaled_image = generate_and_upscale_image(
+        text_prompt=text_prompt,
+        clipdrop_api_key=clipdrop_api_key,
+        stability_api_key=stability_api_key,
+        replicate_api_token=replicate_api_token,
+    )
+
+    # Show the resulting image
+    st.image(upscaled_image, caption='Your upscaled image')
+
+    # Option to download image
+    buffered = io.BytesIO()
+    upscaled_image.save(buffered, format="PNG")
+    img_str = buffered.getvalue()
+    try:
+        st.download_button(
+            label="Download Image",
+            data=img_str,
+            file_name='upscaled_image.png',
+            mime='image/png',
+        )
+    except Exception as e:
+        st.write("Error in downloading the image.")