From 89ee8735e0d038c1dfe817f7a4c0e18f883be6a0 Mon Sep 17 00:00:00 2001
From: Yuanpei-Wu <fiverwyp@gmail.com>
Date: Tue, 2 Jul 2024 11:39:46 +0800
Subject: [PATCH] message

enable origin send_sync
---
 .gitignore |  3 ++-
 main.py    |  8 ++++++++
 server.py  | 16 ++++++++++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 5092c98f4..b212b0297 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,4 +18,5 @@ venv/
 /tests-ui/data/object_info.json
 /user/
 *.log
-web_custom_versions/
\ No newline at end of file
+web_custom_versions/
+.DS_Store
diff --git a/main.py b/main.py
index 479643b84..065a71eb3 100644
--- a/main.py
+++ b/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}")
diff --git a/server.py b/server.py
index 23ca2fd33..411e8355a 100644
--- a/server.py
+++ b/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))