Update app.py

app.py

@@ -12,35 +12,26 @@ 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)
+
+    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", 
@@ -53,8 +44,9 @@ if uploaded_file is not None:
     in_submit = st.button("Submit")
 
     if uploaded_file is not None and in_submit:
-        img_input = Image.open(uploaded_file)
-        
+        # Use BytesIO to avoid decoding error
+        img_input = Image.open(BytesIO(uploaded_file.read()))
+
         with st.spinner("AI is doing magic to your photo. Please wait..."):
             hexmap = {
                 "Transparent (PNG)": "#000000",
@@ -67,11 +59,11 @@ if uploaded_file is not None:
             alpha = 0.0 if in_mode == "Transparent (PNG)" else 1.0
             img_matte = matte(img_input)
             img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=hexmap[in_mode])
-        
+
         with st.expander("Success!", expanded=True):
             st.image(img_output)
             uploaded_name = os.path.splitext(uploaded_file.name)[0]
-            download_button(
+            image_download_button(
                 pil_image=img_output,
                 filename=uploaded_name,
                 fmt="png"