support on prompt event handler (#765)

Co-authored-by: Lt.Dr.Data <lt.dr.data@gmail.com>
This commit is contained in:
Dr.Lt.Data 2023-08-28 13:52:22 +09:00 committed by GitHub
parent b8c7c770d3
commit 0faee1186f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 17 additions and 0 deletions

View File

@ -1,6 +1,8 @@
import os import os
import sys import sys
import asyncio import asyncio
import traceback
import nodes import nodes
import folder_paths import folder_paths
import execution import execution
@ -88,6 +90,8 @@ class PromptServer():
self.last_node_id = None self.last_node_id = None
self.client_id = None self.client_id = None
self.on_prompt_handlers = []
@routes.get('/ws') @routes.get('/ws')
async def websocket_handler(request): async def websocket_handler(request):
ws = web.WebSocketResponse() ws = web.WebSocketResponse()
@ -438,6 +442,7 @@ class PromptServer():
resp_code = 200 resp_code = 200
out_string = "" out_string = ""
json_data = await request.json() json_data = await request.json()
json_data = self.trigger_on_prompt(json_data)
if "number" in json_data: if "number" in json_data:
number = float(json_data['number']) number = float(json_data['number'])
@ -606,3 +611,15 @@ class PromptServer():
if call_on_start is not None: if call_on_start is not None:
call_on_start(address, port) 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