KeCheng/BaiHu/doTask.py

import os
import os.path
import random
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 runWebUI(model_id, json_data, input_image):
    # 提示词
    prompt = json_data['prompt']
    # 图片
    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])

    # 宽度与高度
    width = json_data['width']
    height = json_data['height']
    # 提示词
    json_data['prompt'] = prompt

    # 开始生成图片
    printf("开始生成图片...")
    # 设置目标文件夹
    target_folder = "Out"

    # 创建以当前模型命名的目录
    folder = os.path.join(target_folder, 'Images/User/' + str(model_id) + '/')
    if not os.path.exists(folder):
        os.makedirs(folder)
    # 调用生图服务
    response = submit_post(txt2img_url, json_data)
    # 创建文件名
    png_file = str(uuid.uuid4()) + '.png'
    # 创建文件的完整路径
    file_path = os.path.join(folder, png_file)
    save_encoded_image(response.json()['images'][0], 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张图片,宽度：" + str(width) + ",高度:" + str(height) + "  地址：" + url)
    # 添加到结果集合中
    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))


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 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"]

    return prompt_data


def runComfyUI(model_id, json_data, input_image):
    # 创建目标目录
    output_path = "Out/Images/User/" + 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)

    # 填充输入
    json_data = fill_input(json_data, img)
    # 过滤器,哪些返回节点中获取到的图片是有用的
    myfilter = getMyFilter(json_data)
    # 生成
    files = generate_clip(server_address, json_data, client_id, output_path, myfilter)

    # 上传到云存储，回写数据接口
    for file in files:
        png_file = file + '.png'
        file_path = './Out/Images/User/' + str(model_id) + "/" + png_file
        # 清除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
        res.append(url)


# 获取任务
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 == 1 else "ComfyUI") + "任务,model_id=" + str(
                model_id) + ",task_id=" + str(task_id))

            # WEBUI
            if model_type_id == 1:
                # 开始生成
                runWebUI(model_id=model_id, json_data=vjson_data, input_image=input_image)

            # COMFY_UI
            if model_type_id == 2:
                #  请求
                runComfyUI(model_id=model_id, json_data=vjson_data, input_image=input_image)

            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']
    # 文生图服务地址
    txt2img_url = config['webui']['txt2img_url']
    # 图生图服务地址
    img2img_url = config['webui']['img2img_url']

    # WEBUI服务器地址
    web_url = config['webServer']['web_url']
    #
    # COMFYUI服务器地址
    server_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)
            elif v_task_id == 0:
                # 本轮没有任务
                time.sleep(3)

        except Exception as err:
            printf("发生异常，将休息3秒后重试..." + str(err))
            time.sleep(3)