xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-12 07:46:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

message

Browse Source 
enable origin send_sync
This commit is contained in:
Yuanpei-Wu 2024-07-02 11:39:46 +08:00
parent 369f459b20
commit 89ee8735e0
3 changed files with 26 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

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

8
main.py
View File

@ -7,6 +7,9 @@ import folder_paths
import time import time
from comfy.cli_args import args 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_prestartup_script():
def execute_script(script_path): def execute_script(script_path):
@ -233,6 +236,11 @@ if __name__ == "__main__":
threading.Thread(target=prompt_worker, daemon=True, args=(q, server,)).start() 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: if args.output_directory:
output_dir = os.path.abspath(args.output_directory) output_dir = os.path.abspath(args.output_directory)
logging.info(f"Setting output directory to: {output_dir}") logging.info(f"Setting output directory to: {output_dir}")

16
server.py
View File

@ -1,6 +1,7 @@
import os import os
import sys import sys
import asyncio import asyncio
import time
import traceback import traceback
import nodes import nodes
@ -29,6 +30,7 @@ import node_helpers
from app.frontend_management import FrontendManager from app.frontend_management import FrontendManager
from app.user_manager import UserManager from app.user_manager import UserManager
from custom_nodes.comfy_scheduler.redis_client import RedisClient
class BinaryEventTypes: class BinaryEventTypes:
PREVIEW_IMAGE = 1 PREVIEW_IMAGE = 1
@ -98,6 +100,8 @@ class PromptServer():
self.on_prompt_handlers = [] self.on_prompt_handlers = []
self.redis_op = RedisClient()
@routes.get('/ws') @routes.get('/ws')
async def websocket_handler(request): async def websocket_handler(request):
ws = web.WebSocketResponse() ws = web.WebSocketResponse()
@ -651,6 +655,18 @@ class PromptServer():
await send_socket_catch_exception(self.sockets[sid].send_json, message) await send_socket_catch_exception(self.sockets[sid].send_json, message)
def send_sync(self, event, data, sid=None): 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.loop.call_soon_threadsafe(
self.messages.put_nowait, (event, data, sid)) self.messages.put_nowait, (event, data, sid))