feat: 适配 RTX 3090 (24GB) 环境优化
硬件迁移：从 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:
+20
-30
@@ -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"}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user