From 0b3f5fca9d5c9e76d9608b129d4a68dcde49a3fc Mon Sep 17 00:00:00 2001 From: Austin Mroz Date: Tue, 25 Jun 2024 15:49:21 -0500 Subject: [PATCH] Include orchestration, rework execution wrapping Orchestration code (mostly endpoints for probes) was intended as a seperate plugin, but given the tight tying of asynchronous dependency code, it has been included in this history and will likely be branched off as a dedicated repo not available for non client/server architectures. much of the execution wrapping code has been moved into a wrapper on execute instead of recursive executed. This cleans up logic and fixes issues with collecting all outputs when multiple output nodes exist --- __init__.py | 2 +- orchestration.py | 40 ++++++++++++++++++++++++++++++++++ workflowcheckpointing.py | 47 ++++++++++++++++++---------------------- 3 files changed, 62 insertions(+), 27 deletions(-) create mode 100644 orchestration.py diff --git a/__init__.py b/__init__.py index b2de527..3b58167 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,4 @@ -from . import workflowcheckpointing +from . import workflowcheckpointing, orchestration NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} diff --git a/orchestration.py b/orchestration.py new file mode 100644 index 0000000..37c75cf --- /dev/null +++ b/orchestration.py @@ -0,0 +1,40 @@ +import server +#Check for availability of workflow checkpointing +import importlib +#wcp = importlib.import_module('custom_nodes.ComfyUI-WorkflowCheckpointing.workflowcheckpointing') + +web = server.web +ps = server.PromptServer.instance + +finished_startup = False +original_server_start = ps.start +async def server_start(*args, **kwargs): + args = list(args) + original_on_start= args[3] + def on_start(scheme, address, port): + if original_on_start is not None: + original_on_start(scheme, address, port) + global finished_startup + finished_startup= True + args[3] = on_start + return await original_server_start(*args, **kwargs) +ps.start = server_start + + +@ps.routes.get("/health") +async def heath(request): + #while any of the server endpoints could likely be used + return web.json_response([]) + +@ps.routes.get("/startup") +async def startup(request): + if finished_startup: + return web.json_response([]) + return web.Response(status=503) + +@ps.routes.get("/ready") +async def ready(request): + current_queue = ps.prompt_queue.get_current_queue() + if len(ps[0]) == 0 and len(ps[1]) == 0: + return web.json_response(current_queue) + return web.json_response(current_queue, status=503) diff --git a/workflowcheckpointing.py b/workflowcheckpointing.py index fce0b3b..7a1c1d3 100644 --- a/workflowcheckpointing.py +++ b/workflowcheckpointing.py @@ -189,8 +189,6 @@ class FetchLoop: awaken = asyncio.Future() self.queue.enqueue_unchecked(awaken, priority-1) await awaken - else: - print("got eof") future.set_result(filename) except: future.set_result(None) @@ -201,7 +199,6 @@ class FetchLoop: fetch_loop = FetchLoop() async def prepare_file(url, path, priority): hashloc = await fetch_loop.enqueue(url, priority) - breakpoint() if os.path.exists(path): os.remove(path) #TODO consider if symlinking would be better @@ -386,22 +383,6 @@ def recursive_execute_injection(*args): unique_id = args[3] class_type = args[1][unique_id]['class_type'] extra_data = args[4] - prev_outputs = None - if 'checkpoints' in extra_data: - checkpoint.update(extra_data.pop('checkpoints')) - if 'prompt_checked' not in args[4]: - metadata = checkpoint.get('prompt')[1] - if metadata is None or json.loads(metadata['prompt']) != args[1]: - checkpoint.reset() - checkpoint.store('prompt', {'x': torch.ones(1)}, - {'prompt': json.dumps(args[1])}, priority=2) - args[4]['prompt_checked'] = True - prev_outputs = {} - os.makedirs("temp", exist_ok=True) - #TODO: Consider subdir recursing? - for item in itertools.chain(os.scandir("output"), os.scandir("temp")): - if item.is_file(): - prev_outputs[item.path] = item.stat().st_mtime if class_type in SAMPLER_NODES: data, metadata = checkpoint.get(unique_id) if metadata is not None and 'step' in metadata: @@ -431,17 +412,31 @@ def recursive_execute_injection(*args): data[str(x)] = torch.stack([l['samples'] for l in outputs[x]]) outputs[x] = 'latent' checkpoint.store(unique_id, data, {'completed': json.dumps(outputs)}, priority=1) - if prev_outputs is not None: - outputs = [] - for item in itertools.chain(os.scandir("output"), os.scandir("temp")): - if item.is_file() and prev_outputs.get(item.path, 0) < item.stat().st_mtime: - outputs.append(item.path) - if 'completion_future' in extra_data: - completion_futures[extra_data['completion_future']].set_result(outputs) return res +original_execute = execution.PromptExecutor.execute +def execute_injection(*args, **kwargs): + metadata = checkpoint.get('prompt')[1] + if metadata is None or json.loads(metadata['prompt']) != args[1]: + checkpoint.reset() + checkpoint.store('prompt', {'x': torch.ones(1)}, + {'prompt': json.dumps(args[1])}, priority=2) + prev_outputs = {} + os.makedirs("temp", exist_ok=True) + #TODO: Consider subdir recursing? + for item in itertools.chain(os.scandir("output"), os.scandir("temp")): + if item.is_file(): + prev_outputs[item.path] = item.stat().st_mtime + original_execute(*args, **kwargs) + outputs = [] + for item in itertools.chain(os.scandir("output"), os.scandir("temp")): + if item.is_file() and prev_outputs.get(item.path, 0) < item.stat().st_mtime: + outputs.append(item.path) + if 'completion_future' in args[3]: + completion_futures[args[3]['completion_future']].set_result(outputs) comfy.samplers.KSAMPLER = CheckpointSampler execution.recursive_execute = recursive_execute_injection +execution.PromptExecutor.execute = execute_injection NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {}