import streamlit as st
from main import generate_image_from_text, upscale_image_esrgan, further_upscale_image
import asyncio

# Here we use asyncio to allow tasks to run concurrently
async def run_process():

    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 = st.progress(0)
        st.write("Generating image from text prompt...")
        image_bytes = generate_image_from_text(prompt)
        progress.progress(33)
        st.write("Upscaling image with ESRGAN...")
        upscaled_image_bytes = upscale_image_esrgan(image_bytes)
        progress.progress(66)
        st.write("Further upscaling image with GFPGAN...")
        download_link = await further_upscale_image(upscaled_image_bytes)
        progress.progress(100)
        st.markdown(download_link, unsafe_allow_html=True)

def main():
    loop = asyncio.get_event_loop()
    loop.run_until_complete(run_process())

if __name__ == "__main__":
    main()