Update app.py

app.py

@@ -61,7 +61,7 @@ birefnet.to(device)
 birefnet.eval()
 
 @spaces.GPU
-def predict(image):
+def remove_background(image):
     if image is None:
         raise gr.Error("Please upload an image.")
 
@@ -84,16 +84,59 @@ def predict(image):
 
     torch.cuda.empty_cache()
 
-    # Save as PNG
+    # Save mask as PNG
+    mask_path = "mask.png"
+    pred_pil.save(mask_path)
+
+    # Save output as PNG
     output_path = "output.png"
     image_masked.save(output_path)
 
-    return output_path
+    return mask_path, output_path
+
+css = """
+body {
+    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif, "Apple Color Emoji", "Segoe UI Emoji", "Segoe UI Symbol";
+}
+.gradio-container {
+    background: linear-gradient(
+        135deg,
+        #e0f7fa, #e8f5e9, #fff9c4, #ffebee,
+        #f3e5f5, #e1f5fe, #fff3e0, #e8eaf6
+    );
+    background-size: 400% 400%;
+    animation: gradient-animation 15s ease infinite;
+}
+@keyframes gradient-animation {
+    0% { background-position: 0% 50%; }
+    50% { background-position: 100% 50%; }
+    100% { background-position: 0% 50%; }
+}
+.gradio-button {
+    font-family: inherit;
+    font-size: 16px;
+    font-weight: bold;
+    color: #000000;
+    background: white;
+    border: 2px solid black;
+    border-radius: 10px;
+}
+.gradio-button:hover {
+    background: #f0f0f0;
+}
+"""
 
 iface = gr.Interface(
-    fn=predict,
+    fn=remove_background,
     inputs=gr.Image(type="numpy"),
-    outputs=gr.Image(type="filepath"),
+    outputs=[
+        gr.Image(type="filepath", label="Mask"),
+        gr.Image(type="filepath", label="Output")
+    ],
+    title="<div style='font-size: 36px; font-weight: bold;'>{.Remove Background}</div>",
+    description="Upload an image to remove its background using BiRefNet.",
+    allow_flagging="never",
+    css=css
 )
 
 if __name__ == "__main__":