Compare commits
2
Commits
c5e50de40a
...
326b206d08
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
326b206d08 | ||
|
|
990fae929f |
@@ -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)
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
import io
|
import io
|
||||||
|
import os
|
||||||
import os.path
|
import os.path
|
||||||
import time
|
import time
|
||||||
|
|
||||||
@@ -83,13 +84,13 @@ class ControlnetRequestImg2Img:
|
|||||||
return encoded_image
|
return encoded_image
|
||||||
|
|
||||||
|
|
||||||
def build_body_v2(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7):
|
def build_body_v2(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7, steps=None):
|
||||||
self.body = {
|
self.body = {
|
||||||
"prompt": self.prompt,
|
"prompt": self.prompt,
|
||||||
"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(steps) if steps is not None else 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,
|
||||||
@@ -413,7 +414,7 @@ def encode_numpy_to_base64(img):
|
|||||||
encoded_image = base64.b64encode(bytes).decode('utf-8')
|
encoded_image = base64.b64encode(bytes).decode('utf-8')
|
||||||
return encoded_image
|
return encoded_image
|
||||||
|
|
||||||
def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id=None, lora_material_path=None, tag="", is_hr=False, denoising_strength=0.7, inference_port="57860", refiner_switch_at=0.5):
|
def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id=None, lora_material_path=None, tag="", is_hr=False, denoising_strength=0.7, inference_port="57860", refiner_switch_at=0.5, webui_steps=None):
|
||||||
# url = "http://hairservice.tslead.net:57860/sdapi/v1/img2img"
|
# url = "http://hairservice.tslead.net:57860/sdapi/v1/img2img"
|
||||||
|
|
||||||
neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative'
|
neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative'
|
||||||
@@ -429,7 +430,7 @@ def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id
|
|||||||
control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img)
|
control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img)
|
||||||
# control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength)
|
# control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength)
|
||||||
if not is_hr:
|
if not is_hr:
|
||||||
control_net.build_body_v2(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength)
|
control_net.build_body_v2(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength, steps=webui_steps)
|
||||||
else:
|
else:
|
||||||
control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength, refiner_switch_at=refiner_switch_at)
|
control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength, refiner_switch_at=refiner_switch_at)
|
||||||
|
|
||||||
|
|||||||
@@ -606,6 +606,8 @@ def change_hairstyle_v4():
|
|||||||
else:
|
else:
|
||||||
p_tag = ""
|
p_tag = ""
|
||||||
denoising_strength = float(input_info.get('denoising_strength', 0.6)) # 接口11 可调;默认 0.6
|
denoising_strength = float(input_info.get('denoising_strength', 0.6)) # 接口11 可调;默认 0.6
|
||||||
|
webui_steps_raw = input_info.get('webui_steps', None) # 可选:webui img2img 步数,None用服务端默认
|
||||||
|
webui_steps = int(webui_steps_raw) if webui_steps_raw is not None else None
|
||||||
print(f"功能8:处理发型区域,耗时:{time.time() - start_time:.3f}s")
|
print(f"功能8:处理发型区域,耗时:{time.time() - start_time:.3f}s")
|
||||||
# cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate)
|
# cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate)
|
||||||
# cv2.imwrite(f"{task_id}_final_img.jpg", final_img)
|
# cv2.imwrite(f"{task_id}_final_img.jpg", final_img)
|
||||||
@@ -619,7 +621,7 @@ def change_hairstyle_v4():
|
|||||||
# mask_dilate = cv2.imread("/root/project/hair_service_sd/gt_mask_dilate.jpg")
|
# mask_dilate = cv2.imread("/root/project/hair_service_sd/gt_mask_dilate.jpg")
|
||||||
sd_result = webui_img2img(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id,
|
sd_result = webui_img2img(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id,
|
||||||
hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr,
|
hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr,
|
||||||
denoising_strength=denoising_strength, inference_port="57860")
|
denoising_strength=denoising_strength, inference_port="57860", webui_steps=webui_steps)
|
||||||
print(f"功能:webui,耗时:{time.time() - start_time:.3f}s")
|
print(f"功能:webui,耗时:{time.time() - start_time:.3f}s")
|
||||||
|
|
||||||
# 功能:后处理并上传结果
|
# 功能:后处理并上传结果
|
||||||
|
|||||||
Reference in New Issue
Block a user