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 import requests 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 = "深度图批处理" button = WebDriverWait(driver, 10).until( EC.element_to_be_clickable((By.XPATH, f'//div[contains(text(), "{target_text}")]')) ) button.click() print("选择版本按钮已成功点击") time.sleep(3) # 使用XPath定位包含特定文本的div元素 target_text = "Extra options" button = WebDriverWait(driver, 10).until( EC.element_to_be_clickable((By.XPATH, f'//label[contains(text(), "{target_text}")]')) ) button.click() # 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,styleNumber) 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" time.sleep(5) utils.oss().put_object_from_file(ossPath, file_path) uploadCounts += 1 elif "png" in file: ossPath = f"photos/{pid}/cartoon_deep/{styleNumber}/{styleNumber}.png" time.sleep(5) 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() if pid > 1000: 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 downloadPath = "D:\downloadPhotos" inputFolder = "c:\deep" #"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:\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) if not os.path.exists(downloadPath): os.makedirs(downloadPath) 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.spop('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: #创建临时文件 以pid 和 orderId 和 styleNumber 为文件名 文件内容为空 tempFile = os.path.join(downloadPath, f"{pid}_{orderId}_{styleNumber}.txt") with open(tempFile, 'w') as f: f.write('') main(pid,orderId,styleNumber) 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