app.py

# app.py

import streamlit as st
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.")