From dfaf02a94653b6cb3e6a843b81051aa74228363e Mon Sep 17 00:00:00 2001 From: Austin Mroz Date: Thu, 22 Aug 2024 13:38:37 -0500 Subject: [PATCH] Minimal compatibility with execution inversion Bare bones changes to make individual node caching work post execution inversion. This is likely non-viable moving forward information is pushed into the output cache too late to block prior steps. I plan a more correct implementation that instead wraps the cache. This would mean that no changes to the actual execution pathing are required, but prevents making things forward and backwards compatible --- workflowcheckpointing.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/workflowcheckpointing.py b/workflowcheckpointing.py index 44f10d6..6aea7a8 100644 --- a/workflowcheckpointing.py +++ b/workflowcheckpointing.py @@ -411,17 +411,16 @@ class CheckpointSampler(comfy.samplers.KSAMPLER): if step == int(os.environ['FORCE_CRASH_AT']): raise Exception("Simulated Crash") -original_recursive_execute = execution.recursive_execute +original_recursive_execute = execution.execute def recursive_execute_injection(*args): - unique_id = args[3] - class_type = args[1][unique_id]['class_type'] + class_type = args[1].get_node(unique_id)['class_type'] extra_data = args[4] if class_type in SAMPLER_NODES: data, metadata = checkpoint.get(unique_id) if metadata is not None and 'step' in metadata: - args[1][unique_id]['inputs']['latent_image'] = ['checkpointed'+unique_id, 0] - args[2]['checkpointed'+unique_id] = [[{'samples': data['x']}]] + args[1].get_node(unique_id)['inputs']['latent_image'] = ['checkpointed'+unique_id, 0] + args[2].outputs.set('checkpointed'+unique_id, [[{'samples': data['x']}]]) elif metadata is not None and 'completed' in metadata: outputs = json.loads(metadata['completed']) for x in range(len(outputs)): @@ -429,15 +428,15 @@ def recursive_execute_injection(*args): outputs[x] = list(data[str(x)]) elif outputs[x] == 'latent': outputs[x] = [{'samples': l} for l in data[str(x)]] - args[2][unique_id] = outputs + args[2].outputs.set(unique_id, outputs) return True, None, None res = original_recursive_execute(*args) #Conditionally save node output #TODO: determine which non-sampler nodes are worth saving - if class_type in SAMPLER_NODES and unique_id in args[2]: + if class_type in SAMPLER_NODES and args[2].outputs.get(unique_id) is not None: data = {} - outputs = args[2][unique_id].copy() + outputs = args[2].outputs.get(unique_id).copy() for x in range(len(outputs)): if isinstance(outputs[x][0], torch.Tensor): data[str(x)] = torch.stack(outputs[x]) @@ -469,7 +468,7 @@ def execute_injection(*args, **kwargs): completion_futures[args[3]['completion_future']].set_result(outputs) comfy.samplers.KSAMPLER = CheckpointSampler -execution.recursive_execute = recursive_execute_injection +execution.execute = recursive_execute_injection execution.PromptExecutor.execute = execute_injection NODE_CLASS_MAPPINGS = {}