Singularity666/VisionGPT

1import streamlit as st2import replicate3import os4import requests5from PIL import Image6from io import BytesIO7 8# Set up environment variable for Replicate API Token9os.environ['REPLICATE_API_TOKEN'] = 'r8_3V5WKOBwbbuL0DQGMliP0972IAVIBo62Lmi8I'  # Replace with your actual API token10 11def upscale_image(image_path):12    # Open the image file13    with open(image_path, "rb") as img_file:14        # Run the GFPGAN model15        output = replicate.run(16            "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",17            input={"img": img_file, "version": "v1.4", "scale": 16}18        )19        20        # The output is a URI of the processed image21        # We will retrieve the image data and save it22        response = requests.get(output)23        img = Image.open(BytesIO(response.content))24 25        # Save the upscaled image to a BytesIO object26        img_byte_arr = BytesIO()27        img.save(img_byte_arr, format='PNG')28        img_byte_arr = img_byte_arr.getvalue()29 30        return img, img_byte_arr31 32def main():33    st.title("Image Upscaling")34    st.write("Upload an image and it will be upscaled.")35 36    uploaded_file = st.file_uploader("Choose an image...", type="png")37    if uploaded_file is not None:38        with open("temp_img.png", "wb") as f:39            f.write(uploaded_file.getbuffer())40        st.success("Uploaded image successfully!")41        if st.button("Upscale Image"):42            img, img_bytes = upscale_image("temp_img.png")43            st.image(img, caption='Upscaled Image', use_column_width=True)44 45            # Add download button46            st.download_button(47                label="Download Upscaled Image",48                data=img_bytes,49                file_name="upscaled_image.png",50                mime="image/png"51            )52 53if __name__ == "__main__":54    main()55