Update app.py

app.py

@@ -1,48 +1,29 @@
 # app.py
 
 import streamlit as st
-from main import generate_and_upscale_image
-from PIL import Image
-import io
+import main
 
-# 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.")
+def app():
 
-# User input for text prompt
-text_prompt = st.text_input('Enter your text prompt here:')
+    st.title("Image Generator and Upscaler")
 
-# 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")
+    text_prompt = st.text_input("Enter your text prompt")
+    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'):
+    if st.button("Generate and Upscale Image"):
+        if text_prompt and clipdrop_api_key and stability_api_key and replicate_api_token:
+            final_img = main.generate_and_upscale_image(
+                text_prompt, clipdrop_api_key, stability_api_key, replicate_api_token
+            )
 
-    st.write("Processing... Please wait, this may take a while.")
+            if final_img is not None:
+                st.image(final_img)
+            else:
+                st.write("An error occurred. Please try again.")
+        else:
+            st.write("Please enter all required fields.")
 
-    # 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.")
+if __name__ == "__main__":
+    app()