comfyui 的脚本流程处理

comfyui/main.py

@@ -0,0 +1,236 @@
+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):
+        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
+

comfyui/util/config.toml

@@ -0,0 +1,7 @@
+
+
+[oss]
+    AccessKeyId     = "LTAI5tSReWm8hz7dSYxxth8f"
+    AccessKeySecret = "8ywTDF9upPAtvgXtLKALY2iMYHIxdS"
+    EndPoint  = "oss-cn-shanghai.aliyuncs.com"
+    BucketSuwa3dDataSecret = "suwa3d-securedata"

comfyui/util/utils.py

@@ -0,0 +1,76 @@
+import toml,time,hmac,hashlib,requests
+import oss2,os
+import redis
+#读取配置文件
+def cfg(keyName):
+    try:
+        # 获取当前文件所在目录
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        config_path = os.path.join(current_dir, 'config.toml')
+        
+        with open(config_path, 'r') as f:
+            config = toml.load(f)
+
+        print(config)
+        if "." in keyName:
+            keyName = keyName.split(".")
+            if keyName[0] in config and keyName[1] in config[keyName[0]]:
+                return config[keyName[0]][keyName[1]]
+            else:
+                print(f"配置项 {keyName} 不存在")
+                return None
+        else:
+            if keyName in config:
+                return config[keyName]
+            else:
+                print(f"配置项 {keyName} 不存在")
+                return None
+    except FileNotFoundError:
+        print(f"config.toml 文件不存在，尝试路径: {config_path}")
+        return None
+    except Exception as e:
+        print(f"读取配置文件时发生错误: {e}")
+        return None
+
+
+def oss():
+    AccessKeyId = cfg("oss.AccessKeyId")
+    AccessKeySecret = cfg("oss.AccessKeySecret")
+    Endpoint = cfg("oss.EndPoint")
+    Bucket = cfg("oss.BucketSuwa3dDataSecret")
+    
+    # 检查配置是否完整
+    if not all([AccessKeyId, AccessKeySecret, Endpoint, Bucket]):
+        raise ValueError("OSS配置不完整，请检查config.toml文件中的oss配置项")
+    
+    oss_client = oss2.Bucket(oss2.Auth(AccessKeyId, AccessKeySecret), Endpoint, Bucket)
+    return oss_client
+
+
+def create_redis_connection():
+    """创建 Redis 连接，若连接失败则重试"""
+    while True:
+        try:
+            r = redis.Redis(host="106.14.158.208",password="kcV2000",port=6379,db=6)
+            # 尝试进行一次操作，检查连接是否有效
+            r.ping()  # ping 操作是一个简单的连接测试
+            print("Redis连接成功！")
+            return r
+        except ConnectionError:
+            print("Redis连接失败，正在重试...")
+            time.sleep(5) 
+
+def notify(content):
+
+    if content == "":
+        return "content 不能为空"
+    
+    notify_user_Ids = ["18950403426"]
+    for user_agent_id in notify_user_Ids:
+        data = {
+            'userId': user_agent_id,
+            'message': content,
+        }
+        headers = {'Content-Type': 'application/json'}
+        message_send_url = "https://mp.api.suwa3d.com/api/qyNotify/sendMessage?userId="+user_agent_id+"&message="+content
+        response = requests.post(message_send_url, data=json.dumps(data), headers=headers)