import io
import os
from fastapi import FastAPI, UploadFile, File, Response
from huggingface_hub import AsyncInferenceClient

# Model Hugging Face untuk penghapusan background
MODEL_ID = "briaai/RMBG-1.4"
HF_TOKEN = os.getenv("HF_TOKEN")
client = AsyncInferenceClient(model=MODEL_ID, token=HF_TOKEN)

app = FastAPI(title="Background Removal Proxy")

@app.get("/")
async def root():
    return {
        "ok": True,
        "message": "POST /removebg with multipart field 'image' → returns image/png",
        "model": MODEL_ID,
    }

@app.get("/healthz")
async def health():
    return {"ok": True}

@app.post("/removebg")
async def removebg(image: UploadFile = File(...)):
    img_bytes = await image.read()
    result = await client.image_to_image(img_bytes)
    buf = io.BytesIO()
    result.save(buf, format="PNG")
    return Response(content=buf.getvalue(), media_type="image/png")