Update main.py

main.py

@@ -6,6 +6,13 @@ import aiofiles
 import aiohttp
 import os
 from bs4 import BeautifulSoup
+
+from pydantic import BaseModel
+import base64
+from rembg import remove
+from io import BytesIO
+from PIL import Image
+
 from akenoai import OldAkenoXToJs
 from akenoai.runner import run_fast
 from akenoai.clients import create_pyrogram
@@ -29,6 +36,9 @@ logger.setLevel(logging.DEBUG)
 js = OldAkenoXToJs()
 fast_app = js.get_app()
 
+os.environ["U2NET_HOME"] = "/usr/src/app/.u2net"
+os.environ["XDG_CACHE_HOME"] = "/usr/src/app/.cache"
+
 origins = [
     "https://randydev-meta-ai.hf.space"
 ]
@@ -107,6 +117,23 @@ async def upload_to_catbox(dl_path: str) -> str:
             response.raise_for_status()
             return (await response.text()).strip()
 
+class ImageInput(BaseModel):
+    image_base64: str
+
+@fast_app.post("/remove-bg-base64")
+async def remove_bg_base64(data: ImageInput):
+    try:
+        image_data = base64.b64decode(data.image_base64)
+        input_image = Image.open(BytesIO(image_data))
+        output_image = remove(input_image)
+
+        buffer = BytesIO()
+        output_image.save(buffer, format="PNG")
+        output_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        return {"output_image_base64": output_base64}
+    except Exception as e:
+        return {"error": str(e)}
+
 @fast_app.get("/image/share/{filename}")
 async def get_image(filename: str):
     file_path = f"image/{filename}"