xsl/change_hair_3090
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
change_hair_3090/hair_service_sd/gen_super_image.py
T
colomi 40a512a6c4 移除 majicmixRealistic_v7 依赖，改用 v1-5-pruned-emaonly 
- lora_train_service_1.py: 训练底模改为 v1-5-pruned-emaonly.safetensors
- gen_super_image.py: 注释掉 refiner_checkpoint（majicmix 不存在会导致推理报错）
- setup.sh: 检查项改为 v1-5-pruned-emaonly.safetensors
- README.md: 移除 majicmixRealistic_v7 下载步骤
2026-07-11 18:26:51 +08:00

517 lines
20 KiB
Python
Raw Permalink Blame History

import io
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": 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": "majicmixRealistic_v7.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):
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": "majicmixRealistic_v7.safetensors", # 底模不存在，禁用 refiner
# "refiner_switch_at": 0.5,
"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": "majicmixRealistic_v7.safetensors", # 底模不存在，禁用 refiner
# "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"):
# 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)
# 发送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