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

@@ -1,4 +1,4 @@
-from . import workflowcheckpointing
+from . import workflowcheckpointing, orchestration
 
 NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}

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)

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 = {}