Save as single file, support cache completed nodes
Logic has been further cleaned up in preparation for networking integration. There's not longer unique behaviour for what is being saved and the onus is placed on the caller to provide data in a safetensor appropriate format Supported has been added to cache the output of completed nodes. This code should be trivially expandable to non-sampler nodes, but will need to be cautiously enabled on a per-node basis to better regulate dataflow
This commit is contained in:
Austin Mroz
+48
-41
@@ -1,7 +1,7 @@
|
||||
import torch
|
||||
import os
|
||||
import json
|
||||
import safetensors
|
||||
import safetensors.torch
|
||||
|
||||
import comfy.samplers
|
||||
import execution
|
||||
@@ -9,44 +9,26 @@ import server
|
||||
|
||||
SAMPLER_NODES = ["SamplerCustom", "KSampler", "KSamplerAdvanced", "SamplerCustomAdvanced"]
|
||||
|
||||
def store_checkpoint(unique_id, data, partial_progress_counter=-2, is_tensor=True):
|
||||
def store_checkpoint(unique_id, tensors, metadata, priority=0):
|
||||
"""Swappable interface for saving checkpoints.
|
||||
Implementation must be transactional: Either the whole thing completes,
|
||||
or the prior checkpoint must be valid even if crash occurs mid execution"""
|
||||
file = f"checkpoint/{unique_id}.{partial_progress_counter%2}.safetensor"
|
||||
if is_tensor:
|
||||
comfy.utils.save_torch_file({'x': data},file)
|
||||
else:
|
||||
with open(file, "w") as f:
|
||||
json.dump(data,f)
|
||||
with open(f"checkpoint/{unique_id}.txt", "w") as f:
|
||||
f.write(str(partial_progress_counter)+"\n"+str(1*is_tensor))
|
||||
file = f"checkpoint/{unique_id}.checkpoint"
|
||||
safetensors.torch.save_file(tensors, file, metadata)
|
||||
def get_checkpoint(unique_id):
|
||||
"""Returns the information previously saved"""
|
||||
partial_progress = -2
|
||||
is_tensor = True
|
||||
if os.path.exists(f"checkpoint/{unique_id}.txt"):
|
||||
with open(f"checkpoint/{unique_id}.txt", "r") as f:
|
||||
partial_progress = int(f.readline())
|
||||
is_tensor = f.readline() != '0'
|
||||
file = f"checkpoint/{unique_id}.{partial_progress%2}.safetensor"
|
||||
file = f"checkpoint/{unique_id}.checkpoint"
|
||||
if not os.path.exists(file):
|
||||
return None, None
|
||||
if is_tensor:
|
||||
res = safetensors.torch.load_file(file)['x']
|
||||
else:
|
||||
with open(file, "r") as f:
|
||||
res = json.load(f)
|
||||
return partial_progress, res
|
||||
with safetensors.torch.safe_open(file, framework='pt' ) as f:
|
||||
metadata = f.metadata()
|
||||
tensors = {key:f.get_tensor(key) for key in f.keys()}
|
||||
return tensors, metadata
|
||||
def reset_checkpoints(unique_id=None):
|
||||
"""Clear all checkpoint information."""
|
||||
if unique_id is not None:
|
||||
if os.path.exists(f"checkpoint/{unique_id}.0.safetensor"):
|
||||
os.remove(f"checkpoint/{unique_id}.0.safetensor")
|
||||
if os.path.exists(f"checkpoint/{unique_id}.1.safetensor"):
|
||||
os.remove(f"checkpoint/{unique_id}.1.safetensor")
|
||||
if os.path.exists(f"checkpoint/{unique_id}.txt"):
|
||||
os.remove(f"checkpoint/{unique_id}.txt")
|
||||
if os.path.exists(f"checkpoint/{unique_id}.checkpoint"):
|
||||
os.remove(f"checkpoint/{unique_id}.checkpoint")
|
||||
return
|
||||
for file in os.listdir("checkpoint"):
|
||||
os.remove(os.path.join("checkpoint", file))
|
||||
@@ -55,8 +37,11 @@ class CheckpointSampler(comfy.samplers.KSAMPLER):
|
||||
def sample(self, *args, **kwargs):
|
||||
args = list(args)
|
||||
self.unique_id = server.PromptServer.instance.last_node_id
|
||||
self.step, data = get_checkpoint(self.unique_id)
|
||||
if self.step is not None:
|
||||
self.step = None
|
||||
data, metadata = get_checkpoint(self.unique_id)
|
||||
if metadata is not None and 'step' in metadata:
|
||||
data = data['x']
|
||||
self.step = int(metadata['step'])
|
||||
#checkpoint of execution exists
|
||||
args[5] = data.to(args[4].device)
|
||||
args[1] = args[1][self.step:]
|
||||
@@ -75,29 +60,51 @@ class CheckpointSampler(comfy.samplers.KSAMPLER):
|
||||
def callback(self, step, denoised, x, total_steps):
|
||||
if self.step is not None:
|
||||
step += self.step
|
||||
store_checkpoint(self.unique_id, x, step)
|
||||
data = safetensors.torch.save
|
||||
store_checkpoint(self.unique_id, {'x':x}, {'step':str(step)})
|
||||
|
||||
original_recursive_execute = execution.recursive_execute
|
||||
def recursive_execute_injection(*args):
|
||||
|
||||
unique_id = args[3]
|
||||
class_type = args[1][unique_id]['class_type']
|
||||
#Imperfect, is checked for each bubble down step
|
||||
#Only applied once, but has unnecessary loads
|
||||
if len(args[5]) == 0:
|
||||
_, prev_prompt = get_checkpoint('prompt')
|
||||
if prev_prompt != args[1]:
|
||||
metadata = get_checkpoint('prompt')[1]
|
||||
if metadata is None or json.loads(metadata['prompt']) != args[1]:
|
||||
reset_checkpoints()
|
||||
store_checkpoint('prompt', args[1], is_tensor=False)
|
||||
store_checkpoint('prompt', {'x': torch.ones(1)},
|
||||
{'prompt': json.dumps(args[1])}, priority=2)
|
||||
if class_type in SAMPLER_NODES:
|
||||
step, x = get_checkpoint(unique_id)
|
||||
if step is not None and step>0:
|
||||
data, metadata = get_checkpoint(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': x}]]
|
||||
args[2]['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)):
|
||||
if outputs[x] == 'tensor':
|
||||
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
|
||||
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:
|
||||
pass
|
||||
#output = args[2][unique_id]
|
||||
if class_type in SAMPLER_NODES and unique_id in args[2]:
|
||||
data = {}
|
||||
outputs = args[2][unique_id].copy()
|
||||
for x in range(len(outputs)):
|
||||
if isinstance(outputs[x][0], torch.Tensor):
|
||||
data[str(x)] = torch.stack(outputs[x])
|
||||
outputs[x] = 'tensor'
|
||||
elif isinstance(outputs[x][0], dict):
|
||||
data[str(x)] = torch.stack([l['samples'] for l in outputs[x]])
|
||||
outputs[x] = 'latent'
|
||||
store_checkpoint(unique_id, data, {'completed': json.dumps(outputs)}, priority=1)
|
||||
return res
|
||||
|
||||
comfy.samplers.KSAMPLER = CheckpointSampler
|
||||
|
||||
Reference in New Issue
Block a user