Update app.py

app.py

@@ -1,10 +1,11 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from PIL import Image
-import io
-import torch
-from transformers import pipeline
+from fastapi.responses import FileResponse
+from playwright.async_api import async_playwright
+import os
 import base64
+import io
+import time
 
 app = FastAPI()
 
@@ -13,29 +14,61 @@ app.add_middleware(
     CORSMiddleware,
     allow_origins=["https://*.figma.com"],
     allow_credentials=True,
-    allow_methods=["POST"],
+    allow_methods=["POST", "GET"],
     allow_headers=["*"],
 )
 
-# Load the model
-pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+# 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(...)):
     try:
-        # Read image
+        # Save uploaded image
         contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        temp_path = "temp_image.png"
+        with open(temp_path, "wb") as f:
+            f.write(contents)
+
+        # Use Playwright to process image via index.html
+        async with async_playwright() as p:
+            browser = await p.chromium.launch()
+            page = await browser.new_page()
+
+            # Load the Space's frontend
+            await page.goto("http://localhost:7860/")
+
+            # Upload image
+            await page.set_input_files("#upload", temp_path)
+
+            # Wait for processing (adjust based on JS behavior)
+            await page.wait_for_timeout(5000)  # Adjust if needed
+
+            # Extract processed image (assumes JS sets it in #status or similar)
+            # This depends on index-E_M5nW8h.js output (e.g., canvas or img)
+            processed_image = await page.evaluate("""
+                () => {
+                    const canvas = document.querySelector('canvas'); // Adjust selector
+                    if (canvas) {
+                        return canvas.toDataURL('image/png').split(',')[1];
+                    }
+                    return null;
+                }
+            """)
+
+            await browser.close()
 
-        # Process image
-        result = pipe(image)
-        output = result["output"]
+        # Clean up
+        os.remove(temp_path)
 
-        # Convert to base64
-        buffered = io.BytesIO()
-        output.save(buffered, format="PNG")
-        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        if not processed_image:
+            raise HTTPException(status_code=500, detail="Failed to process image")
 
-        return {"image": img_str}
+        return {"image": processed_image}
     except Exception as e:
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
         raise HTTPException(status_code=500, detail=str(e))