You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

889 lines
33 KiB

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

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)