Files
change_hair/project/hair_service_sd/gen_super_image.py
T
xsl 326b206d08 feat(swapHair): webui_steps参数可外部传入(接口2调试页可调)
- gen_super_image.py: webui_img2img + build_body_v2 接收可选 steps 参数
- run_copy_cost_colorb64.py: swapHair 接口接收 webui_steps,透传给 webui_img2img
- 未传时保持环境变量 WEBUI_STEPS 默认(15),向后兼容
2026-07-26 16:49:21 +08:00

518 lines
20 KiB
Python

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, steps=None):
self.body = {
"prompt": self.prompt,
"negative_prompt": self.neg_prompt,
"sampler_name": "DPM++ 2M Karras",
"batch_size": 1,
"steps": int(steps) if steps is not None else 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, webui_steps=None):
# 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, steps=webui_steps)
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)