|
|
import os
|
|
|
import random
|
|
|
import threading
|
|
|
from urllib.parse import urlparse
|
|
|
from tqdm import tqdm
|
|
|
from Util import ConfigUtil
|
|
|
from Util.AesUtil import *
|
|
|
from Util.SDUtil import *
|
|
|
from Util.CommonUtil import *
|
|
|
from Util.Logger import Logger
|
|
|
from Util.OssUtil import uploadOss
|
|
|
from Util.PngUtil import *
|
|
|
|
|
|
# 定义三种模式
|
|
|
MS_RELAX = 'relax'
|
|
|
MS_FAST = 'fast'
|
|
|
MS_TURBO = 'turbo'
|
|
|
|
|
|
|
|
|
# 获取指定任务号的完成情况
|
|
|
def get_task_info(api_token, result_id):
|
|
|
# 获取任务结果
|
|
|
url = "https://goapi.gptnb.me/mj/task/" + str(result_id) + "/fetch"
|
|
|
payload = {}
|
|
|
headers = {
|
|
|
'Authorization': 'Bearer ' + api_token
|
|
|
}
|
|
|
# 国内网址，不使用代理
|
|
|
res = requests.request("GET", url, headers=headers, data=payload)
|
|
|
# 获取任务结果
|
|
|
return json.loads(res.text)
|
|
|
|
|
|
|
|
|
# 下载Midjourney生成的图片，外网地址，需要用VPN挂代理下载
|
|
|
def down_mj_img(image_url, file_path):
|
|
|
# 拆分目录和文件名
|
|
|
directory = os.path.dirname(file_path)
|
|
|
|
|
|
# 检查目录是否存在
|
|
|
if not os.path.exists(directory):
|
|
|
# 如果目录不存在，则创建它
|
|
|
os.makedirs(directory)
|
|
|
|
|
|
while True:
|
|
|
try:
|
|
|
# 使用代理进行下载
|
|
|
response = requests.get(image_url, proxies=proxies)
|
|
|
# 获取文件总大小
|
|
|
total_size_in_bytes = int(response.headers.get('content-length', 0))
|
|
|
log.info("开始下载文件，文件大小:" + str(round(total_size_in_bytes / 1024 / 1024, 2)) + "MB")
|
|
|
# 创建一个进度条
|
|
|
progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
|
|
|
|
|
|
# 定义一个函数来处理数据块
|
|
|
def callback_chunk(chunk):
|
|
|
progress_bar.update(len(chunk))
|
|
|
|
|
|
# 打开一个本地文件用于写入
|
|
|
with open(file_path, 'wb') as file:
|
|
|
# 迭代响应内容
|
|
|
for chunk in response.iter_content(chunk_size=1024):
|
|
|
file.write(chunk)
|
|
|
callback_chunk(chunk) # 更新进度条
|
|
|
# 关闭进度条
|
|
|
progress_bar.close()
|
|
|
# log.info("恭喜，文件下载完成!")
|
|
|
break
|
|
|
except Exception as err:
|
|
|
log.error("发现网络异常，将休息2秒后再试,详细错误信息：" + str(err))
|
|
|
time.sleep(2)
|
|
|
|
|
|
|
|
|
# 根据任务编号下载生成的图片
|
|
|
def try_down_mj_img(api_token, task_id, file_path):
|
|
|
jo = get_task_info(api_token=api_token, result_id=task_id)
|
|
|
# 图片生成情况
|
|
|
if jo['progress'] == '100%' and len(jo['imageUrl']) > 0:
|
|
|
image_url = jo['imageUrl']
|
|
|
down_mj_img(image_url, file_path)
|
|
|
return 1, '生成完毕！'
|
|
|
elif jo['progress'] == '100%' and len(jo['imageUrl']) == 0:
|
|
|
return -1, "生成失败，需要重新生成" + jo['failReason']
|
|
|
else:
|
|
|
return 0, "没生成完！progress = " + jo['progress']
|
|
|
|
|
|
|
|
|
# 创建图片
|
|
|
def imagine(api_server, api_token, prompt, type_code=MS_RELAX, input_img=None):
|
|
|
# 地址
|
|
|
if type_code == MS_RELAX:
|
|
|
url = api_server + "/mj-relax/mj/submit/imagine" # relax
|
|
|
elif type_code == MS_FAST:
|
|
|
url = api_server + "/mj/submit/imagine" # fast
|
|
|
elif type_code == MS_TURBO:
|
|
|
url = api_server + "/mj-turbo/mj/submit/imagine" # turbo
|
|
|
|
|
|
data = {
|
|
|
"botType": "MID_JOURNEY",
|
|
|
"prompt": prompt,
|
|
|
"base64Array": []
|
|
|
}
|
|
|
# 垫图转为base64
|
|
|
if input_img is not None:
|
|
|
for e in input_img.split(','):
|
|
|
img_base64_str = get_image_base64_str(e)
|
|
|
data["base64Array"].append(img_base64_str)
|
|
|
|
|
|
headers = {
|
|
|
'content-type': 'application/json', 'Authorization': 'Bearer ' + api_token
|
|
|
}
|
|
|
# 提交的是国内的站点，不需要代理
|
|
|
res = requests.request("POST", url, headers=headers, data=json.dumps(data))
|
|
|
return json.loads(res.text)
|
|
|
|
|
|
|
|
|
# 检查指定的风格与形态图是不是存在
|
|
|
def checkStylePoseIdExists(user_id, model_id, style_id, pose_id):
|
|
|
# 需要调用远程WEB SERVER的接口，获取是不是此用户+此模型+此风格+此姿势的记录
|
|
|
enData = getenData(machine_id)
|
|
|
# 请求的地址
|
|
|
url = web_url + '/QingLong/HuiYa/checkStylePoseIdExists'
|
|
|
# 要发送的数据，可以是字典形式
|
|
|
data = {
|
|
|
'enData': enData,
|
|
|
'user_id': user_id,
|
|
|
'model_id': model_id,
|
|
|
'style_id': style_id,
|
|
|
'pose_id': pose_id
|
|
|
}
|
|
|
# 发送GET请求
|
|
|
response = requests.get(url, data=data)
|
|
|
|
|
|
# 检查请求是否成功
|
|
|
if response.status_code == 200:
|
|
|
exists = json.loads(response.text)['success']
|
|
|
else:
|
|
|
log.error('获取用户模型风格与姿势是否存在失败，状态码：', response.status_code)
|
|
|
return False
|
|
|
|
|
|
return exists
|
|
|
|
|
|
|
|
|
# 回写数据库样式与姿势键值对
|
|
|
def record_style_pose_id(user_id, model_id, style_id, pose_id):
|
|
|
if not checkStylePoseIdExists(user_id, model_id, style_id, pose_id):
|
|
|
# 需要调用远程WEB SERVER的接口，获取是不是此用户+此模型+此风格+此姿势的记录
|
|
|
enData = getenData(machine_id)
|
|
|
# 请求的地址
|
|
|
url = web_url + '/QingLong/HuiYa/record_style_pose_id'
|
|
|
# 要发送的数据，可以是字典形式
|
|
|
data = {
|
|
|
'enData': enData,
|
|
|
'user_id': user_id,
|
|
|
'model_id': model_id,
|
|
|
'style_id': style_id,
|
|
|
'pose_id': pose_id
|
|
|
}
|
|
|
# 发送GET请求
|
|
|
response = requests.post(url, data=data)
|
|
|
|
|
|
# 检查请求是否成功
|
|
|
if response.status_code == 200:
|
|
|
log.error('设置用户模型风格与姿势时成功!')
|
|
|
else:
|
|
|
log.error('设置用户模型风格与姿势时失败，状态码：', response.status_code)
|
|
|
|
|
|
|
|
|
# 获取Model对象
|
|
|
def getModel(model_id):
|
|
|
file_name = r'JSON/' + str(model_id) + '.json'
|
|
|
# 打开JSON文件
|
|
|
if not os.path.exists(file_name):
|
|
|
log.error("文件：" + file_name + "不存在，跳过！")
|
|
|
return
|
|
|
with open(file_name, 'r', encoding='utf-8') as file:
|
|
|
data = json.load(file)
|
|
|
return data
|
|
|
|
|
|
|
|
|
# 获取模型名称
|
|
|
def getModelName():
|
|
|
return data['override_settings']['sd_model_checkpoint']
|
|
|
|
|
|
|
|
|
# 整体的生成任务
|
|
|
def runWebUI(model_id, task_type_code, input_image, total_tasks, user_id, prompt):
|
|
|
global data
|
|
|
data = getModel(model_id) # data是全局变量
|
|
|
# 生成多少张
|
|
|
data['batch_size'] = total_tasks
|
|
|
# 图片
|
|
|
args = data['alwayson_scripts']['ControlNet']['args']
|
|
|
pose_image = 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])
|
|
|
|
|
|
# 宽度与高度
|
|
|
width = data['width']
|
|
|
height = data['height']
|
|
|
# 提示词
|
|
|
data['prompt'] = prompt
|
|
|
|
|
|
# 开始生成图片
|
|
|
log.info("开始生成图片...")
|
|
|
# 设置目标文件夹
|
|
|
target_folder = "Out"
|
|
|
|
|
|
# 创建以当前日期命名的目录
|
|
|
folder = os.path.join(target_folder, 'Images/' + task_type_code + "/" + str(model_id) + '/')
|
|
|
if not os.path.exists(folder):
|
|
|
os.makedirs(folder)
|
|
|
# 调用生图服务
|
|
|
response = submit_post(txt2img_url, data)
|
|
|
# 保存为图片
|
|
|
for i in range(0, total_tasks):
|
|
|
# 创建文件名
|
|
|
png_file = str(uuid.uuid4()) + '.png'
|
|
|
# 创建文件的完整路径
|
|
|
file_path = os.path.join(folder, png_file)
|
|
|
save_encoded_image(response.json()['images'][i], file_path)
|
|
|
|
|
|
# 清除Exif信息
|
|
|
clearExif(file_path, file_path) # 清除Exif信息
|
|
|
# 上传到OSS
|
|
|
key = 'Images/' + task_type_code + "/" + str(model_id) + '/' + png_file
|
|
|
uploadOss(key, file_path)
|
|
|
# # 生成下载链接
|
|
|
url = 'http://hzkc.oss-cn-beijing.aliyuncs.com/' + key
|
|
|
# 生成预览链接
|
|
|
# previewSuffix = '?x-oss-process=image/resize,w_150,limit_0'
|
|
|
# 删除本地文件
|
|
|
os.remove(file_path)
|
|
|
log.info("成功生成第" + str(i + 1) + "张图片,宽度：" + str(width) + ",高度:" + str(height) + " 地址：" + url)
|
|
|
# 添加到结果集合中
|
|
|
res.append(url)
|
|
|
log.info("本轮图片生成成功, 共完成" + str(total_tasks) + "张图片")
|
|
|
|
|
|
|
|
|
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:
|
|
|
log.error('图片下载失败，状态码：', response.status_code)
|
|
|
|
|
|
|
|
|
def generate_random_number(length):
|
|
|
return ''.join(random.choices('0123456789', k=length))
|
|
|
|
|
|
|
|
|
# 指定目录下有多少个指定前缀的文件
|
|
|
def getPrefixFileCount(directory_path, prefix):
|
|
|
# 使用os.listdir()获取目录下的所有文件和文件夹
|
|
|
entries = os.listdir(directory_path)
|
|
|
# 使用os.path.isfile()检查每个条目是否是文件
|
|
|
files = [entry for entry in entries if os.path.isfile(os.path.join(directory_path, entry))]
|
|
|
cnt = 0
|
|
|
for file in files:
|
|
|
if file.startswith(prefix):
|
|
|
cnt = cnt + 1
|
|
|
return cnt
|
|
|
|
|
|
|
|
|
def getMyFilter(prompt_data):
|
|
|
myfilter = []
|
|
|
# 遍历prompt_data中的所有节点
|
|
|
for key, value in prompt_data.items():
|
|
|
if 'inputs' in value:
|
|
|
if 'filename_prefix' in value['inputs']:
|
|
|
if value['inputs']['filename_prefix'] == 'ComfyUI':
|
|
|
myfilter.append(key)
|
|
|
return myfilter
|
|
|
|
|
|
|
|
|
def C16(prompt_data, img, model_id, user_id, prompt):
|
|
|
# 图生图
|
|
|
prompt_data["22"]["inputs"]["image"] = img[0]
|
|
|
# 替换关键词
|
|
|
prompt_data["331"]["inputs"]["text"] = prompt
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
def C17(prompt_data, img, model_id, user_id, prompt):
|
|
|
# 上传姿势
|
|
|
# random_number = random.randint(1, 9) # 9张姿势图
|
|
|
# input_image_pose = './Image/C17/' + str(random_number) + '.png'
|
|
|
input_image_pose = './Image/C17/1.jpg'
|
|
|
with open(input_image_pose, "rb") as f:
|
|
|
imgPose = upload_file(server_address, f, "", True)
|
|
|
|
|
|
# 图生图
|
|
|
prompt_data["331"]["inputs"]["image"] = imgPose
|
|
|
prompt_data["432"]["inputs"]["image"] = img[0]
|
|
|
|
|
|
# 指定数字长度，例如16位
|
|
|
number_length = 16
|
|
|
random_number = generate_random_number(number_length)
|
|
|
prompt_data["354"]["inputs"]["seed"] = random_number
|
|
|
random_number = generate_random_number(number_length)
|
|
|
prompt_data["425"]["inputs"]["seed"] = random_number
|
|
|
|
|
|
# 替换关键词
|
|
|
prompt_data["331"]["inputs"]["text"] = prompt
|
|
|
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
def C19(prompt_data, img, model_id, user_id, prompt):
|
|
|
#
|
|
|
prompt_data["12"]["inputs"]["image"] = img[0]
|
|
|
|
|
|
if len(img) > 1:
|
|
|
prompt_data["27"]["inputs"]["image"] = img[1]
|
|
|
else:
|
|
|
prompt_data["27"]["inputs"]["image"] = img[0]
|
|
|
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
def C20(prompt_data, img, model_id, user_id, prompt):
|
|
|
#
|
|
|
prompt_data["12"]["inputs"]["image"] = img[0]
|
|
|
if len(img) > 1:
|
|
|
prompt_data["27"]["inputs"]["image"] = img[1]
|
|
|
else:
|
|
|
input_image_style = './Image/C20/meixi.png' # 只上传一张就与梅西握手
|
|
|
with open(input_image_style, "rb") as f:
|
|
|
image_style = upload_file(server_address, f, "", True)
|
|
|
prompt_data["27"]["inputs"]["image"] = image_style
|
|
|
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
def C22(prompt_data, img, model_id, user_id, prompt):
|
|
|
####------------------------------------------------------------####
|
|
|
# 记录下用户的风格图+姿势图的组合，下次生成时，尽时避开已经输出过的组合，让用户更好的玩耍
|
|
|
# 获取目录下有多少个pose,多少个style
|
|
|
directory_path = './Image/C' + str(model_id) + '/'
|
|
|
pose_count = getPrefixFileCount(directory_path, 'pose_')
|
|
|
style_count = getPrefixFileCount(directory_path, 'style_')
|
|
|
style_id = random.randint(1, style_count) # 查找未使用过的姿态图与姿势图，如果找不到就随机用一个
|
|
|
pose_id = random.randint(1, max(pose_count, 1))
|
|
|
for i in range(1, style_count + 1):
|
|
|
for j in range(1, pose_count + 1):
|
|
|
if not checkStylePoseIdExists(user_id, model_id, i, j):
|
|
|
style_id = i
|
|
|
pose_id = j
|
|
|
# 回写数据库表，标识这些组合已经用过
|
|
|
record_style_pose_id(user_id, model_id, style_id, pose_id)
|
|
|
####------------------------------------------------------------####
|
|
|
# 风格图
|
|
|
input_image_style = './Image/C' + str(model_id) + '/style_' + str(style_id) + '.png'
|
|
|
with open(input_image_style, "rb") as f:
|
|
|
image_style = upload_file(server_address, f, "", True)
|
|
|
prompt_data["118"]["inputs"]["image"] = image_style
|
|
|
|
|
|
# 样式图
|
|
|
input_image_pose = './Image/C' + str(model_id) + '/pose_' + str(pose_id) + '.png'
|
|
|
with open(input_image_pose, "rb") as f:
|
|
|
image_pose = upload_file(server_address, f, "", True)
|
|
|
prompt_data["98"]["inputs"]["image"] = image_pose
|
|
|
|
|
|
# 原图
|
|
|
prompt_data["49"]["inputs"]["image"] = img[0]
|
|
|
|
|
|
# 提示词
|
|
|
prompt_data["131"]["inputs"]["String"] = prompt
|
|
|
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
def C23(prompt_data, img, model_id, user_id, prompt):
|
|
|
# 原图
|
|
|
prompt_data["52"]["inputs"]["image"] = img[0]
|
|
|
|
|
|
# prompt_data["331"]["inputs"]["text"] = prompt
|
|
|
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
def C24(prompt_data, img, model_id, user_id, prompt):
|
|
|
# 原图
|
|
|
prompt_data["12"]["inputs"]["image"] = img[0]
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
def C25(prompt_data, img, model_id, user_id, prompt):
|
|
|
# 原图
|
|
|
prompt_data["111"]["inputs"]["image"] = img[0]
|
|
|
|
|
|
####------------------------------------------------------------####
|
|
|
# 记录下用户的风格图+姿势图的组合，下次生成时，尽时避开已经输出过的组合，让用户更好的玩耍
|
|
|
# 获取目录下有多少个pose,多少个style
|
|
|
directory_path = './Image/C' + str(model_id) + '/'
|
|
|
pose_count = getPrefixFileCount(directory_path, 'pose_')
|
|
|
style_count = getPrefixFileCount(directory_path, 'style_')
|
|
|
style_id = random.randint(1, style_count) # 查找未使用过的姿态图与姿势图，如果找不到就随机用一个
|
|
|
pose_id = random.randint(1, max(pose_count, 1))
|
|
|
for i in range(1, style_count + 1):
|
|
|
for j in range(1, pose_count + 1):
|
|
|
flag = checkStylePoseIdExists(user_id, model_id, i, j)
|
|
|
if not flag:
|
|
|
style_id = i
|
|
|
pose_id = j
|
|
|
# 回写数据库表，标识这些组合已经用过
|
|
|
record_style_pose_id(user_id, model_id, style_id, pose_id)
|
|
|
####------------------------------------------------------------####
|
|
|
|
|
|
with open('./Image/C' + str(model_id) + '/style_' + str(style_id) + ".png", "rb") as f:
|
|
|
y = upload_file(server_address, f, "", True)
|
|
|
|
|
|
# 风格(随机)
|
|
|
prompt_data["113"]["inputs"]["image"] = y
|
|
|
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
# 获取C26中有哪些样式
|
|
|
def C26_get_sdxl_styles():
|
|
|
files = ['./JSON/C26/sdxl_styles_sai.json', './JSON/C26/sdxl_styles_twri.json']
|
|
|
SDXL_STYLES = []
|
|
|
for json_file in files:
|
|
|
with open(json_file, 'r', encoding="utf-8") as fi:
|
|
|
data = json.load(fi)
|
|
|
for x in data:
|
|
|
SDXL_STYLES.append(x['name'])
|
|
|
return SDXL_STYLES
|
|
|
|
|
|
|
|
|
def C26(prompt_data, img, model_id, user_id, prompt):
|
|
|
# 原图
|
|
|
prompt_data["13"]["inputs"]["image"] = img[0]
|
|
|
|
|
|
# 风格图片
|
|
|
SDXL_STYLES = C26_get_sdxl_styles()
|
|
|
style_id = random.randint(0, len(SDXL_STYLES) - 1)
|
|
|
pose_id = 1
|
|
|
for i in range(0, len(SDXL_STYLES)):
|
|
|
if not checkStylePoseIdExists(user_id, model_id, i, 1):
|
|
|
style_id = i
|
|
|
|
|
|
# 回写数据库表，标识这些组合已经用过
|
|
|
record_style_pose_id(user_id, model_id, style_id, pose_id)
|
|
|
prompt_data["79"]["inputs"]["style"] = SDXL_STYLES[style_id]
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
def C27(prompt_data, img, model_id, user_id, prompt):
|
|
|
# 原图
|
|
|
if len(img) == 1:
|
|
|
with open('../Image/Backup/huanghai_3.png', "rb") as f:
|
|
|
x = upload_file(server_address, f, "", True)
|
|
|
prompt_data["82"]["inputs"]["image"] = x
|
|
|
prompt_data["78"]["inputs"]["image"] = img[0] # 女
|
|
|
else:
|
|
|
prompt_data["13"]["inputs"]["image"] = img[0] # 男
|
|
|
prompt_data["78"]["inputs"]["image"] = img[1] # 女
|
|
|
|
|
|
# 风格图片
|
|
|
style_id = random.randint(1, 2)
|
|
|
pose_id = random.randint(1, 3)
|
|
|
|
|
|
for i in range(1, 3):
|
|
|
for j in range(1, 4):
|
|
|
if not checkStylePoseIdExists(user_id, model_id, i, j):
|
|
|
style_id = i
|
|
|
pose_id = j
|
|
|
|
|
|
# 回写数据库表，标识这些组合已经用过
|
|
|
record_style_pose_id(user_id, model_id, style_id, pose_id)
|
|
|
|
|
|
# 姿势图有3个 pose_id
|
|
|
|
|
|
pose_img_file = './Image/C27/pose_' + str(pose_id) + ".jpg"
|
|
|
# 上传姿势图
|
|
|
with open(pose_img_file, "rb") as f:
|
|
|
y = upload_file(server_address, f, "", True)
|
|
|
prompt_data["82"]["inputs"]["image"] = y
|
|
|
|
|
|
# 模型有2个 style_id
|
|
|
if style_id == 1:
|
|
|
prompt_data["4"]["inputs"]["ckpt_name"] = "samaritan3dCartoon_v40SDXL.safetensors"
|
|
|
else:
|
|
|
prompt_data["4"]["inputs"]["ckpt_name"] = "turbovisionxlSuperFastXLBasedOnNew_tvxlV431Bakedvae.safetensors"
|
|
|
|
|
|
# 过滤器,哪些返回节点中获取到的图片是有用的
|
|
|
myfilter = getMyFilter(prompt_data)
|
|
|
return prompt_data, myfilter
|
|
|
|
|
|
|
|
|
def runComfyUI(json_file, model_id, task_type_code, input_image, user_id, prompt):
|
|
|
# 创建目标目录
|
|
|
output_path = "Out/Images/" + task_type_code + "/" + str(model_id) + "/"
|
|
|
if not os.path.exists(output_path):
|
|
|
os.makedirs(output_path)
|
|
|
|
|
|
# 生成一个唯一的客户端ID
|
|
|
client_id = str(uuid.uuid4())
|
|
|
|
|
|
# 上传图片
|
|
|
img = []
|
|
|
for x in input_image:
|
|
|
with open(x, "rb") as f:
|
|
|
y = upload_file(server_address, f, "", True)
|
|
|
img.append(y)
|
|
|
|
|
|
with open(json_file, 'r', encoding="utf-8") as fi:
|
|
|
prompt_data = json.load(fi)
|
|
|
|
|
|
# 调用不同的实现函数
|
|
|
number = json_file[json_file.rfind('/') + 1:-5] # 从最后一个'/'之后开始，到倒数第五个字符结束
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '16':
|
|
|
prompt_data, myfilter = C16(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '17':
|
|
|
prompt_data, myfilter = C17(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '19':
|
|
|
prompt_data, myfilter = C19(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '20':
|
|
|
prompt_data, myfilter = C20(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '22':
|
|
|
prompt_data, myfilter = C22(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '23':
|
|
|
prompt_data, myfilter = C23(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '24':
|
|
|
prompt_data, myfilter = C24(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '25':
|
|
|
prompt_data, myfilter = C25(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '26':
|
|
|
prompt_data, myfilter = C26(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 分支判断
|
|
|
if number == '27':
|
|
|
prompt_data, myfilter = C27(prompt_data, img, model_id, user_id, prompt)
|
|
|
|
|
|
# 生成
|
|
|
files = generate_clip(server_address, prompt_data, client_id, output_path, myfilter)
|
|
|
|
|
|
# 上传到云存储，回写数据接口
|
|
|
for file in files:
|
|
|
png_file = file + '.png'
|
|
|
file_path = './Out/Images/' + task_type_code + "/" + str(model_id) + "/" + png_file
|
|
|
# 清除Exif信息
|
|
|
clearExif(file_path, file_path) # 清除Exif信息
|
|
|
# 上传到OSS
|
|
|
key = 'Images/' + task_type_code + "/" + str(model_id) + '/' + png_file
|
|
|
uploadOss(key, file_path)
|
|
|
# # 生成下载链接
|
|
|
url = 'http://hzkc.oss-cn-beijing.aliyuncs.com/' + key
|
|
|
res.append(url)
|
|
|
|
|
|
|
|
|
# 获取任务
|
|
|
def get_task():
|
|
|
# 获取签名
|
|
|
enData = getenData(machine_id)
|
|
|
# 请求的地址
|
|
|
url = web_url + '/QingLong/HuiYa/getTask'
|
|
|
# 要发送的数据，可以是字典形式
|
|
|
data = {
|
|
|
'enData': enData
|
|
|
}
|
|
|
# 发送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']
|
|
|
# 任务数量
|
|
|
total_count = data["count"]
|
|
|
# 用户ID
|
|
|
user_id = data['user_id']
|
|
|
|
|
|
# prompt_template 模板内容
|
|
|
prompt_template = data['prompt_template']
|
|
|
|
|
|
# prompt_en
|
|
|
prompt_en = data['prompt_en']
|
|
|
if prompt_en is None:
|
|
|
prompt_en = ''
|
|
|
# 提示词
|
|
|
prompt = prompt_template.replace('{sample}', prompt_en)
|
|
|
|
|
|
# 这个返回应该是一个 类似这样的地址： source_img_url = 'http://hzkc.oss-cn-beijing.aliyuncs.com/Images/Upload/wife.jpg',为了扩展支持多张图片上传，所以，多张图片之间用逗号分隔
|
|
|
source_img_url = data["source_img_url"]
|
|
|
# 模板ID
|
|
|
model_id = data["model_id"]
|
|
|
# 任务类型代码
|
|
|
task_type_code = data["task_type_code"]
|
|
|
# 模板类型
|
|
|
model_type_id = data["model_type_id"]
|
|
|
# 为了扩展支持多张图片，所以这里需要对,号进行拆解
|
|
|
input_image = []
|
|
|
for x in source_img_url.split(','):
|
|
|
# 解析URL
|
|
|
parsed_url = urlparse(x)
|
|
|
# 获取URL路径中的最后一个部分，通常就是文件名
|
|
|
filename = parsed_url.path.split('/')[-1]
|
|
|
# 下载这个图片
|
|
|
input_image.append('Temp/' + filename) # input_image扩展为数组，支持多张图片
|
|
|
download_image(x, 'Temp/' + filename)
|
|
|
|
|
|
# WEBUI
|
|
|
if model_type_id == 1:
|
|
|
# 开始生成
|
|
|
log.info("发现需要处理的WebUI任务,task_id=" + str(task_id))
|
|
|
runWebUI(model_id, task_type_code, input_image, total_count,
|
|
|
user_id, prompt) # 原来的input_image是单张图片，现在修改为支持多张图片
|
|
|
|
|
|
# COMFY_UI
|
|
|
if model_type_id == 2:
|
|
|
log.info("发现需要处理的ComfyUI任务,task_id=" + str(task_id))
|
|
|
# 使用哪个文件
|
|
|
json_file = './JSON/' + str(model_id) + '.json'
|
|
|
# 请求
|
|
|
runComfyUI(json_file=json_file,
|
|
|
model_id=model_id,
|
|
|
task_type_code=task_type_code,
|
|
|
input_image=input_image,
|
|
|
user_id=user_id, prompt=prompt)
|
|
|
|
|
|
# Midjourney
|
|
|
if model_type_id == 3:
|
|
|
log.info("发现需要处理的Midjourney任务...")
|
|
|
# 生图
|
|
|
ret = imagine(api_server=ApiServer, api_token=ApiToken, prompt=prompt, type_code=MS_RELAX,
|
|
|
input_img=input_image[0])
|
|
|
if ret['code'] == 1:
|
|
|
mj_task_id = ret['result']
|
|
|
# 回写WEB服务器数据，表示任务处理中,防止再次被取回来
|
|
|
write_mj_task_id(task_id, mj_task_id)
|
|
|
log.info("生图任务创建成功:mj_task_id = " + mj_task_id)
|
|
|
else:
|
|
|
log.error("生图任务创建生成失败," + str(ret))
|
|
|
|
|
|
return model_type_id, task_id
|
|
|
elif success == 0:
|
|
|
log.info("现在没有待处理的任务,休息10秒后再试...")
|
|
|
return 0, 0
|
|
|
else:
|
|
|
message = response.json()["message"]
|
|
|
log.warning('发生了异常:' + message)
|
|
|
return success
|
|
|
else:
|
|
|
log.error('请求失败，状态码：', response.status_code)
|
|
|
return 0, -1
|
|
|
|
|
|
|
|
|
# 回写MJ任务已发出状态
|
|
|
def write_mj_task_id(task_id, mj_task_id):
|
|
|
# 获取签名
|
|
|
enData = getenData(machine_id)
|
|
|
# 请求的地址
|
|
|
url = web_url + '/QingLong/HuiYa/writeMjTaskId'
|
|
|
# 要发送的数据，可以是字典形式
|
|
|
data = {
|
|
|
'task_id': task_id,
|
|
|
'mj_task_id': mj_task_id,
|
|
|
'enData': enData
|
|
|
}
|
|
|
# 发送POST请求
|
|
|
response = requests.post(url, data=data)
|
|
|
|
|
|
# 检查请求是否成功
|
|
|
if response.status_code == 200:
|
|
|
log.info("成功完成MJ任务发送状态数据回写！")
|
|
|
else:
|
|
|
log.error('请求失败，状态码：', response.status_code)
|
|
|
|
|
|
|
|
|
# 回复服务器，生图任务完成
|
|
|
def finish_task(task_id, target_img_urls):
|
|
|
# 获取签名
|
|
|
enData = getenData(machine_id)
|
|
|
|
|
|
# 请求的地址
|
|
|
url = web_url + '/QingLong/HuiYa/finishTask'
|
|
|
|
|
|
# 要发送的数据，可以是字典形式
|
|
|
data = {
|
|
|
'task_id': task_id,
|
|
|
'enData': enData,
|
|
|
'target_img_urls': target_img_urls
|
|
|
}
|
|
|
# 发送POST请求
|
|
|
response = requests.post(url, data=data)
|
|
|
# 检查请求是否成功
|
|
|
if response.status_code == 200:
|
|
|
log.info("成功完成数据回写！")
|
|
|
else:
|
|
|
log.error('请求失败，状态码：', response.status_code)
|
|
|
|
|
|
|
|
|
# 删除云端服务器关于MJ的任务
|
|
|
def delete_mj_task_id(task_id):
|
|
|
# 获取签名
|
|
|
enData = getenData(machine_id)
|
|
|
|
|
|
# 请求的地址
|
|
|
url = web_url + '/QingLong/HuiYa/deleteMjTaskId'
|
|
|
|
|
|
# 要发送的数据，可以是字典形式
|
|
|
data = {
|
|
|
'task_id': task_id,
|
|
|
'enData': enData
|
|
|
}
|
|
|
# 发送POST请求
|
|
|
response = requests.post(url, data=data)
|
|
|
|
|
|
# 检查请求是否成功
|
|
|
if response.status_code == 200:
|
|
|
log.info("成功完成MJ任务重发！task_id=" + str(task_id))
|
|
|
else:
|
|
|
log.error('MJ任务重发请求失败，状态码：', response.status_code)
|
|
|
|
|
|
|
|
|
# 后台线程用来跑midjourney的任务
|
|
|
def do_midjourmey():
|
|
|
while True:
|
|
|
try:
|
|
|
# 获取签名
|
|
|
enData = getenData(machine_id)
|
|
|
# 请求的地址
|
|
|
url = web_url + '/QingLong/HuiYa/getMjFetchRecord'
|
|
|
# 要发送的数据，可以是字典形式
|
|
|
data = {
|
|
|
'enData': enData,
|
|
|
}
|
|
|
# 发送GET请求
|
|
|
response = requests.get(url, data=data)
|
|
|
|
|
|
# 检查请求是否成功
|
|
|
if response.status_code == 200:
|
|
|
jo = json.loads(response.text)
|
|
|
|
|
|
if jo['success']:
|
|
|
# 绘鸭的任务id
|
|
|
task_id = jo['data']['task_id']
|
|
|
# MJ API的任务id
|
|
|
mj_task_id = jo['data']['mj_task_id']
|
|
|
# task_type_code
|
|
|
task_type_code = jo['data']['task_type_code']
|
|
|
# model_id
|
|
|
model_id = jo['data']['model_id']
|
|
|
|
|
|
# 取回来保存的文件名称
|
|
|
file_name = str(uuid.uuid4()) + '.png'
|
|
|
file_path = 'Out/Images/' + task_type_code + '/' + str(model_id) + '/' + file_name
|
|
|
ret = try_down_mj_img(ApiToken, mj_task_id, file_path)
|
|
|
|
|
|
# 图片地址
|
|
|
target_img_urls = ''
|
|
|
|
|
|
if ret[0] == 1:
|
|
|
dir_path = 'Out/Images/' + task_type_code + '/' + str(model_id)
|
|
|
images = split_4_image(file_path, dir_path)
|
|
|
for img in images:
|
|
|
# 将生成图片上传到OSS
|
|
|
key = 'Images/' + task_type_code + "/" + str(model_id) + '/' + img
|
|
|
uploadOss(key, os.path.join(dir_path, img))
|
|
|
# # 生成下载链接
|
|
|
url = 'http://hzkc.oss-cn-beijing.aliyuncs.com/' + key
|
|
|
target_img_urls = url + ',' + target_img_urls
|
|
|
# 删除本地文件
|
|
|
os.remove(os.path.join(dir_path, img))
|
|
|
log.info("成功生成MJ图片,地址：" + url)
|
|
|
# 删除4合1图片
|
|
|
os.remove(file_path)
|
|
|
# 记录到数据库
|
|
|
target_img_urls = target_img_urls[0:len(target_img_urls) - 1]
|
|
|
finish_task(task_id, target_img_urls)
|
|
|
elif ret[0] == 0: # 还没有生成完毕
|
|
|
log.info(str(mj_task_id) + " " + ret[1])
|
|
|
time.sleep(3)
|
|
|
else: # 生成失败,休息10秒，可能原因：[Queue full] Your job queue is full. Please wait for a job to finish first, then resubmit this one.
|
|
|
log.error("【MJ】生成失败，" + ret[1])
|
|
|
time.sleep(20)
|
|
|
# 通知云端服务器，此任务需要重新生成
|
|
|
delete_mj_task_id(task_id)
|
|
|
else:
|
|
|
# log.info("【MJ】没有找到取回来的任务...")
|
|
|
time.sleep(3)
|
|
|
else:
|
|
|
log.error('请求失败，状态码：', response.status_code)
|
|
|
|
|
|
|
|
|
except Exception as err:
|
|
|
log.error(str(err))
|
|
|
time.sleep(3)
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
# 禁止代理
|
|
|
set_http_proxy("")
|
|
|
# 配置文件
|
|
|
config = ConfigUtil.getConfig()
|
|
|
# 处理机编号
|
|
|
machine_id = config['system']['machine_id']
|
|
|
# 生图服务地址
|
|
|
txt2img_url = config['webui']['txt2img_url']
|
|
|
# WEB服务器地址
|
|
|
web_url = config['webui']['web_url']
|
|
|
# 配置文件
|
|
|
server_address = config.get('comfyui', 'server_address')
|
|
|
# 第三方接入TOKEN
|
|
|
ApiToken = config['api']['ApiToken']
|
|
|
# 第三方接入服务器的Server地址
|
|
|
ApiServer = config['api']['ApiServer']
|
|
|
|
|
|
# 绘鸭日志文件
|
|
|
LOG = Logger('../Log/HuiYa.log', level='debug')
|
|
|
log = LOG.logger
|
|
|
|
|
|
# 创建后台线程执行日志记录
|
|
|
thread = threading.Thread(target=do_midjourmey)
|
|
|
thread.daemon = True # 设置为后台线程
|
|
|
thread.start()
|
|
|
|
|
|
while True:
|
|
|
try:
|
|
|
res = [] # 所有图片的地址路径
|
|
|
model_type_id, task_id = get_task()
|
|
|
# 回复上位生成完毕
|
|
|
if task_id > 0 and model_type_id < 3:
|
|
|
# 图片地址
|
|
|
target_img_urls = ''
|
|
|
for r in res:
|
|
|
target_img_urls = r + ',' + target_img_urls
|
|
|
target_img_urls = target_img_urls[0:len(target_img_urls) - 1]
|
|
|
finish_task(task_id, target_img_urls)
|
|
|
elif task_id == 0:
|
|
|
time.sleep(10)
|
|
|
except Exception as err:
|
|
|
printf("发生异常，将休息10秒后重试..." + str(err))
|
|
|
time.sleep(10)
|
