Update app.py

app.py

@@ -1,31 +1,45 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from transformers import pipeline
+from fastapi.responses import FileResponse
+import os
+import base64
 import io
-from PIL import Image
 
 app = FastAPI()
 
-# Allow CORS for Figma plugin
+# Enable CORS for Figma
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], # Update to specific origins in production
+    allow_origins=["https://*.figma.com"],
     allow_credentials=True,
-    allow_methods=["*"],
+    allow_methods=["POST", "GET"],
     allow_headers=["*"],
 )
 
+# Serve static files (index.html, assets)
+@app.get("/")
+async def serve_index():
+    return FileResponse("index.html")
+
+# API endpoint for background removal
 @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()}
+    try:
+        # Save uploaded image temporarily
+        contents = await file.read()
+        temp_path = "temp_image.png"
+        with open(temp_path, "wb") as f:
+            f.write(contents)
+
+        # 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)
+
+        img_base64 = base64.b64encode(img_data).decode("utf-8")
+        return {"image": img_base64}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))