52 lines
1.6 KiB
Python
52 lines
1.6 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 "flux_fill"
|
|
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)
|
|
call_kw: Dict[str, Any] = {"strength": req.strength, "max_side": req.max_side}
|
|
if req.guidance_scale is not None:
|
|
call_kw["guidance_scale"] = req.guidance_scale
|
|
if req.num_inference_steps is not None:
|
|
call_kw["num_inference_steps"] = req.num_inference_steps
|
|
out = predictor(
|
|
pil,
|
|
mask,
|
|
req.prompt or "",
|
|
req.negative_prompt or "",
|
|
**call_kw,
|
|
)
|
|
|
|
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)}
|