Update app.py

app.py

@@ -12,26 +12,6 @@ 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 st.download_button(
-        label=label,
-        data=buf.getvalue(),
-        file_name=f'{filename}.{file_format}',
-        mime=mime,
-    )
-
-
 
 uploaded_file = st.file_uploader(
     label="Upload your photo here", 
@@ -40,8 +20,6 @@ uploaded_file = st.file_uploader(
 
 if uploaded_file is not None:
 
-   
-
     in_mode = "Transparent (PNG)"
     in_submit = st.button("Submit")
 
@@ -64,8 +42,9 @@ 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]
-            image_download_button(
-                pil_image=img_output,
-                filename=uploaded_name,
-                fmt="png"
-            )
+            st.download_button(
+                label="Save Image",
+                data=img_output,
+                file_name=f'{uploaded_name}.png',
+                mime="image/png"
+            )