Update app.py

app.py

@@ -5,25 +5,26 @@ import io
 import numpy as np
 
 def remove_bg(image):
-    # Convert PIL Image to bytes for rembg
+    # Convert input (could be numpy array from Gradio or file path from API)
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    elif isinstance(image, str):
+        image = Image.open(image)
+    
+    # Convert to bytes for rembg
     img_byte_arr = io.BytesIO()
     image.save(img_byte_arr, format='PNG')
-    img_bytes = img_byte_arr.getvalue()
-    
-    # Process with rembg and return as PIL Image
-    output_bytes = remove(img_bytes)
+    output_bytes = remove(img_byte_arr.getvalue())
     return Image.open(io.BytesIO(output_bytes))
 
-with gr.Blocks() as demo:
-    gr.Markdown("## U²Net Background Remover API")
-    with gr.Row():
-        inp = gr.Image(type="pil", label="Input Image")
-        out = gr.Image(label="Output (Background Removed)")
-    submit = gr.Button("Remove Background")
-    submit.click(remove_bg, inputs=inp, outputs=out)
+# Create interface - this works better for APIs than Blocks
+iface = gr.Interface(
+    fn=remove_bg,
+    inputs=gr.Image(type="filepath", label="Input Image"),
+    outputs=gr.Image(label="Output"),
+    title="U²Net Background Remover API",
+    description="Upload an image to remove the background"
+)
 
-# Key changes for API support:
-# 1. api_name="predict" enables the API endpoint
-# 2. Added proper image type conversion
-# 3. Better error handling
-demo.launch(api_name="predict")  # ← This is critical for API access
+# For Hugging Face Spaces, use this launch configuration
+iface.launch(server_name="0.0.0.0", server_port=7860)