'commit'

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: