# app.py

import streamlit as st
import base64
from main import generate_and_upscale_image

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("Generate and Upscale Image")

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")

if st.button("Generate and Upscale Image"):
    final_image_bytes, error = generate_and_upscale_image(text_prompt, clipdrop_api_key, stability_api_key, replicate_api_token)
    if error:
        st.error(error)
    else:
        st.markdown(get_image_download_link(final_image_bytes, "gfpgan_upscaled_image.png"), unsafe_allow_html=True)