48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Dict
|
|
|
|
from fastapi import APIRouter
|
|
|
|
from ..deps import OUTPUT_DIR
|
|
from ..schemas import InpaintRequest
|
|
from ..services.image_io import b64_to_pil_image, default_half_mask, pil_image_to_png_b64
|
|
from ..services.predictors import get_inpaint_predictor
|
|
|
|
router = APIRouter(tags=["inpaint"])
|
|
|
|
|
|
@router.post("/inpaint")
|
|
def inpaint(req: InpaintRequest) -> Dict[str, Any]:
|
|
try:
|
|
model_name = req.model_name or "sdxl_inpaint"
|
|
pil = b64_to_pil_image(req.image_b64).convert("RGB")
|
|
|
|
if req.mask_b64:
|
|
mask = b64_to_pil_image(req.mask_b64).convert("L")
|
|
else:
|
|
mask = default_half_mask(pil)
|
|
|
|
predictor = get_inpaint_predictor(model_name)
|
|
out = predictor(
|
|
pil,
|
|
mask,
|
|
req.prompt or "",
|
|
req.negative_prompt or "",
|
|
strength=req.strength,
|
|
max_side=req.max_side,
|
|
)
|
|
|
|
out_dir = OUTPUT_DIR / "inpaint"
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
out_path = out_dir / f"{model_name}_inpaint.png"
|
|
out.save(out_path)
|
|
|
|
return {
|
|
"success": True,
|
|
"output_path": str(out_path),
|
|
"output_image_b64": pil_image_to_png_b64(out),
|
|
}
|
|
except Exception as e:
|
|
return {"success": False, "error": str(e)}
|