app/server.py

"""Nukki Studio — FastAPI 서버."""
import base64
import io
from pathlib import Path

from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import HTMLResponse, JSONResponse, Response
from fastapi.staticfiles import StaticFiles
from PIL import Image

from . import config, pipeline
from .device import device_name
from .registry import REGISTRY

WEB = Path(__file__).parent / "web"
app = FastAPI(title="Nukki Studio")


@app.get("/", response_class=HTMLResponse)
def index():
    return (WEB / "index.html").read_text(encoding="utf-8")


@app.get("/api/health")
def health():
    return {
        "device": device_name(),
        "status": REGISTRY.status,
        "seg_models": list(config.SEG_MODELS.keys()),
        "modes": list(config.MODES.keys()),
        "default_mode": config.DEFAULT_MODE,
    }


@app.post("/api/warmup")
def warmup(seg_model: str = Form("general"), matte: bool = Form(True)):
    REGISTRY.segmenter(seg_model)
    if matte:
        REGISTRY.matte()
    return {"ok": True, "loaded": REGISTRY.status["loaded"]}


def _b64(image: Image.Image) -> str:
    return "data:image/png;base64," + base64.b64encode(pipeline.encode_png(image)).decode()


@app.post("/api/remove")
async def remove(
    file: UploadFile = File(...),
    mode: str = Form(config.DEFAULT_MODE),
    subject: str = Form("auto"),
    seg_model: str = Form(""),
    matte_size: str = Form(config.DEFAULT_VITMATTE),
    feather: float = Form(0.0),
    edge_shift: float = Form(0.0),
    band_scale: float = Form(1.0),
    max_side: int = Form(0),
    matte_max_side: int = Form(0),
    want_mask: bool = Form(True),
):
    raw = await file.read()
    try:
        image = Image.open(io.BytesIO(raw))
        image.load()
    except Exception as e:
        return JSONResponse({"error": f"이미지를 열 수 없습니다: {e}"}, status_code=400)

    rgba, info = pipeline.process(
        image,
        mode=mode,
        subject=subject,
        seg_model=seg_model or None,
        matte_size=matte_size,
        feather=feather,
        edge_shift=edge_shift,
        band_scale=band_scale,
        max_side=max_side,
        matte_max_side=matte_max_side,
    )

    payload = {"result": _b64(rgba), "info": info}
    if want_mask:
        mask = rgba.getchannel("A").convert("L")
        payload["mask"] = _b64(mask.convert("RGB"))
    return JSONResponse(payload)


@app.post("/api/remove.png")
async def remove_png(file: UploadFile = File(...), mode: str = Form(config.DEFAULT_MODE),
                     subject: str = Form("auto")):
    """순수 PNG 바이너리 반환 (CLI/자동화용)."""
    raw = await file.read()
    image = Image.open(io.BytesIO(raw))
    rgba, _ = pipeline.process(image, mode=mode, subject=subject)
    return Response(content=pipeline.encode_png(rgba), media_type="image/png")


# 정적 자산
app.mount("/static", StaticFiles(directory=str(WEB)), name="static")