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

perf(swapHair): webui降步数 + 用户图预处理内存缓存

Browse Source 
优化1 - webui img2img steps 20→15 (gen_super_image.py):
- build_body_v2 的 steps 改为环境变量 WEBUI_STEPS 可配，默认15
- DPM++ 2M Karras 15步对换发型质量影响可忽略，省~0.1s

优化2 - infer_hairstyle_diy_jy 用户图预处理内存缓存 (hairstyle_model.py):
- 新增 _user_prepare_cache 进程内LRU缓存(8张图上限)
- key=图片md5哈希+ratio，命中时跳过landmark检测+get_prepare_user_768_data整条GPU管线
- 接口2女性多发型场景: 同一张用户图第2个发型起命中，功能6从1.8s降至1.1s(省0.7s)
- 缓存命中时补写磁盘文件(task_id每次不同,下游功能7仍从磁盘读)
- 顺带修复: user_matting_8uc3_bald_orisize 为None时写user_orig_mask.png的crash
- 加分步计时日志(主GAN/融合耗时),便于定位热点

实测: 同图连续请求 swapHair 从4.28s降至3.5s(省18%)
This commit is contained in:
xsl
2026-07-26 16:41:19 +08:00
parent c5e50de40a
commit 990fae929f
2 changed files with 108 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
project/hair_service_sd/core/hairstyle_model.py
+106 -39
View File
@@ -90,6 +90,14 @@ class HairStyle_Model(object):
self.keypoints_processor = hair_init.keypoints_processor self.keypoints_processor = hair_init.keypoints_processor
self.human_keypoint = hair_init.human_keypoint self.human_keypoint = hair_init.human_keypoint
# 用户图预处理产物内存缓存（一级缓存，避免重复 GPU 推理 + 磁盘 IO）。
# key = 图片字节哈希 + ratio；value = 7 个产物 + landmarks_1k 的 dict。
# 接口2 女性多发型场景：同一张用户图连续请求，第二次起命中缓存省 ~1.6s。
import collections
self._user_prepare_cache = {} # {cache_key: {产物dict}}
self._user_prepare_cache_keys = collections.deque() # LRU 顺序
self._USER_CACHE_MAX = 8 # 最多缓存 8 张图（防止内存膨胀）
# worker_id = int(os.environ.get('APP_WORKER_ID', 1)) # worker_id = int(os.environ.get('APP_WORKER_ID', 1))
# rand_max = 9527 # rand_max = 9527
@@ -2268,48 +2276,99 @@ class HairStyle_Model(object):
ref_landmark_f1k2_768) ref_landmark_f1k2_768)
cv2.imwrite(another_pose_hair_image_dir, another_pose_hair_image*255) cv2.imwrite(another_pose_hair_image_dir, another_pose_hair_image*255)
landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt')
if not osp.exists(landmark1k_dir):
landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward_diy(user_rgb_8uc3_orisize) # ===== 内存缓存（一级）：key = 图片哈希 + ratio =====
# 命中则跳过 landmark 检测 + get_prepare_user_768_data 整条 GPU 管线（省 ~1.6s）
import hashlib as _hashlib
_img_hash = _hashlib.md5(user_rgb_8uc3_orisize.tobytes()).hexdigest()[:16]
_cache_key = f"{_img_hash}_r{ratio}"
_cached = self._user_prepare_cache.get(_cache_key)
if _cached is not None:
# 命中内存缓存：直接取所有产物，跳过 landmark 检测和预处理管线
landmarks_origin_img_1k = _cached['landmarks_1k']
user_bald_res_8uc3_orisize = _cached['bald_res']
user_baldseg_8uc3_orisize = _cached['baldseg_ori']
user_baldseg_8uc3_768 = _cached['baldseg_768']
user_bald_8uc3_768 = _cached['bald_768']
user_landmark_f1k2_768 = _cached['lmk_768']
user_hairstyle_M = _cached['M']
user_matting_8uc3_bald_orisize = _cached['matting_ori']
# 补写磁盘文件：下游代码（功能7等）仍从 userinfo_dir 读这些文件，
# 而 task_id 每次不同导致 userinfo_dir 不同，必须补写保证下游可用
os.makedirs(userinfo_dir, exist_ok=True)
np.savetxt(landmark1k_dir, landmarks_origin_img_1k)
cv2.imwrite(osp.join(userinfo_dir, 'bald_res_ori.png'), user_bald_res_8uc3_orisize)
cv2.imwrite(osp.join(userinfo_dir, 'bald_seg_ori.png'), user_baldseg_8uc3_orisize)
cv2.imwrite(osp.join(userinfo_dir, 'user_baldseg_768.png'), user_baldseg_8uc3_768)
cv2.imwrite(osp.join(userinfo_dir, 'bald_seg_768.png'), user_bald_8uc3_768)
np.savetxt(osp.join(userinfo_dir, 'landmark_f1k2_768.txt'), user_landmark_f1k2_768)
np.savetxt(osp.join(userinfo_dir, 'hairstyle_M.txt'), user_hairstyle_M)
if user_matting_8uc3_bald_orisize is not None:
cv2.imwrite(osp.join(userinfo_dir, 'user_orig_mask.png'), user_matting_8uc3_bald_orisize)
self.logger_process.info(f"内存缓存命中 key={_cache_key}，跳过用户图预处理(补写磁盘文件)")
else:
# 未命中：走原有逻辑（landmark 检测 + 预处理管线 + 磁盘缓存）
if not osp.exists(landmark1k_dir):
landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward_diy(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None:
return None, 10001, None, None, None
np.savetxt(landmark1k_dir, landmarks_origin_img_1k)
else:
landmarks_origin_img_1k = np.loadtxt(landmark1k_dir)
# landmarks_origin_img_1k, _, _ = self.get_landmark.forward(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None: if landmarks_origin_img_1k is None:
return None, 10001, None, None, None return None, 10001, None, None, None
np.savetxt(landmark1k_dir, landmarks_origin_img_1k)
else:
landmarks_origin_img_1k = np.loadtxt(landmark1k_dir)
# landmarks_origin_img_1k, _, _ = self.get_landmark.forward(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None:
return None, 10001, None, None, None
user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png') user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png')
user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png') user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png')
user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png') user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png')
user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png') user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png')
user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt') user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt')
user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt') user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt')
condition_exist2 = True condition_exist2 = True
pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir, pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir,
user_landmark_f1k2_768_dir, user_hairstyle_M_dir] user_landmark_f1k2_768_dir, user_hairstyle_M_dir]
for tmp_dir in pre_list: for tmp_dir in pre_list:
if not osp.exists(tmp_dir): if not osp.exists(tmp_dir):
condition_exist2 = False condition_exist2 = False
user_matting_8uc3_bald_orisize = None user_matting_8uc3_bald_orisize = None
if not condition_exist2: if not condition_exist2:
user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \ user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \
user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio)
cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize) cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize)
cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize) cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize)
cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768) cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768)
cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768) cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768)
np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768) np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768)
np.savetxt(user_hairstyle_M_dir, user_hairstyle_M) np.savetxt(user_hairstyle_M_dir, user_hairstyle_M)
else: else:
user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir) user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir)
user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir) user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir)
user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir) user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir)
user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir) user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir)
user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir) user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir)
user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir) user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir)
# 写入内存缓存（含 landmarks_1k 和 matting，供后续同图请求命中）
self._user_prepare_cache[_cache_key] = {
'landmarks_1k': landmarks_origin_img_1k,
'bald_res': user_bald_res_8uc3_orisize,
'baldseg_ori': user_baldseg_8uc3_orisize,
'baldseg_768': user_baldseg_8uc3_768,
'bald_768': user_bald_8uc3_768,
'lmk_768': user_landmark_f1k2_768,
'M': user_hairstyle_M,
'matting_ori': user_matting_8uc3_bald_orisize,
}
self._user_prepare_cache_keys.append(_cache_key)
# LRU 淘汰：超过上限删除最老的
while len(self._user_prepare_cache_keys) > self._USER_CACHE_MAX:
_old = self._user_prepare_cache_keys.popleft()
self._user_prepare_cache.pop(_old, None)
self.logger_process.info(f"内存缓存写入 key={_cache_key}，当前缓存 {len(self._user_prepare_cache)} 张")
# show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1)
# ratio = 1536. / max(show_concat.shape[:2]) # ratio = 1536. / max(show_concat.shape[:2])
@@ -2318,18 +2377,26 @@ class HairStyle_Model(object):
# cv2.waitKey() # cv2.waitKey()
user_orig_mask_path = os.path.join(userinfo_dir, "user_orig_mask.png") user_orig_mask_path = os.path.join(userinfo_dir, "user_orig_mask.png")
if not os.path.exists(user_orig_mask_path): if user_matting_8uc3_bald_orisize is not None and not os.path.exists(user_orig_mask_path):
cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize) cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize)
# 换发型 # 换发型（主 GAN）
import time as _time
_t_gan0 = _time.perf_counter()
hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image,
user_baldseg_8uc3_768, user_baldseg_8uc3_768,
user_bald_8uc3_768, user_bald_8uc3_768,
user_landmark_f1k2_768, gender) user_landmark_f1k2_768, gender)
_t_gan = _time.perf_counter() - _t_gan0
# 融合（第3次matte + 融合GAN）
_t_fuse0 = _time.perf_counter()
hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste(
user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M)
_t_fuse = _time.perf_counter() - _t_fuse0
self.logger_process.info(
f"功能6分步计时: 主GAN={_t_gan:.3f}s 融合={_t_fuse:.3f}s (缓存={'命中' if _cached is not None else '未命中'})")
gen_hair_mask_path_2 = os.path.join(userinfo_dir, "hair_mask_2.png") gen_hair_mask_path_2 = os.path.join(userinfo_dir, "hair_mask_2.png")
cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize) cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize)
project/hair_service_sd/gen_super_image.py
+2 -1
View File
@@ -1,4 +1,5 @@
import io import io
import os
import os.path import os.path
import time import time
@@ -89,7 +90,7 @@ class ControlnetRequestImg2Img:
"negative_prompt": self.neg_prompt, "negative_prompt": self.neg_prompt,
"sampler_name": "DPM++ 2M Karras", "sampler_name": "DPM++ 2M Karras",
"batch_size": 1, "batch_size": 1,
"steps": 20, "steps": int(os.environ.get("WEBUI_STEPS", "15")),
"width": dst_width, "width": dst_width,
"height": dst_height, "height": dst_height,
"cfg_scale": cfg_scale, "cfg_scale": cfg_scale,