Update app.py

app.py

@@ -50,11 +50,20 @@ def load_meta():
             return {}
     return {}
 
+def premultiply_alpha(bgra):
+    """Reduce gray/white halos on edges for nicer blending."""
+    bgr = bgra[:, :, :3].astype(np.float32) / 255.0
+    a = (bgra[:, :, 3:4].astype(np.float32) / 255.0)
+    bgr_pm = (bgr * a * 255.0).astype(np.uint8)
+    return np.dstack([bgr_pm, bgra[:, :, 3]])
+
 def load_hair_png(name):
     path = os.path.join(HAIR_DIR, name)
     hair = cv2.imread(path, cv2.IMREAD_UNCHANGED)  # BGRA
     if hair is None or hair.shape[2] != 4:
         raise ValueError(f"Invalid hair asset: {name} (must be RGBA PNG)")
+    # Improve visual quality
+    hair = premultiply_alpha(hair)
     return hair
 
 def detect_face_keypoints(img_bgr):
@@ -84,7 +93,7 @@ def hair_reference_points(hair_bgra, filename, meta):
         pts = np.array(meta[filename], dtype=np.float32)
         if pts.shape == (3, 2):
             return pts
-    # Defaults (ok for many styles). For perfect fit, put 3 points per file in meta.json.
+    # Defaults (ok for many styles). For perfect fit, add 3 points per file to meta.json.
     pL = np.array([0.30*w, 0.60*h], dtype=np.float32)
     pR = np.array([0.70*w, 0.60*h], dtype=np.float32)
     pM = np.array([0.50*w, 0.40*h], dtype=np.float32)
@@ -143,7 +152,7 @@ def apply_tryon(image, hairstyle, scale_pct, rot_deg, dx, dy, opacity, meta):
     return out_rgb, "OK"
 
 def save_png_to_tmp(img, filename="output_tryon.png"):
-    """DownloadButton needs a file path string. We save to /tmp and return that path."""
+    """Create a file in /tmp and return the path (works reliably on Spaces)."""
     if img is None:
         raise gr.Error("No image to download. Click Apply first.")
     out_path = os.path.join(tempfile.gettempdir(), filename)
@@ -205,7 +214,9 @@ def build_ui():
                         interactive=True
                     )
                     apply_btn = gr.Button("✨ Apply (Align & Overlay)")
-                    download_btn = gr.DownloadButton("⬇️ Download")  # NOTE: no file_name arg
+                    # Reliable download: a button that fills a File link.
+                    download_btn = gr.Button("⬇️ Download")
+                    download_file = gr.File(label="download", visible=False)
                 status = gr.Markdown()
 
                 with gr.Row():
@@ -236,12 +247,11 @@ def build_ui():
                     outputs=[out_img, status]
                 )
 
-                # Return a *file path* so the browser downloads with that name
-                download_btn.click(
-                    fn=lambda im: save_png_to_tmp(im, "output_tryon.png"),
-                    inputs=[out_img],
-                    outputs=[download_btn]
-                )
+                def prepare_download(im):
+                    path = save_png_to_tmp(im, "output_tryon.png")
+                    return gr.File.update(value=path, visible=True)
+
+                download_btn.click(fn=prepare_download, inputs=[out_img], outputs=[download_file])
 
                 def do_refresh():
                     files = load_hairstyles()
@@ -275,7 +285,8 @@ def build_ui():
                 out2 = gr.Image(label="Live result", height=360)
                 state_live = gr.State(None)
                 snap = gr.Button("📸 Snapshot")
-                save_live = gr.DownloadButton("⬇️ Download Snapshot")
+                save_live_btn = gr.Button("⬇️ Download Snapshot")
+                save_live_file = gr.File(label="snapshot", visible=False)
 
                 def live(im, h, s, r, dxv, dyv, op, meta):
                     res, _ = apply_tryon(im, h, s, r, dxv, dyv, op, meta)
@@ -289,11 +300,11 @@ def build_ui():
 
                 snap.click(lambda x: x, inputs=[state_live], outputs=[out2])
 
-                save_live.click(
-                    fn=lambda im: save_png_to_tmp(im, "tryon_webcam.png"),
-                    inputs=[state_live],
-                    outputs=[save_live]
-                )
+                def prepare_webcam_download(im):
+                    path = save_png_to_tmp(im, "tryon_webcam.png")
+                    return gr.File.update(value=path, visible=True)
+
+                save_live_btn.click(fn=prepare_webcam_download, inputs=[state_live], outputs=[save_live_file])
 
         return demo