diff --git a/hairline/service.py b/hairline/service.py index f91e006..e226533 100644 --- a/hairline/service.py +++ b/hairline/service.py @@ -68,7 +68,7 @@ _REDRAW_PROMPT = os.getenv("REDRAW_PROMPT", "填充遮罩区域的头发") _REDRAW_MAX_SIDE = int(os.getenv("REDRAW_MAX_SIDE", "1024")) def _call_local_redraw(image_png_bytes, mask_png_bytes, timeout=300.0, - max_side=None, unet_name=None): + max_side=None, unet_name=None, prompt=None): """直接调 ComfyUI 重绘(替代原 local_test HTTP 服务)。 传 final 图 + 纯红遮罩 PNG,返回重绘后的 PNG bytes。 @@ -76,6 +76,7 @@ def _call_local_redraw(image_png_bytes, mask_png_bytes, timeout=300.0, max_side:送 ComfyUI 前长边压到多少像素,None 用全局默认 _REDRAW_MAX_SIDE。 unet_name:非 None 时切换 Flux 模型,None 用工作流内置默认。 + prompt:None 用默认 _REDRAW_PROMPT,否则用传入的提示词。 """ from .redraw import run_redraw eff_side = _REDRAW_MAX_SIDE if max_side is None else max_side @@ -97,7 +98,8 @@ def _call_local_redraw(image_png_bytes, mask_png_bytes, timeout=300.0, orig_w, orig_h, nw, nh, eff_side) # front=True:接口2 时延敏感,插到 ComfyUI 队列最前,避免排在接口3/5 的批量任务后面 out = run_redraw(image_png_bytes, mask_png_bytes, timeout=timeout, - prompt=_REDRAW_PROMPT, front=True, unet_name=unet_name) + prompt=prompt if prompt is not None else _REDRAW_PROMPT, + front=True, unet_name=unet_name) if scale < 1.0 and out: out = _upscale_png_to(out, orig_w, orig_h) return out