xsl/hair
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 适配 RTX 3090 (24GB) 环境优化

Browse Source 
硬件迁移：从 RTX 5090 (32GB) 迁移到 RTX 3090 (24GB)

主要改动：
1. hairline/comfyui.py: 轮询间隔从 0.2s 降到 0.05s
2. hairline/service.py: PNG 编码 compress_level=1，节省 ~240ms
3. add_hair.json: 工作流使用 4B FP8 模型 + steps=4
4. static/test_interface3.html: 修复图片显示（添加 data:image/jpeg;base64, 前缀）

性能优化后接口3响应时间：6.6-7.3s（之前 8.66s）
This commit is contained in:
xsl
2026-07-18 18:58:40 +08:00
parent 659c037270
commit 74ccab0ff8
9 changed files with 311 additions and 563 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hairline/comfyui.py
+2 -2
View File
@@ -18,7 +18,7 @@ import uuid
import httpx
COMFYUI_URL = os.getenv("COMFYUI_URL", "http://10.60.74.221:8188").rstrip("/")
COMFYUI_URL = os.getenv("COMFYUI_URL", "http://127.0.0.1:8188").rstrip("/")
_WORKFLOW_DEFAULT = os.getenv(
"ADD_HAIR_WORKFLOW",
os.path.join(os.path.dirname(os.path.dirname(__file__)), "add_hair.json"),
@@ -145,7 +145,7 @@ def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = Non
outputs = entry.get("outputs")
if outputs and output_node in outputs:
break
time.sleep(1.0)
time.sleep(0.05)
if not outputs or output_node not in outputs:
raise TimeoutError(f"ComfyUI 出图超时({timeout}s) prompt_id={prompt_id}")
hairline/service.py
+20 -30
View File
@@ -40,27 +40,14 @@ _REPO = os.path.dirname(os.path.dirname(__file__))
_TEXTURE_DIR = os.path.join(_REPO, "hairline_texture")
_BLACK_TEXTURE_DIR = os.path.join(_REPO, "hairline_texture_black")
# 外部重绘服务（local_test，0716add-hair.json 工作流）。接口2 female 用它替代原 Flux-2 重绘。
# 重绘服务已独立到远程机器；可用 HAIR_LOCAL_REDRAW_URL 覆盖。
# gpu_worker 同内网走 10.60.74.221，外网（本地开发机）需走公网 117.50.183.232。
_LOCAL_REDRAW_URL = os.getenv("HAIR_LOCAL_REDRAW_URL", "http://10.60.74.221:8899").rstrip("/")
def _call_local_redraw(image_png_bytes, mask_png_bytes, timeout=300.0):
"""调外部重绘服务（local_test /api/generate）：传 final 图 + 纯红遮罩 PNG，
返回重绘后的 PNG bytes。失败抛异常（调用方负责 try/except 跳过）。
"""直接调 ComfyUI 重绘（替代原 local_test HTTP 服务）。
传 final 图 + 纯红遮罩 PNG，返回重绘后的 PNG bytes。
失败抛异常（调用方负责 try/except 跳过）。
"""
import requests
files = {
"image": ("final.jpg", image_png_bytes, "image/jpeg"),
"mask": ("mask.png", mask_png_bytes, "image/png"),
}
resp = requests.post(_LOCAL_REDRAW_URL + "/api/generate", files=files, timeout=timeout)
resp.raise_for_status()
if "image/" not in resp.headers.get("Content-Type", ""):
# 服务返回了 JSON 错误
raise RuntimeError(f"local_test 返回非图片: {resp.text[:200]}")
return resp.content
from .redraw import run_redraw
return run_redraw(image_png_bytes, mask_png_bytes, timeout=timeout)
# 发际线贴图档位：middle=默认(hairline_texture/)，high/low 各自独立文件夹。
_TEXTURE_DIRS = {
@@ -69,9 +56,8 @@ _TEXTURE_DIRS = {
"low": os.path.join(_REPO, "hairline_texture_low"),
}
# ⚠️ 本 worker 是 RTX 5090(sm_120)，torch 2.2.2(cu121) 只编到 sm_90，CUDA 跑算子会报
# "no kernel image"。SegFormer 默认走 CPU（~2.5s/张）。换 torch cu128 后可设 SEG_DEVICE=cuda。
_SEG_DEVICE = os.getenv("SEG_DEVICE", "cpu")
# torch 2.7.1+cu128 已支持 RTX 5090 (sm_120)，SegFormer 走 GPU（~0.05s/张）
_SEG_DEVICE = os.getenv("SEG_DEVICE", "cuda")
_landmarker = None
_parser = None
@@ -205,7 +191,7 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
try:
h, w = image_bgr.shape[:2]
buf = io.BytesIO()
compose_comfy_rgba(image_bgr, np.zeros((h, w), np.uint8)).save(buf, format="PNG")
compose_comfy_rgba(image_bgr, np.zeros((h, w), np.uint8)).save(buf, format="PNG", compress_level=1)
shared_grown = comfyui.run(buf.getvalue(), prompt=prompt, workflow_path=workflow_path)
except Exception as e: # noqa: BLE001
logger.warning("接口2 生发图失败(无遮罩)：%s", e)
@@ -225,7 +211,7 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
marked, mask = build_inpaint_mask(
image_bgr, ctx["landmarks"], ctx["parse_map"], ctx["points"], black)
buf = io.BytesIO()
compose_comfy_rgba(marked, mask).save(buf, format="PNG")
compose_comfy_rgba(marked, mask).save(buf, format="PNG", compress_level=1)
grown_png = comfyui.run(buf.getvalue(), prompt=prompt, workflow_path=workflow_path)
except Exception as e: # noqa: BLE001 单张失败不拖垮整请求
logger.warning("接口2 生发图失败 type=%s：%s", key, e)
@@ -241,20 +227,24 @@ def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | N
grown 图来源（新流程）：对每个选中发型把 female key 映射到 change_hair 的 chang_* hair_id，
调 face_analysis.hairline_grow.generate_hairline_redraw（= 接口12 final 管线，参数用
redraw_defaults）拿到 ④ final（接缝融合基底）+ ⑤-② 纯红遮罩 PNG，再**后端调外部重绘服务
local_test**（0716add-hair.json 工作流）完成发际线带重绘，重绘结果作为生发图。
redraw_defaults）拿到 ④ final（接缝融合基底）+ ⑤-② 纯红遮罩 PNG，再**后端直接调
ComfyUI**（0716add-hair-api.json 工作流）完成发际线带重绘，重绘结果作为生发图。
overlay 仍是发际线曲线透明层（与 generate_grow_results 完全一致）。
Returns: list[dict] {"hairline_type","order","overlay","grown_png"(jpg bytes 或 None)}；
无人脸返回 None。单个发型换发型/重绘失败时 grown_png=None，不抛异常。
"""
from face_analysis.hairline_grow import generate_hairline_redraw, NoFaceError
from face_analysis.head_mask import SEGFORMER_HAIR
ctx = extract_context(image_bgr)
if ctx is None:
return None
uv, ext_faces = load_ext_mesh()
# 复用 extract_context 已算好的 SegFormer parse_map，避免 generate_hairline_redraw 内部重复分割
hair_mask_reuse = (ctx["parse_map"] == SEGFORMER_HAIR)
textures = get_texture_map()["female"] # [(key, path), ...] 已排序
if hair_styles is not None:
items = [(s, textures[s - 1]) for s in hair_styles]
@@ -273,7 +263,7 @@ def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | N
logger.warning("接口2 换发型：female key=%s 无对应 chang_id，跳过生发图", key)
else:
try:
data = generate_hairline_redraw(image_bgr, chang_id, **redraw_defaults)
data = generate_hairline_redraw(image_bgr, chang_id, hair_mask=hair_mask_reuse, **redraw_defaults)
steps = data.get("steps") or {}
# ④ final（接缝融合基底）+ ⑤-② 纯红遮罩 PNG
final_b64 = steps.get("final_base64") or ""
@@ -289,7 +279,7 @@ def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | N
mask_b64 = mask_b64.split(",", 1)[1]
final_bytes = base64.b64decode(final_b64)
mask_bytes = base64.b64decode(mask_b64)
# 后端调外部重绘服务（local_test），返回重绘后的 PNG
# 后端直接调 ComfyUI 重绘，返回重绘后的 PNG
grown_png = _call_local_redraw(final_bytes, mask_bytes)
if grown_png is None:
logger.warning("接口2 换发型：type=%s 重绘结果为空", key)
@@ -319,7 +309,7 @@ def _grow_from_texture(image_bgr: np.ndarray, ctx: dict, white_path: str | None,
h, w = image_bgr.shape[:2]
marked, mask = image_bgr, np.zeros((h, w), np.uint8)
buf = io.BytesIO()
compose_comfy_rgba(marked, mask).save(buf, format="PNG")
compose_comfy_rgba(marked, mask).save(buf, format="PNG", compress_level=1)
return comfyui.run(buf.getvalue(), prompt=prompt)
except Exception as e: # noqa: BLE001 单张失败不拖垮整请求
logger.warning("接口5 生发图失败：%s", e)
@@ -416,7 +406,7 @@ def generate_grow_b(marked_bgr: np.ndarray, use_mask: bool = True, prompt: str =
mask = np.zeros((h, w), np.uint8) # 空遮罩：alpha 全 255，跳过检测
buf = io.BytesIO()
compose_comfy_rgba(marked_bgr, mask).save(buf, format="PNG") # marked 原图 + 遮罩
compose_comfy_rgba(marked_bgr, mask).save(buf, format="PNG", compress_level=1) # marked 原图 + 遮罩
grown_png = comfyui.run(buf.getvalue(), prompt=prompt)
return {"grown_png": grown_png, "status": "ok"}