Update app.py

app.py

@@ -12,6 +12,35 @@ apply_prod_style(st)  # NOTE: Uncomment this for production!
 def V_SPACE(lines):
     for _ in range(lines):
         st.write(' ')
+        
+def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
+    if fmt not in ["jpg", "png"]:
+        raise Exception(f"Unknown image format (Available: {fmt} - case sensitive)")
+    
+    pil_format = "JPEG" if fmt == "jpg" else "PNG"
+    file_format = "jpg" if fmt == "jpg" else "png"
+    mime = "image/jpeg/?target=external" if fmt == "jpg" else "image/png/?target=external"
+    
+    buf = BytesIO()
+    pil_image.save(buf, format=pil_format)
+    
+    # Return the binary data as a string
+    return buf.getvalue()
+    
+def download_button(pil_image, filename: str, fmt: str, label="Download"):
+    if fmt not in ["jpg", "png"]:
+        raise Exception(f"Unknown image format (Available: {fmt} - case sensitive)")
+    
+    pil_format = "JPEG" if fmt == "jpg" else "PNG"
+    file_format = "jpg" if fmt == "jpg" else "png"
+    mime = "image/jpeg/?target=external" if fmt == "jpg" else "image/png/?target=external"
+    
+    buf = BytesIO()
+    pil_image.save(buf, format=pil_format)
+    
+    # Display the download button
+    href = f'<a href="data:{mime};base64,{buf.getvalue().decode("utf-8")}" download="{filename}.{file_format}">{label}</a>'
+    st.markdown(href, unsafe_allow_html=True)
 
 uploaded_file = st.file_uploader(
     label="Upload your photo here", 
@@ -42,9 +71,8 @@ if uploaded_file is not None:
         with st.expander("Success!", expanded=True):
             st.image(img_output)
             uploaded_name = os.path.splitext(uploaded_file.name)[0]
-            st.download_button(
-                label="Save Image",
-                data=img_output,
-                file_name=f'{uploaded_name}.png',
-                mime="image/png"
+            download_button(
+                pil_image=img_output,
+                filename=uploaded_name,
+                fmt="png"
             )