Update app.py

app.py

@@ -1,41 +1,24 @@
 import streamlit as st
-import replicate
-import os
-import requests
-from PIL import Image
-from io import BytesIO
+from main import upscale_image
+import base64
 
-# 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))
-        img.save("upscaled.png")  # Save the upscaled image
-        return img
-
-def main():
-    st.title("Image Upscaling")
-    st.write("Upload an image and it will be upscaled.")
+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
 
+def app():
     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 = upscale_image("temp_img.png")
-            st.image(img, caption='Upscaled Image', use_column_width=True)
+        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)
 
 if __name__ == "__main__":
-    main()
+    app()