diff --git a/checkpointsampling.py b/checkpointsampling.py index 2842365..b3adfdd 100644 --- a/checkpointsampling.py +++ b/checkpointsampling.py @@ -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