import streamlit as st
import replicate
import os
import requests
from PIL import Image
from io import BytesIO

# Set up environment variable for Replicate API Token
os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'  # Replace with your actual API token

def upscale_image(image_path):
    # Open the image file
    with open(image_path, "rb") as img_file:
        # Run the GFPGAN model
        output = replicate.run(
            "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
            input={"img": img_file, "version": "v1.4", "scale": 16}
        )
        
        # The output is a URI of the processed image
        # We will retrieve the image data and save it
        response = requests.get(output)
        img = Image.open(BytesIO(response.content))

        # Save the upscaled image to a BytesIO object
        img_byte_arr = BytesIO()
        img.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()

        return img, img_byte_arr

def main():
    st.title("Image Upscaling")
    st.write("Upload an image and it will be upscaled.")

    uploaded_file = st.file_uploader("Choose an image...", type="png")
    if uploaded_file is not None:
        with open("temp_img.png", "wb") as f:
            f.write(uploaded_file.getbuffer())
        st.success("Uploaded image successfully!")
        if st.button("Upscale Image"):
            img, img_bytes = upscale_image("temp_img.png")
            st.image(img, caption='Upscaled Image', use_column_width=True)

            # Add download button
            st.download_button(
                label="Download Upscaled Image",
                data=img_bytes,
                file_name="upscaled_image.png",
                mime="image/png"
            )

if __name__ == "__main__":
    main()