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'] = [''] if 'mask' in print_body: print_body['mask'] = '' 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 = ',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)