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.
240 lines
9.5 KiB
240 lines
9.5 KiB
from selenium import webdriver |
|
from selenium.webdriver.common.by import By |
|
from selenium.webdriver.common.keys import Keys |
|
from selenium.webdriver.chrome.service import Service |
|
from selenium.webdriver.support.ui import WebDriverWait |
|
from selenium.webdriver.support import expected_conditions as EC |
|
import time,os,sys,oss2,json,shutil |
|
from util import utils |
|
|
|
def main(pid,orderId,styleNumber): |
|
|
|
#进行一些初始化的处理 |
|
flag = init(pid,orderId,styleNumber) |
|
if flag == False: |
|
print(f"初始化步骤失败 pid:{pid} orderId:{orderId} styleNumber:{styleNumber}") |
|
utils.notify(f"初始化步骤失败 pid:{pid} orderId:{orderId} styleNumber:{styleNumber}") |
|
return |
|
|
|
# 初始化浏览器 |
|
chrome_driver_path = driver_path |
|
service = Service(executable_path=chrome_driver_path) |
|
driver = webdriver.Chrome(service=service) |
|
# 打开目标网页 |
|
driver.get('http://127.0.0.1:8188') # 替换为你的目标网页地址 |
|
# 等待页面加载完成 |
|
time.sleep(5) # 等待2秒，确保页面加载完成 |
|
# 定位按钮并点击 |
|
try: |
|
# 假设按钮的HTML元素可以通过id定位 |
|
button = driver.find_element(By.CSS_SELECTOR, '#comfy-load-button .pysssss-workflow-arrow') |
|
button.click() |
|
print("弹出版本选择按钮已成功点击") |
|
time.sleep(3) |
|
# 使用XPath定位包含特定文本的div元素 |
|
target_text = "3.小鱼AI（灰度图）图片浮雕流程-V3" |
|
button = WebDriverWait(driver, 10).until( |
|
EC.element_to_be_clickable((By.XPATH, f'//div[contains(text(), "{target_text}")]')) |
|
) |
|
button.click() |
|
print("选择版本按钮已成功点击") |
|
# |
|
time.sleep(5) |
|
button = driver.find_element(By.ID, 'queue-button') |
|
print(button) |
|
button.click() |
|
|
|
#获取当前时间戳 |
|
timestamp = int(time.time()) |
|
|
|
|
|
#检测是否完成数据，一个死循环 |
|
while True: |
|
boolV = checkIsFinish(driver, pid, orderId) |
|
if boolV: |
|
break |
|
else: |
|
#计算时间戳是否超过5分钟了 |
|
now = int(time.time()) |
|
if now - timestamp >= 300: |
|
print("超过5分钟，任务还没结束，强行停止执行下一个任务") |
|
utils.notify(f"pid:{pid} orderId:{orderId} 超过5分钟，任务还没结束，强行停止执行下一个任务") |
|
break |
|
time.sleep(5) |
|
continue |
|
|
|
|
|
except Exception as e: |
|
print(f"发生错误：{e}") |
|
utils.notify(f"发生错误：{e}") |
|
|
|
|
|
|
|
def down_obj_from_oss(fileName,pid,styleNumber): |
|
path = os.path.join(inputFolder, fileName) |
|
# 根据前缀获取文件列表 |
|
prefix = f'photos/{pid}/cartoon/cartoon_image' |
|
filelist = oss2.ObjectIteratorV2(utils.oss(), prefix=prefix) |
|
print(f"filelist:{filelist} prefix{prefix}") |
|
isDownload = False |
|
for file in filelist: |
|
if file and hasattr(file, 'key') and f"{styleNumber}.png" in file.key: |
|
utils.oss().get_object_to_file(file.key, path) |
|
isDownload = True |
|
print(f"pid:{pid} styleNumber:{styleNumber} 文件下载成功") |
|
break |
|
return isDownload |
|
|
|
|
|
def init(pid,orderId,styleNumber): |
|
#清除输入输出里的一些文件 |
|
clearFolder(outputFolder) |
|
clearFolder(inputFolder) |
|
#下载oss的文件到指定目录 修改指定json文件里的内容为指定内容 |
|
imageName = f"{pid}_{orderId}_{styleNumber}.png" |
|
downloadRes = down_obj_from_oss(imageName,pid,styleNumber) |
|
if downloadRes == False: |
|
return False |
|
#修改配置文件 3.小鱼AI（灰度图）图片浮雕流程-V3.json, 替换数据 id = 277 的 LoadImage 的 inputs[0] 的 value 为 imageName |
|
with open("3.小鱼AI（灰度图）图片浮雕流程-V3.json", 'r', encoding='utf-8') as f: |
|
data = json.load(f) |
|
for item in data['nodes']: |
|
if item['id'] == 277: |
|
item['widgets_values'][0] = imageName |
|
|
|
jsonFilePathNew = os.path.join(jsonFilePath, "3.小鱼AI（灰度图）图片浮雕流程-V3.json") |
|
with open(jsonFilePathNew, 'w', encoding='utf-8') as f: |
|
json.dump(data, f, ensure_ascii=False, indent=4) |
|
print(f"pid:{pid} orderId:{orderId} styleNumber:{styleNumber} 配置文件修改成功") |
|
#复制到指定路劲下 jsonFilePath |
|
#shutil.copy(jsonFilePath, os.path.join(jsonFilePath, "3.小鱼AI（灰度图）图片浮雕流程-V3.json")) |
|
return True |
|
|
|
def checkIsFinish(driver, pid, orderId,styleNumber): |
|
# 检测指定目录下是否有文件产生 |
|
folder = outputFolder |
|
fileLength = os.listdir(folder) |
|
uploadCounts = 0 |
|
files = "" |
|
if len(fileLength) == 2: |
|
#上传文件到指定的目录 |
|
for file in os.listdir(folder): |
|
files += file+"," |
|
file_path = os.path.join(folder, file) |
|
if "tiff" in file: |
|
ossPath = f"photos/{pid}/cartoon_deep/{styleNumber}/{styleNumber}.tiff" |
|
utils.oss().put_object_from_file(ossPath, file_path) |
|
uploadCounts += 1 |
|
elif "png" in file: |
|
ossPath = f"photos/{pid}/cartoon_deep/{styleNumber}/{styleNumber}.png" |
|
utils.oss().put_object_from_file(ossPath, file_path) |
|
uploadCounts += 1 |
|
|
|
if uploadCounts == 2: |
|
print(f"pid={pid},orderid={orderId},style_number:{styleNumber} 任务处理完成,深度图已上传到OSS") |
|
#推入到队列中 |
|
r = utils.create_redis_connection() |
|
r.lpush('model:badge_cartoon_build', json.dumps({'pid': pid, 'order_id': orderId, 'style_number': styleNumber})) |
|
|
|
driver.quit() |
|
return True |
|
else: |
|
print(f"pid={pid},orderid={orderId},style_number={styleNumber} 任务处理失败,当前文件数据为{files}") |
|
utils.notify(f"pid={pid},orderid={orderId},style_number={styleNumber} 任务处理失败,当前文件数据为{files}") |
|
return False |
|
|
|
#清除指定文件夹下的所有文件 |
|
def clearFolder(folder): |
|
for file in os.listdir(folder): |
|
#判断如果是文件夹，则整个删除 |
|
if os.path.isdir(os.path.join(folder, file)): |
|
shutil.rmtree(os.path.join(folder, file)) |
|
else: |
|
os.remove(os.path.join(folder, file)) |
|
|
|
|
|
#检测输入目录里是否有输入的文件存在 |
|
def checkInputFolder(): |
|
folder = inputFolder |
|
fileLength = os.listdir(folder) |
|
if len(fileLength) == 1: |
|
return True |
|
return False |
|
|
|
if __name__ == '__main__': |
|
pid = 0 |
|
orderId = 0 |
|
styleNumber = 0 |
|
inputFolder = "C:\StableDiffusion_xiaoyu_V3\ComfyUI-aki-v1.6-XY\ComfyUI\input" |
|
outputFolder = "C:\StableDiffusion_xiaoyu_V3\ComfyUI-aki-v1.6-XY\ComfyUI\output" |
|
jsonFilePath = "C:\StableDiffusion_xiaoyu_V3\ComfyUI-aki-v1.6-XY\ComfyUI\pysssss-workflows" |
|
# 设置浏览器驱动路径（以Chrome为例） |
|
driver_path = 'D:\chrome\chromedriver-win64\chromedriver.exe' # 替换为你的chromedriver路径 |
|
|
|
#检查文件是否存在，不存在就创建， 测试使用 |
|
if not os.path.exists(inputFolder): |
|
os.makedirs(inputFolder) |
|
if not os.path.exists(outputFolder): |
|
os.makedirs(outputFolder) |
|
if not os.path.exists(jsonFilePath): |
|
os.makedirs(jsonFilePath) |
|
|
|
|
|
|
|
arrArgs = sys.argv |
|
if len(arrArgs) == 4: |
|
pid = arrArgs[1] |
|
orderId = arrArgs[2] |
|
styleNumber = arrArgs[3] |
|
#判断是否是数字 |
|
if pid.isdigit() and orderId.isdigit() and styleNumber.isdigit(): |
|
main(pid,orderId,styleNumber) |
|
else: |
|
print("输入的参数不是数字") |
|
else: |
|
r = utils.create_redis_connection() |
|
while True: |
|
try: |
|
if r.llen('model:badge_cartoon') == 0: |
|
print('队列为空，等待10秒') |
|
time.sleep(10) |
|
continue |
|
|
|
info = r.lpop('model:badge_cartoon') |
|
if info is None: |
|
print('队列为空，等待10秒') |
|
time.sleep(10) |
|
continue |
|
|
|
# 确保info是字符串类型 |
|
if isinstance(info, bytes): |
|
info = info.decode('utf-8') |
|
|
|
info = json.loads(info) |
|
print(info) |
|
|
|
# 检查必要的字段是否存在 |
|
if not all(key in info for key in ['pid', 'order_id', 'style_number']): |
|
print(f"Redis数据格式错误，缺少必要字段: {info}") |
|
utils.notify(f"Redis数据格式错误，缺少必要字段: {info}") |
|
continue |
|
|
|
pid = int(info['pid']) |
|
orderId = int(info['order_id']) |
|
styleNumber = int(info['style_number']) |
|
print(f"pid:{pid} orderId:{orderId} style_number:{styleNumber}") |
|
if pid > 0 and orderId > 0 and styleNumber > 0: |
|
main(pid,orderId,styleNumber) |
|
break |
|
except json.JSONDecodeError as e: |
|
print(f'JSON解析错误：{e}') |
|
utils.notify(f'JSON解析错误：{e}') |
|
time.sleep(10) |
|
continue |
|
except Exception as e: |
|
print(f'错误：{e}') |
|
utils.notify(f'错误：{e}') |
|
time.sleep(10) |
|
r = utils.create_redis_connection() |
|
continue |
|
|
|
|