|
|
import os
|
|
|
import random
|
|
|
from Util.PngUtil import *
|
|
|
from Util import ConfigUtil
|
|
|
from Util.SDUtil import *
|
|
|
from Util.CommonUtil import *
|
|
|
from Util.OssUtil import uploadOss
|
|
|
|
|
|
|
|
|
# 获取Model对象
|
|
|
def getModel(model_id):
|
|
|
file_name = r'../JSON/' + str(model_id) + '.json'
|
|
|
# 打开JSON文件
|
|
|
if not os.path.exists(file_name):
|
|
|
printf("文件：" + file_name + "不存在，跳过！")
|
|
|
return
|
|
|
with open(file_name, 'r', encoding='utf-8') as file:
|
|
|
data = json.load(file)
|
|
|
return data
|
|
|
|
|
|
|
|
|
# 获取模型有多少个提示词
|
|
|
def getPromptCount():
|
|
|
if 'prompts' in data:
|
|
|
return len(data['prompts'])
|
|
|
return 1
|
|
|
|
|
|
|
|
|
# 运行一次提示词
|
|
|
def runWebUIOnce(model_id, task_type_code, input_image, prompt_index, batch_size):
|
|
|
# 生成多少张
|
|
|
data['batch_size'] = batch_size
|
|
|
# 图片
|
|
|
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'] = data["prompts"][prompt_index]
|
|
|
|
|
|
# 开始生成图片
|
|
|
printf("开始生成图片...")
|
|
|
# 设置目标文件夹
|
|
|
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, batch_size):
|
|
|
# 创建文件名
|
|
|
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)
|
|
|
printf("成功生成第" + str(i + 1) + "张图片,宽度：" + str(width) + ",高度:" + str(height) + " 地址：" + url)
|
|
|
printf("本轮图片生成成功, 共完成" + str(batch_size) + "张图片")
|
|
|
|
|
|
|
|
|
# 获取模型名称
|
|
|
def getModelName():
|
|
|
return data['override_settings']['sd_model_checkpoint']
|
|
|
|
|
|
|
|
|
def generate_tasks(total_tasks, prompt_count):
|
|
|
# 确保提示词的数量大于0
|
|
|
if prompt_count <= 0:
|
|
|
raise ValueError("提示词的数量必须大于0")
|
|
|
|
|
|
# 初始化任务分配列表，每个提示词分配0个任务
|
|
|
task_allocation = [0] * prompt_count
|
|
|
|
|
|
# 循环直到所有任务都被分配
|
|
|
while sum(task_allocation) < total_tasks:
|
|
|
# 随机选择一个提示词
|
|
|
prompt_index = random.randint(0, prompt_count - 1)
|
|
|
|
|
|
# 为选中的提示词分配一个随机数量的任务，确保总数不会超过total_tasks
|
|
|
additional_tasks = min(1, total_tasks - sum(task_allocation))
|
|
|
task_allocation[prompt_index] += additional_tasks
|
|
|
|
|
|
return task_allocation
|
|
|
|
|
|
|
|
|
# 整体的生成任务
|
|
|
def runWebUI(model_id, task_type_code, input_image, total_tasks):
|
|
|
global data
|
|
|
data = getModel(model_id) # data是全局变量
|
|
|
|
|
|
# 示例使用
|
|
|
prompt_count = getPromptCount() # 提示词数量
|
|
|
task_allocation = generate_tasks(total_tasks, prompt_count)
|
|
|
|
|
|
for i in range(len(task_allocation)):
|
|
|
if task_allocation[i] == 0:
|
|
|
continue
|
|
|
printf(f"模型 {getModelName()}，第{i + 1}个提示词，生成{task_allocation[i]}张图片")
|
|
|
runWebUIOnce(model_id, task_type_code, input_image, i, task_allocation[i])
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
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')
|
|
|
model_id = 21
|
|
|
task_type_code = 'System'
|
|
|
input_image = ['../Image/xiaoqiao.jpg']
|
|
|
total_count = 10
|
|
|
runWebUI(model_id, task_type_code, input_image, total_count) # 原来的input_image是单张图片，现在修改为支持多张图片
|
