from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline
import io
from PIL import Image

app = FastAPI()

# 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(...)):
    # 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()}