Update app.py

app.py

@@ -1,46 +1,24 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.responses import StreamingResponse
+import gradio as gr
 from rembg import remove
+import numpy as np
 from io import BytesIO
-import gradio as gr
-
-# Create FastAPI app
-app = FastAPI()
+from fastapi import FastAPI, Response
 
-# Background removal function
-def remove_background(image_bytes):
-    try:
-        output_bytes = remove(image_bytes)
-        return output_bytes
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-# Gradio interface
-def gradio_interface(input_image):
-    if input_image is None:
-        return None
-    output_bytes = remove_background(input_image)
-    return output_bytes
-
-# Create Gradio app
+# Create both Gradio and FastAPI apps
 gradio_app = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.Image(type="filepath", label="Input Image"),
-    outputs=gr.Image(type="numpy", label="Output Image"),
+    fn=lambda img: remove(img),
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Image(type="numpy"),
     title="Background Remover",
     description="Upload an image to remove the background"
 )
 
-# API endpoint
+app = FastAPI()
+
 @app.post("/api/remove-bg")
-async def api_remove_bg(file: UploadFile = File(...)):
-    if not file.content_type.startswith('image/'):
-        raise HTTPException(status_code=400, detail="File must be an image")
-    
-    input_image = await file.read()
-    output_bytes = remove_background(input_image)
-    
-    return StreamingResponse(BytesIO(output_bytes), media_type="image/png")
+async def api_remove_bg(file: bytes = File(...)):
+    output_bytes = remove(file)
+    return Response(content=output_bytes, media_type="image/png")
 
-# Mount Gradio app to FastAPI
+# Mount Gradio to the root path
 app = gr.mount_gradio_app(app, gradio_app, path="/")