xsl/change_hair
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
990fae929fd8f6415af0ca8a190b67de2cc90e43
change_hair/project/hair_service_sd/gen_super_image.py
T
xsl 990fae929f perf(swapHair): webui降步数 + 用户图预处理内存缓存 
优化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%)
2026-07-26 16:41:19 +08:00

518 lines
20 KiB
Python
Raw Blame History

import io
import os
import os.path
import time
import cv2
import base64
import requests
from PIL import Image
import numpy as np
import json
from utils.call_hair_inter import call_hair_infer
from utils.call_hair_inter import call_hair_infer_diy
from common.logger import config
from uuid import uuid4
user_img_tmp_dir = config.get('default', 'tmp_dir')
version = config.get('default', 'version')
if version == "local":
current_webui_url = 'http://192.168.1.57:57860/'
else:
current_webui_url = 'http://0.0.0.0:57860/'
class WebUISupersuperResolution:
def __init__(self):
self.url = f"{current_webui_url}sdapi/v1/extra-single-image"
self.body = None
def encode_image_to_base64(self, img):
retval, bytes = cv2.imencode('.png', img)
encoded_image = base64.b64encode(bytes).decode('utf-8')
return encoded_image
def send_request(self):
response = requests.post(url=self.url, json=self.body)
return response.json()
# 图像初步超分
def build_body(self, base_img):
self.body = {
# "show_extras_results": True,
# "gfpgan_visibility": 0,
"codeformer_visibility": 1,
"codeformer_weight": 1,
"upscaling_resize": 2,
# "upscaling_resize_w": 512,
# "upscaling_resize_h": 512,
"upscaling_crop": True,
"upscaler_1": "8x_NMKD-Superscale_150000_G",
"upscaler_2": "None",
"extras_upscaler_2_visibility": 0,
"image": self.encode_image_to_base64(base_img)
}
def interrogate(img):
url_interrogate = current_webui_url + 'sdapi/v1/interrogate'
payload = json.dumps({
# "model": "deepdanbooru",
"image": img
})
headers = {
'Content-Type': 'application/json'
}
response = requests.request("POST", url_interrogate, headers=headers, data=payload)
result = response.json()['caption']
return result
class ControlnetRequestImg2Img:
def __init__(self, prompt, net_prompt, mask_img):
self.url = f"{current_webui_url}sdapi/v1/img2img"
self.prompt = prompt
self.neg_prompt = net_prompt
self.body = None
self.mask = mask_img
def read_mask(self):
img = self.mask
retval, bytes = cv2.imencode('.png', img)
encoded_image = base64.b64encode(bytes).decode('utf-8')
return encoded_image
def build_body_v2(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7):
self.body = {
"prompt": self.prompt,
"negative_prompt": self.neg_prompt,
"sampler_name": "DPM++ 2M Karras",
"batch_size": 1,
"steps": int(os.environ.get("WEBUI_STEPS", "15")),
"width": dst_width,
"height": dst_height,
"cfg_scale": cfg_scale,
"seed": 123456789,
"mask_blur": 11,
"init_images": [
base_img
],
"inpaint_full_res": False,
"inpainting_fill": 1,
"inpainting_mask_invert": 0,
"mask": self.read_mask(),
# "refiner_checkpoint": "v1-5-pruned-emaonly.safetensors",
# "refiner_switch_at": 0.5,
"denoising_strength": denoising_strength,
"alwayson_scripts": {
# "controlnet": {
# "args": [
# {
# "enabled": True,
# "module": "openpose_full",
# "model": "openpose",
# "weight": 1.0,
# # "image": self.read_image(),
# "resize_mode": "Crop and Resize",
# "low_vram": False,
# "processor_res": 512,
# "guidance_start": 0.0,
# "guidance_end": 1.0,
# "control_mode": "Balanced",
# "pixel_perfect": True
# }
# ]
# }
# "controlnet": {
# "args": [
# {
# "enabled": True,
# "module": "openpose_full",
# "model": "openpose",
# "weight": 1.0,
# "resize_mode": 1,
# "lowvram": False,
# # "processor_res": 512,
# # "guidance_start": 0.0,
# # "guidance_end": 1.0,
# # "control_mode": 0,
# # "pixel_perfect": True
# },
# ]
# },
}
}
# 打印去除掉图像的body
self.print_body_without_images()
def print_body_without_images(self):
"""打印body内容，但不包含init_images和mask字段"""
import copy
import json
# 深拷贝body，避免修改原始数据
print_body = copy.deepcopy(self.body)
# 移除图像相关字段
if 'init_images' in print_body:
print_body['init_images'] = ['<base64_image_data>']
if 'mask' in print_body:
print_body['mask'] = '<base64_mask_data>'
print("Request body (without images):")
print(json.dumps(print_body, indent=2, ensure_ascii=False))
def build_body_hr(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7, refiner_switch_at=0.5):
self.body = {
"prompt": self.prompt,
"negative_prompt": self.neg_prompt,
"sampler_name": "DPM++ 2M Karras",
"batch_size": 1,
"steps": 20,
"width": dst_width,
"height": dst_height,
"cfg_scale": cfg_scale,
"seed": 123456789,
"mask_blur": 11,
"init_images": [
base_img
],
"inpaint_full_res": False,
"inpainting_fill": 1,
"inpainting_mask_invert": 0,
"mask": self.read_mask(),
"refiner_checkpoint": "v1-5-pruned-emaonly.safetensors",
"refiner_switch_at": refiner_switch_at,
"denoising_strength": denoising_strength,
"alwayson_scripts": {
}
}
# 打印去除掉图像的body
self.print_body_without_images()
def build_body_full_inpaint(self, dst_width, dst_height, cfg_scale, base_img):
self.body = {
"prompt": self.prompt,
"negative_prompt": self.neg_prompt,
"sampler_name": "DPM++ 2M Karras",
"batch_size": 1,
"steps": 20,
"width": dst_width,
"height": dst_height,
"cfg_scale": cfg_scale,
"seed": 123456789,
"mask_blur": 11,
"init_images": [
base_img
],
"inpaint_full_res": False,
"inpainting_fill": 1,
"inpainting_mask_invert": 0,
# "mask": self.read_mask(),
"refiner_checkpoint": "v1-5-pruned-emaonly.safetensors",
"refiner_switch_at": 0.5,
"denoising_strength": 0.7,
"alwayson_scripts": {
}
}
# 打印去除掉图像的body
self.print_body_without_images()
def build_body(self, dst_width, dst_height, cfg_scale, base_img):
self.body = {
"prompt": self.prompt,
"negative_prompt": self.neg_prompt,
"sampler_name": "DPM++ 2M Karras",
"batch_size": 1,
"steps": 30,
"width": dst_width,
"height": dst_height,
"cfg_scale": cfg_scale,
"seed": -1,
"mask_blur": 4,
"init_images": [
base_img
],
"inpaint_full_res": False,
"inpainting_fill": 1,
"inpainting_mask_invert": 1,
"mask": self.read_mask(),
"denoising_strength": 0.5,
"alwayson_scripts": {
"controlnet": {
"args": [
{
"enabled": True,
"module": "openpose_full",
"model": "openpose",
"weight": 1.0,
# "image": self.read_image(),
"resize_mode": "Crop and Resize",
"low_vram": False,
"processor_res": 512,
"guidance_start": 0.0,
"guidance_end": 1.0,
"control_mode": "Balanced",
"pixel_perfect": True
}
]
}
# "controlnet": {
# "args": [
# {
# "enabled": True,
# "module": "openpose_full",
# "model": "openpose",
# "weight": 1.0,
# "resize_mode": 1,
# "lowvram": False,
# # "processor_res": 512,
# # "guidance_start": 0.0,
# # "guidance_end": 1.0,
# # "control_mode": 0,
# # "pixel_perfect": True
# },
# ]
# },
}
}
# 打印去除掉图像的body
self.print_body_without_images()
def send_request(self):
response = requests.post(url=self.url, json=self.body)
return response.json()
def encode_image_to_base64(self, img):
retval, bytes = cv2.imencode('.png', img)
encoded_image = base64.b64encode(bytes).decode('utf-8')
return encoded_image
def get_high_train_img(img_path, in_gender):
img = cv2.imread(img_path)
out_path = img_path
# 如果图像长边尺寸小于1000，做超分
if max(img.shape[1], img.shape[0]) < 1000:
# cv2.imshow("img orig", img)
# 发送超分请求
img_super_res = WebUISupersuperResolution()
img_super_res.build_body(img)
print('sent hr request')
result = img_super_res.send_request()['image']
print('Super resolution done!')
image_array = np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8)
img = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
print(img.shape[1], img.shape[0])
# cv2.imshow("img super", img)
# cv2.waitKey(0)
# 做全图重绘到2000
img_scale = 2000 / (max(img.shape[1], img.shape[0]))
if img_scale < 1.0:
img = cv2.resize(img, (0, 0), fx=img_scale, fy=img_scale, interpolation=cv2.INTER_LANCZOS4)
print(img.shape[1], img.shape[0])
# cv2.imshow("orig", img)
# 存储超分后的图片
task_id = str(uuid4())
super_image_save_path = os.path.join(user_img_tmp_dir, task_id + "_super.png")
cv2.imwrite(super_image_save_path, img)
out_path = super_image_save_path
# todo can be del
# cv2.imshow("super", img)
# temp1 = os.path.join("/home/data/hair/data/test_tmp", str(uuid4()) + ".png")
# cv2.imwrite(temp1, img)
# 直接请求增强接口
# out = call_hair_enhance(super_image_save_path, "", task_id, in_gender)
# out_path = out["result"]
print("out_path:", out_path)
return out_path
def super_process(in_img=None, in_mask_img=None, in_gender=None, material_save_path=None, train_lora_material_path=None, task_id=None, hair_id=None):
img = in_img
mask_img = in_mask_img
# 如果图像长边尺寸大于1000，缩放到1000且不做超分
if max(img.shape[1], img.shape[0]) > 1024:
scale = 1024 / max(img.shape[1], img.shape[0])
img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LANCZOS4)
# 发送超分请求
img_super_res = WebUISupersuperResolution()
img_super_res.build_body(img)
print('sent hr request')
result = img_super_res.send_request()['image']
print('Super resolution done!')
image_array = np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8)
img = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
print(img.shape[1], img.shape[0])
img_scale = 1500 / (max(img.shape[1], img.shape[0]))
if img_scale < 1.0:
img = cv2.resize(img, (0, 0), fx=img_scale, fy=img_scale, interpolation=cv2.INTER_LANCZOS4)
print(img.shape[1], img.shape[0])
mask_img = cv2.resize(mask_img, dsize=(img.shape[1], img.shape[0]))
# cv2.imshow("super image", img)
super_image_save_path = os.path.join(material_save_path, "super.png")
cv2.imwrite(super_image_save_path, img)
# cv2.imshow("mask image", mask_img)
gen_img_mask_save_path = os.path.join(material_save_path, "gen_img_mask.png")
cv2.imwrite(gen_img_mask_save_path, mask_img)
# cv2.waitKey(0)
# 图像编码
retval, bytes = cv2.imencode('.png', img)
encoded_image = base64.b64encode(bytes).decode('utf-8')
prompt = interrogate(encoded_image)
prompt = ""
# print("prompt: ", prompt)
# prompt = '<lora:5b05d5eeee0188f436d7131c4f0ff52b:0.8>,easyphoto_face, easyphoto, 1person,face,suit'
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'
if in_gender == "boy":
neg_prompt = '(nsfw:1.5),(worst quality:2),(low quality:2),(normal quality:2),lowers,normal quality,(monochrome:1.2),(grayscale:1.2),skin spots,acnes,skin blemishes,age spot,ugly face,glans,fat,missing fingers,extra fingers,extra arms,extra legs,watermark,text,error,blurry,jpeg artifacts,cropped,bad anatomy,double navel,muscle,nsfw,nude,no nipple,hair ornaments,bad_pictures,badhandv4,easynegative'
control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img)
control_net.build_body(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image)
# 发送inpainting请求
print('sent inpainting request')
output = call_hair_infer(task_id, hair_id, train_lora_material_path, control_net.body)
print('Img2img done!')
# print(output)
result = output['images'][0]
res_img_encode = result.split(",", 1)[0]
# image_array = np.frombuffer(base64.b64decode(res_img_encode), np.uint8)
# img_res = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
# cv2.imwrite("/mnt/database2/online-server/hair-online/res_dir/90f21793-819f-46f6-91a9-d9a5259471101111.png", img_res)
# cv2.imshow("res_img:", img_res)
# cv2.waitKey(0)
return res_img_encode
def encode_numpy_to_base64(img):
retval, bytes = cv2.imencode('.png', img)
encoded_image = base64.b64encode(bytes).decode('utf-8')
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):
# 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'
if in_gender == "boy":
neg_prompt = '(nsfw:1.5),(worst quality:2),(low quality:2),(normal quality:2),lowers,normal quality,(monochrome:1.2),(grayscale:1.2),skin spots,acnes,skin blemishes,age spot,ugly face,glans,fat,missing fingers,extra fingers,extra arms,extra legs,watermark,text,error,blurry,jpeg artifacts,cropped,bad anatomy,double navel,muscle,nsfw,nude,no nipple,hair ornaments,bad_pictures,badhandv4,easynegative'
# 图像编码
retval, bytes = cv2.imencode('.png', img)
encoded_image = base64.b64encode(bytes).decode('utf-8')
prompt = tag
print("prompt:", prompt)
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)
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)
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)
# 发送inpainting请求
print('sent inpainting request')
start_inter = time.time()
output = call_hair_infer(task_id, hair_id, lora_material_path, control_net.body, is_hr, inference_port)
print('Img2img done!')
# print("--------------------- infer:", time.time() - start_inter)
# print(output)
result = output['images'][0]
image = Image.open(io.BytesIO(base64.b64decode(result.split(",", 1)[0])))
img_rgb = np.array(image)
img = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
return img
def webui_img2img_diy(img=None, mask_img=None, in_gender=None, task_id=None, tag="", inference_port="57860"):
# 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'
if in_gender == "boy":
neg_prompt = '(nsfw:1.5),(worst quality:2),(low quality:2),(normal quality:2),lowers,normal quality,(monochrome:1.2),(grayscale:1.2),skin spots,acnes,skin blemishes,age spot,ugly face,glans,fat,missing fingers,extra fingers,extra arms,extra legs,watermark,text,error,blurry,jpeg artifacts,cropped,bad anatomy,double navel,muscle,nsfw,nude,no nipple,hair ornaments,bad_pictures,badhandv4,easynegative'
# 图像编码
retval, bytes = cv2.imencode('.png', img)
encoded_image = base64.b64encode(bytes).decode('utf-8')
prompt = tag
print("prompt:", prompt)
denoising_strength = 0.3
print(f"diy strength:{denoising_strength}")
control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img)
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)
# 发送inpainting请求
print('sent inpainting request')
start_inter = time.time()
output = call_hair_infer_diy(task_id, control_net.body, inference_port)
print('Img2img done!')
# print("--------------------- infer:", time.time() - start_inter)
# print(output)
result = output['images'][0]
image = Image.open(io.BytesIO(base64.b64decode(result.split(",", 1)[0])))
img_rgb = np.array(image)
img = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
return img
def webui_super_res_img(img, ratio):
url = f"{current_webui_url}sdapi/v1/extra-single-image"
request_dict = {
"resize_mode": 0,
"show_extras_results": False,
"gfpgan_visibility": 0,
"codeformer_visibility": 1,
"codeformer_weight": 1,
"upscaling_resize": ratio,
"upscaler_1": "8x_NMKD-Superscale_150000_G",
"upscale_first": False,
"image": encode_numpy_to_base64(img)
}
response = requests.post(url=url, json=request_dict)
ret_json = response.json()
result = ret_json['image']
img = cv2.imdecode(np.frombuffer(base64.b64decode(result), np.uint8), cv2.IMREAD_COLOR)
return img
if __name__ == '__main__':
in_img = cv2.imread("/mnt/database2/online-server/hair-online/res_dir/90f21793-819f-46f6-91a9-d9a525947110.png")
in_mask_img = cv2.imread("/mnt/database2/online-server/hair-online/ref_hairstyle/5cc660db-0970-4467-9ccb-8f895fcdf5be/hull_mask.png")
# cv2.imshow("in_img", in_img)
# cv2.imshow("in_mask_img", in_mask_img)
# cv2.waitKey(0)
super_process(in_img=in_img, in_mask_img=in_mask_img)
Reference in New Issue View Git Blame Copy Permalink