- 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 下载步骤
517 lines
20 KiB
Python
517 lines
20 KiB
Python
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)
|