From a5c3bc42be2b5b321f4006935b9da2f12f81cb07 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E6=B5=B7?= <10402852@qq.com>
Date: Tue, 2 Jul 2024 10:26:06 +0800
Subject: [PATCH] 'commit'

---
 BaiHu/doTask.py | 35 +++++++++++------------------------
 1 file changed, 11 insertions(+), 24 deletions(-)

diff --git a/BaiHu/doTask.py b/BaiHu/doTask.py
index 0334055b..6709b4c3 100644
--- a/BaiHu/doTask.py
+++ b/BaiHu/doTask.py
@@ -1,5 +1,6 @@
 import os
 import os.path
+import shutil
 from urllib.parse import urlparse
 
 from Util import ConfigUtil
@@ -10,7 +11,8 @@ from Util.OssUtil import uploadOss
 from Util.PngUtil import *
 
 
-def upload_img_file(png_file, file_path, model_id):
+def upload_img_file(file_path, model_id):
+    png_file = os.path.basename(file_path)
     # 清除Exif信息
     clearExif(file_path, file_path)  # 清除Exif信息
     # 上传到OSS
@@ -46,12 +48,10 @@ def webui_txt_2_img(model_id, json_data, input_image, target_folder):
     printf("开始生成图片...")
     # 调用生图服务
     response = submit_post(txt2img_url, json_data)
-    # 创建文件名
-    png_file = str(uuid.uuid4()) + '.png'
     # 创建文件的完整路径
-    file_path = os.path.join(target_folder, png_file)
+    file_path = os.path.join(target_folder, str(uuid.uuid4()) + '.png')
     save_encoded_image(response.json()['images'][0], file_path)
-    url = upload_img_file(png_file, file_path, model_id)
+    url = upload_img_file(file_path, model_id)
     # 添加到结果集合中
     res.append(url)
 
@@ -86,15 +86,14 @@ def webui_img_2_img(model_id, prompt_id, json_data, input_image, target_folder):
 
     # 调用生图服务
     response = submit_post(img2img_url, json_data)
-    # 创建文件名
-    png_file = str(uuid.uuid4()) + '.png'
     # 创建文件的完整路径
-    file_path = os.path.join(target_folder, png_file)
+    file_path = os.path.join(target_folder, str(uuid.uuid4()) + '.png')
     save_encoded_image(response.json()['images'][0], file_path)
-    url = upload_img_file(png_file, file_path, model_id)
+    url = upload_img_file(file_path, model_id)
     res.append(url)  # 添加到结果集合中
 
 
+# ComfyUI服务调用
 def runComfyUI(model_id, json_data, input_image, target_folder):
     # 生成一个唯一的客户端ID
     client_id = str(uuid.uuid4())
@@ -116,24 +115,13 @@ def runComfyUI(model_id, json_data, input_image, target_folder):
 
     # 上传到云存储，回写数据接口
     for file in files:
-        png_file = file + '.png'
-        file_path = './Out/Images/User/' + str(model_id) + "/" + png_file
-        url = upload_img_file(png_file, file_path, model_id)
+        file_path = './Out/Images/User/' + str(model_id) + "/" + file + '.png'
+        url = upload_img_file(file_path, model_id)
         res.append(url)
 
 
 def download_image(image_url, file_path):
-    # 发送GET请求
-    response = requests.get(image_url)
-    # 检查请求是否成功
-    if response.status_code == 200:
-        # 打开一个文件用于写入
-        with open(file_path, 'wb') as f:
-            # 将响应的内容写入文件中
-            f.write(response.content)
-    else:
-        printf('图片下载失败，状态码：' + str(response.status_code))
-
+    os.system('./Tools/wget ' + image_url + ' -O ' + file_path)
 
 def getMyFilter(prompt_data):
     myfilter = []
@@ -184,7 +172,6 @@ def fill_input(prompt_data, file_array):
     return prompt_data
 
 
-
 # 获取任务
 def get_task():
     # 获取签名