message

enable origin send_sync

.gitignore

@@ -18,4 +18,5 @@ venv/
 /tests-ui/data/object_info.json
 /user/
 *.log
-web_custom_versions/
+web_custom_versions/
+.DS_Store

main.py

@@ -7,6 +7,9 @@ import folder_paths
 import time
 from comfy.cli_args import args
 
+from custom_nodes.comfy_scheduler.redis_client import RedisClient
+from custom_nodes.comfy_scheduler.task_consumer import task_daemon_scheduler
+
 
 def execute_prestartup_script():
     def execute_script(script_path):
@@ -233,6 +236,11 @@ if __name__ == "__main__":
 
     threading.Thread(target=prompt_worker, daemon=True, args=(q, server,)).start()
 
+    # comfy scheduler
+    comfy_execution = execution.PromptExecutor(server)
+    redis_op = RedisClient()
+    threading.Thread(target=task_daemon_scheduler, daemon=True, args=(server, comfy_execution, redis_op,)).start()
+
     if args.output_directory:
         output_dir = os.path.abspath(args.output_directory)
         logging.info(f"Setting output directory to: {output_dir}")

server.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import asyncio
+import time
 import traceback
 
 import nodes
@@ -29,6 +30,7 @@ import node_helpers
 from app.frontend_management import FrontendManager
 from app.user_manager import UserManager
 
+from custom_nodes.comfy_scheduler.redis_client import RedisClient
 
 class BinaryEventTypes:
     PREVIEW_IMAGE = 1
@@ -98,6 +100,8 @@ class PromptServer():
 
         self.on_prompt_handlers = []
 
+        self.redis_op = RedisClient()
+
         @routes.get('/ws')
         async def websocket_handler(request):
             ws = web.WebSocketResponse()
@@ -651,6 +655,18 @@ class PromptServer():
             await send_socket_catch_exception(self.sockets[sid].send_json, message)
 
     def send_sync(self, event, data, sid=None):
+        if hasattr(self, "last_prompt_id") is not True:
+            logging.warning("server has not attr last_prompt_id")
+            return
+        prompt_id = self.last_prompt_id
+        key = "comfy_response_" + prompt_id
+        message = {
+            "type": event,
+            "data": data,
+        }
+        res = self.redis_op.zadd_with_expire(key, json.dumps(message), int(time.time()), 120)
+        if not res:
+            logging.warning(f"Failed to send {event}")
         self.loop.call_soon_threadsafe(
             self.messages.put_nowait, (event, data, sid))