Update main.py

main.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import os
 import requests
+import base64
 from PIL import Image
 from io import BytesIO
 import replicate
@@ -90,13 +91,21 @@ def further_upscale_image(image_bytes):
     try:
         print("Saving upscaled image...")
         img = Image.open(BytesIO(response.content))
-        img.save("upscaled.png")  # Save the upscaled image
+        output_file = "upscaled.png"
+        img.save(output_file)  # Save the upscaled image
     except Exception as e:
         print("Error saving upscaled image: ", e)
         raise e
 
-    return img
+    # Create a function to make download link
+    def create_download_link(file, filename):
+        with open(file, 'rb') as f:
+            bytes = f.read()
+            b64 = base64.b64encode(bytes).decode()
+            href = f'<a href="data:file/octet-stream;base64,{b64}" download="{filename}">Download File</a>'
+            return href
 
+    return create_download_link(output_file, "upscaled_image.png")
 
 def main():
     st.title("Image Generation and Upscaling")
@@ -112,9 +121,9 @@ def main():
         upscaled_image_bytes = upscale_image_esrgan(image_bytes)
         
         st.success("Further upscaling image with GFPGAN...")
-        img = further_upscale_image(upscaled_image_bytes)
+        download_link = further_upscale_image(upscaled_image_bytes)
         
-        st.image(img, caption='Upscaled Image', use_column_width=True)
+        st.markdown(download_link, unsafe_allow_html=True)
 
 if __name__ == "__main__":
-    main()
+    main()