Update app.py

app.py

@@ -1,9 +1,10 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse
-import os
-import base64
+from PIL import Image
 import io
+import torch
+from transformers import pipeline
+import base64
 
 app = FastAPI()
 
@@ -12,34 +13,29 @@ app.add_middleware(
     CORSMiddleware,
     allow_origins=["https://*.figma.com"],
     allow_credentials=True,
-    allow_methods=["POST", "GET"],
+    allow_methods=["POST"],
     allow_headers=["*"],
 )
 
-# Serve static files (index.html, assets)
-@app.get("/")
-async def serve_index():
-    return FileResponse("index.html")
+# Load the model
+pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
 
-# API endpoint for background removal
 @app.post("/remove-background")
 async def remove_background(file: UploadFile = File(...)):
     try:
-        # Save uploaded image temporarily
+        # Read image
         contents = await file.read()
-        temp_path = "temp_image.png"
-        with open(temp_path, "wb") as f:
-            f.write(contents)
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+
+        # Process image
+        result = pipe(image)
+        output = result["output"]
 
-        # The client-side JS (index-E_M5nW8h.js) handles processing
-        # This endpoint assumes the processed image is saved by JS
-        # For simplicity, return the uploaded image as base64
-        # In practice, you'd need a way to trigger JS and get the result
-        with open(temp_path, "rb") as f:
-            img_data = f.read()
-        os.remove(temp_path)
+        # Convert to base64
+        buffered = io.BytesIO()
+        output.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-        img_base64 = base64.b64encode(img_data).decode("utf-8")
-        return {"image": img_base64}
+        return {"image": img_str}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))