From 0faee1186f49d7e1166401fd694d4ab3fa949a8d Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" <128333288+ltdrdata@users.noreply.github.com> Date: Mon, 28 Aug 2023 13:52:22 +0900 Subject: [PATCH] support on prompt event handler (#765) Co-authored-by: Lt.Dr.Data --- server.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/server.py b/server.py index 87abcc67..538f9075 100644 --- a/server.py +++ b/server.py @@ -1,6 +1,8 @@ import os import sys import asyncio +import traceback + import nodes import folder_paths import execution @@ -88,6 +90,8 @@ class PromptServer(): self.last_node_id = None self.client_id = None + self.on_prompt_handlers = [] + @routes.get('/ws') async def websocket_handler(request): ws = web.WebSocketResponse() @@ -438,6 +442,7 @@ class PromptServer(): resp_code = 200 out_string = "" json_data = await request.json() + json_data = self.trigger_on_prompt(json_data) if "number" in json_data: number = float(json_data['number']) @@ -606,3 +611,15 @@ class PromptServer(): if call_on_start is not None: call_on_start(address, port) + def add_on_prompt_handler(self, handler): + self.on_prompt_handlers.append(handler) + + def trigger_on_prompt(self, json_data): + for handler in self.on_prompt_handlers: + try: + json_data = handler(json_data) + except Exception as e: + print(f"[ERROR] An error occurred during the on_prompt_handler processing") + traceback.print_exc() + + return json_data