Update app.py

app.py

@@ -1,52 +1,31 @@
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import Response
-from PIL import Image
+from fastapi.middleware.cors import CORSMiddleware
+from transformers import pipeline
 import io
-import numpy as np
-from transformers import AutoModelForImageSegmentation, AutoProcessor
-import torch
+from PIL import Image
 
 app = FastAPI()
 
-# Load the RMBG V1.4 model and processor with trust_remote_code=True
-model = AutoModelForImageSegmentation.from_pretrained(
-    "briaai/RMBG-1.4", trust_remote_code=True
-)
-processor = AutoProcessor.from_pretrained(
-    "briaai/RMBG-1.4", trust_remote_code=True
+# Allow CORS for Figma plugin
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"], # Update to specific origins in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
 @app.post("/remove-background")
 async def remove_background(file: UploadFile = File(...)):
-    try:
-        # Read uploaded image
-        image_data = await file.read()
-        image = Image.open(io.BytesIO(image_data)).convert("RGB")
-
-        # Preprocess image
-        inputs = processor(images=image, return_tensors="pt")
-
-        # Run model
-        with torch.no_grad():
-            outputs = model(**inputs)
-
-        # Post-process to get mask
-        mask = outputs.logits
-        mask = torch.sigmoid(mask).cpu().numpy()
-        mask = (mask > 0.5).astype(np.uint8) * 255
-        mask = mask.squeeze()
-
-        # Apply mask to remove background
-        image_np = np.array(image)
-        alpha_channel = mask
-        result = np.dstack((image_np, alpha_channel))
-        result_image = Image.fromarray(result, mode="RGBA")
-
-        # Save result to bytes
-        output_buffer = io.BytesIO()
-        result_image.save(output_buffer, format="PNG")
-        output_bytes = output_buffer.getvalue()
-
-        return Response(content=output_bytes, media_type="image/png")
-    except Exception as e:
-        return {"error": str(e)}
+    # Load the RMBG-1.4 model
+    pipe = pipeline('image-segmentation', 'briaai/RMBG-1.4')
+    # Read image
+    image_data = await file.read()
+    image = Image.open(io.BytesIO(image_data))
+    # Process image
+    output = pipe(image)
+    # Convert output to bytes
+    output_image = Image.fromarray(output.data)
+    output_bytes = io.BytesIO()
+    output_image.save(output_bytes, format='PNG')
+    return {"image": output_bytes.getvalue()}