diff --git a/BaiHu/doTask.py b/BaiHu/doTask.py index aea54c42..df621d38 100644 --- a/BaiHu/doTask.py +++ b/BaiHu/doTask.py @@ -129,6 +129,45 @@ def runComfyUI(model_id, json_data, input_image, target_folder): res.append(url) +# 风格转绘 +def webui_convert(model_id, json_data, input_image, target_folder): + # 1、需要反推提示词 + prompt = get_wd_14(wd_url, input_image[0]) + + # 2、需要获取原始图片的大小 + img = Image.open(input_image[0]) + width = img.size[0], height = img.size[1] + img.close() + + # 然后需要干点啥? + # TODO + + # 图片 + args = json_data['alwayson_scripts']['ControlNet']['args'] + pose_image = json_data['pose_image'] + args[0]['input_image'] = encode_image(input_image[0]) + # 如果有特殊的姿态要求,那就用模板指定的姿态图 + if len(pose_image) > 0: + if len(args) > 1: + args[1]['input_image'] = encode_image(pose_image) + else: # 如果没有特殊的姿态要求,那么就统一用一张图片 + if len(args) > 1: + args[1]['input_image'] = encode_image(input_image[0]) + + # 提示词 + json_data['prompt'] = prompt + # 开始生成图片 + printf("开始生成图片...") + # 调用生图服务 + response = submit_post(img2img_url, json_data) + # 创建文件的完整路径 + file_path = os.path.join(target_folder, str(uuid.uuid4()) + '.png') + save_encoded_image(response.json()['images'][0], file_path) + url = upload_img_file(file_path, model_id) + # 添加到结果集合中 + res.append(url) + + def download_image(image_url, file_path): # 发送GET请求 response = requests.get(image_url) @@ -257,6 +296,11 @@ def get_task(): runComfyUI(model_id=model_id, json_data=vjson_data, input_image=input_image, target_folder=output_path) + # 风格转绘 + if model_type_id == 4: + webui_convert(model_id=model_id, prompt_id=prompt_id, json_data=vjson_data, + input_image=input_image, target_folder=output_path) + # 清理服务器显存 release_sd(webui_address, comfyui_address) @@ -318,8 +362,8 @@ if __name__ == '__main__': comfyui_address = config.get('comfyui', 'server_address') # 反推接口地址 wd_url = 'http://' + webui_address + '/tagger/v1/interrogate' - #image_path = r'D:\KeCheng\BaiHu\Backup\mote2.png' - #print(get_wd_14(wd_url, image_path)) + # image_path = r'D:\KeCheng\BaiHu\Backup\mote2.png' + # print(get_wd_14(wd_url, image_path)) while True: try: