|
|
import json
|
|
|
import os
|
|
|
import os.path
|
|
|
from urllib.parse import urlparse
|
|
|
|
|
|
from Util import ConfigUtil
|
|
|
from Util.AesUtil import *
|
|
|
from Util.ComfyUIUtil import *
|
|
|
from Util.CommonUtil import *
|
|
|
from Util.OssUtil import uploadOss
|
|
|
from Util.PngUtil import *
|
|
|
|
|
|
|
|
|
def upload_img_file(file_path, model_id):
|
|
|
png_file = os.path.basename(file_path)
|
|
|
# 清除Exif信息
|
|
|
clearExif(file_path, file_path) # 清除Exif信息
|
|
|
# 上传到OSS
|
|
|
key = 'Images/User/' + str(model_id) + '/' + png_file
|
|
|
uploadOss(key, file_path)
|
|
|
# # 生成下载链接
|
|
|
url = 'http://hzkc.oss-cn-beijing.aliyuncs.com/' + key
|
|
|
# 删除本地文件
|
|
|
os.remove(file_path)
|
|
|
printf("成功生成1张图片,地址：" + url)
|
|
|
return url
|
|
|
|
|
|
|
|
|
# 文生图服务调用
|
|
|
def webui_txt_2_img(model_id, json_data, input_image, target_folder):
|
|
|
# 提示词
|
|
|
prompt = json_data['prompt']
|
|
|
# 图片
|
|
|
args = json_data['alwayson_scripts']['ControlNet']['args']
|
|
|
|
|
|
#pretty_json = json.dumps(json_data, indent=4,ensure_ascii=False)
|
|
|
#print(pretty_json)
|
|
|
|
|
|
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(txt2img_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 webui_img_2_img(model_id, prompt_id, json_data, input_image, target_folder):
|
|
|
source_img = encode_image(input_image[0])
|
|
|
# 模型固定的参考图片
|
|
|
refer_source_img = r'./Image/' + str(model_id) + '/' + str(prompt_id) + '.jpg'
|
|
|
if not os.path.exists(refer_source_img):
|
|
|
refer_source_img = r'./Image/' + str(model_id) + '/' + str(prompt_id) + '.png'
|
|
|
# 都找不到，报错
|
|
|
if not os.path.exists(refer_source_img):
|
|
|
print("模型:" + str(model_id) + ",参考图:" + str(prompt_id) + "不存在，请检查！")
|
|
|
exit(0)
|
|
|
|
|
|
# 找到参考图
|
|
|
refer_img = encode_image(refer_source_img)
|
|
|
|
|
|
# 1.5与XL通用
|
|
|
# 在ControlNet的第一个图片中使用用户输入的图片，用此图片的人脸进行换脸操作
|
|
|
json_data["alwayson_scripts"]["ControlNet"]["args"][0]["input_image"] = source_img
|
|
|
# 图生图，风格参考图是输入
|
|
|
json_data["init_images"] = [refer_img]
|
|
|
|
|
|
# XL 专享
|
|
|
if len(json_data["alwayson_scripts"]["ControlNet"]["args"]) == 2:
|
|
|
json_data["alwayson_scripts"]["ControlNet"]["args"][1]["input_image"] = refer_img
|
|
|
|
|
|
# 开始生成图片
|
|
|
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) # 添加到结果集合中
|
|
|
|
|
|
|
|
|
# ComfyUI服务调用
|
|
|
def runComfyUI(model_id, json_data, input_image, target_folder):
|
|
|
# 生成一个唯一的客户端ID
|
|
|
client_id = str(uuid.uuid4())
|
|
|
|
|
|
# 上传图片
|
|
|
img = []
|
|
|
for x in input_image:
|
|
|
with open(x, "rb") as f:
|
|
|
y = upload_file(comfyui_address, f, "", True)
|
|
|
img.append(y)
|
|
|
# 如果存在多个姿势图和风格图，那么在配置文件中写死，比如pose=post_1.jpg style=style_1.jpg
|
|
|
# 同时，由于ComfyUI需要把这个文件上传到服务器，所以，需要在json文件中指定哪个文件需要上传
|
|
|
# 填充输入
|
|
|
json_data = fill_input(json_data, img)
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = []
|
|
|
# 遍历prompt_data中的所有节点
|
|
|
for key, value in json_data.items():
|
|
|
if 'inputs' in value:
|
|
|
if 'filename_prefix' in value['inputs']:
|
|
|
if value['inputs']['filename_prefix'] == 'ComfyUI':
|
|
|
myfilter.append(key)
|
|
|
# 生成
|
|
|
files = generate_clip(comfyui_address, json_data, client_id, target_folder, myfilter)
|
|
|
|
|
|
# 上传到云存储，回写数据接口
|
|
|
for file in files:
|
|
|
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):
|
|
|
os.system('./Tools/wget ' + image_url + ' -O ' + file_path)
|
|
|
|
|
|
|
|
|
# 根据输入的图片集合，填充到
|
|
|
def fill_input(prompt_data, file_array):
|
|
|
# 是不是不存在配置节
|
|
|
if "source_images" not in prompt_data:
|
|
|
print("配置文件不包含source_images配置节，程序无法继续！")
|
|
|
exit(0)
|
|
|
# 是不是个数不匹配
|
|
|
if len(prompt_data["source_images"]) != len(file_array):
|
|
|
print("输入的图片数量与预替换的图片数量不一致，程序无法继续！")
|
|
|
exit(0)
|
|
|
# 替换
|
|
|
for i in range(len(prompt_data["source_images"])):
|
|
|
x = prompt_data["source_images"][i]
|
|
|
# 替换配置节内容
|
|
|
array = x.split(",")
|
|
|
prompt_data[array[0]][array[1]][array[2]] = file_array[i]
|
|
|
|
|
|
# 删除多余的配置节
|
|
|
if "source_images" in prompt_data:
|
|
|
del prompt_data["source_images"]
|
|
|
|
|
|
# 处理style,pose 的参考图片
|
|
|
if 'style_pose_images' in prompt_data:
|
|
|
for x in prompt_data['style_pose_images']:
|
|
|
key = x['key']
|
|
|
value = x['value']
|
|
|
with open(value, "rb") as f:
|
|
|
y = upload_file(comfyui_address, f, "", True)
|
|
|
# 需要
|
|
|
array = key.split(",")
|
|
|
prompt_data[array[0]][array[1]][array[2]] = y
|
|
|
# 删除多余的配置节
|
|
|
if "style_pose_images" in prompt_data:
|
|
|
del prompt_data["style_pose_images"]
|
|
|
|
|
|
return prompt_data
|
|
|
|
|
|
|
|
|
# 获取任务
|
|
|
def get_task():
|
|
|
# 获取签名
|
|
|
en_data = getenData(machine_id)
|
|
|
# 请求的地址
|
|
|
url = web_url + '/QingLong/HuiYa/getTask'
|
|
|
# 要发送的数据，可以是字典形式
|
|
|
data = {
|
|
|
'enData': en_data
|
|
|
}
|
|
|
# 发送POST请求
|
|
|
response = requests.post(url, data=data)
|
|
|
|
|
|
# 检查请求是否成功
|
|
|
if response.status_code == 200:
|
|
|
success = response.json()["success"]
|
|
|
if success == 1:
|
|
|
data = response.json()["data"]
|
|
|
# 任务编号
|
|
|
task_id = data['task_id']
|
|
|
# 模板ID
|
|
|
model_id = data["model_id"]
|
|
|
# prompt_template 模板内容
|
|
|
prompt_id = data['prompt_id']
|
|
|
# 源图片
|
|
|
source_img_url = data["source_img_url"]
|
|
|
# 模板类型
|
|
|
model_type_id = data["model_type_id"]
|
|
|
|
|
|
# json文件
|
|
|
json_file = r'./JSON/' + str(model_id) + '_' + str(prompt_id) + '.json'
|
|
|
|
|
|
# 为了扩展支持多张图片，所以这里需要对,号进行拆解
|
|
|
input_image = []
|
|
|
for x in source_img_url.split(','):
|
|
|
# 解析URL
|
|
|
parsed_url = urlparse(x)
|
|
|
# 获取URL路径中的最后一个部分，通常就是文件名
|
|
|
filename = parsed_url.path.split('/')[-1]
|
|
|
# 下载这个图片
|
|
|
t_file = 'Temp/' + filename
|
|
|
input_image.append(t_file) # input_image扩展为数组，支持多张图片
|
|
|
# 这里可以加一个优化：如果本地存在这个文件就不用再次去下载了
|
|
|
if not os.path.exists(t_file):
|
|
|
download_image(x, t_file)
|
|
|
|
|
|
# 获取提示词
|
|
|
with open(json_file, 'r', encoding='utf-8') as file:
|
|
|
vjson_data = json.load(file)
|
|
|
|
|
|
# 利用三元表达式输出
|
|
|
printf("发现需要处理的" + ("WebUI" if model_type_id in [1, 2] else "ComfyUI") + "任务,model_type_id=" + str(
|
|
|
model_type_id) + ",model_id=" + str(
|
|
|
model_id) + ",prompt_id=" + str(prompt_id) + ",task_id=" + str(task_id))
|
|
|
|
|
|
# 创建目标目录
|
|
|
output_path = "Out/Images/User/" + str(model_id) + "/"
|
|
|
if not os.path.exists(output_path):
|
|
|
os.makedirs(output_path)
|
|
|
|
|
|
# SD_TXT2IMG
|
|
|
if model_type_id == 1:
|
|
|
webui_txt_2_img(model_id=model_id, json_data=vjson_data, input_image=input_image,
|
|
|
target_folder=output_path)
|
|
|
|
|
|
# SD_IMG2IMG
|
|
|
if model_type_id == 2:
|
|
|
webui_img_2_img(model_id=model_id, prompt_id=prompt_id, json_data=vjson_data,
|
|
|
input_image=input_image, target_folder=output_path)
|
|
|
# COMFY_UI
|
|
|
if model_type_id == 3:
|
|
|
runComfyUI(model_id=model_id, json_data=vjson_data, input_image=input_image,
|
|
|
target_folder=output_path)
|
|
|
|
|
|
return task_id
|
|
|
elif success == 0:
|
|
|
printf("现在没有待处理的任务,休息3秒后再试...")
|
|
|
return 0
|
|
|
else:
|
|
|
message = response.json()["message"]
|
|
|
printf('发生了异常:' + message)
|
|
|
return success
|
|
|
else:
|
|
|
printf('请求失败，状态码：' + response.status_code)
|
|
|
return -1
|
|
|
|
|
|
|
|
|
# 回复服务器，生图任务完成
|
|
|
def finish_task(task_id, target_img_urls):
|
|
|
# 获取签名
|
|
|
en_data = getenData(machine_id)
|
|
|
|
|
|
# 请求的地址
|
|
|
url = web_url + '/QingLong/HuiYa/finishTask'
|
|
|
|
|
|
# 要发送的数据，可以是字典形式
|
|
|
data = {
|
|
|
'task_id': task_id,
|
|
|
'enData': en_data,
|
|
|
'target_img_urls': target_img_urls
|
|
|
}
|
|
|
# 发送POST请求
|
|
|
response = requests.post(url, data=data)
|
|
|
# 检查请求是否成功
|
|
|
if response.status_code == 200:
|
|
|
printf("成功完成数据回写！")
|
|
|
else:
|
|
|
printf('请求失败，状态码：' + str(response.status_code))
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
# 禁止代理
|
|
|
set_http_proxy("")
|
|
|
|
|
|
# 配置文件
|
|
|
config = ConfigUtil.getConfig()
|
|
|
# 处理机编号
|
|
|
machine_id = config['system']['machine_id']
|
|
|
|
|
|
# webui 服务器地址
|
|
|
webui_address = config['webui']['webui_address']
|
|
|
# 文生图服务地址
|
|
|
txt2img_url = 'http://' + webui_address + config['webui']['txt2img_url']
|
|
|
# 图生图服务地址
|
|
|
img2img_url = 'http://' + webui_address + config['webui']['img2img_url']
|
|
|
|
|
|
# WEB服务器地址
|
|
|
web_url = config['webServer']['web_url']
|
|
|
# COMFYUI服务器地址
|
|
|
comfyui_address = config.get('comfyui', 'server_address')
|
|
|
|
|
|
while True:
|
|
|
try:
|
|
|
res = [] # 所有图片的地址路径
|
|
|
v_task_id = get_task()
|
|
|
# 回复上位生成完毕
|
|
|
if v_task_id > 0:
|
|
|
# 图片地址
|
|
|
v_target_img_urls = ''
|
|
|
for r in res:
|
|
|
v_target_img_urls = r + ',' + v_target_img_urls
|
|
|
v_target_img_urls = v_target_img_urls[0:len(v_target_img_urls) - 1]
|
|
|
# 生成完毕
|
|
|
finish_task(v_task_id, v_target_img_urls)
|
|
|
|
|
|
# 清理SD缓存
|
|
|
clear_webui_cache(webui_address)
|
|
|
get_webui_used(webui_address)
|
|
|
|
|
|
# 清理GPU缓存
|
|
|
clear_comfyui_cache(comfyui_address)
|
|
|
get_comfyui_used(comfyui_address)
|
|
|
|
|
|
elif v_task_id == 0:
|
|
|
# 本轮没有任务
|
|
|
time.sleep(3)
|
|
|
|
|
|
except Exception as err:
|
|
|
printf("发生异常，将休息3秒后重试..." + str(err))
|
|
|
time.sleep(3)
|
