Update app.py

app.py

@@ -1,24 +1,54 @@
 import streamlit as st
-from main import upscale_image
-import base64
+import replicate
+import os
+import requests
+from PIL import Image
+from io import BytesIO
 
-def get_image_download_link(img_path, filename="upscaled_image.png"):
-    with open(img_path, 'rb') as f:
-        bytes = f.read()
-        b64 = base64.b64encode(bytes).decode()
-        href = f'<a href="data:file/octet-stream;base64,{b64}" download=\'{filename}\'>Click here to download the upscaled image</a>'
-        return href
+# 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.")
 
-def app():
     uploaded_file = st.file_uploader("Choose an image...", type="png")
     if uploaded_file is not None:
-        with open("uploaded.png", "wb") as f:
-            f.write(uploaded_file.getvalue())
-        
-        upscale_image("uploaded.png")  # This will create "upscaled.png"
-        
-        st.image("upscaled.png", caption="Upscaled Image", use_column_width=True)
-        st.markdown(get_image_download_link("upscaled.png"), unsafe_allow_html=True)
+        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__":
-    app()
+    main()