From 3adb344fe3b7e55342eecafc82e63fc82ad0e5f8 Mon Sep 17 00:00:00 2001
From: InconsolableCellist
 <23345188+InconsolableCellist@users.noreply.github.com>
Date: Fri, 28 Apr 2023 00:00:38 -0600
Subject: [PATCH] Added nodes to save and load latent data. Also started adding
 an incrementing frame counter node

---
 comfy/utils.py | 14 ++++++++++++++
 nodes.py       | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/comfy/utils.py b/comfy/utils.py
index 68f93403c..0e768b0d7 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -1,4 +1,8 @@
+import os
 import torch
+import folder_paths
+import numpy as np
+
 
 def load_torch_file(ckpt, safe_load=False):
     if ckpt.lower().endswith(".safetensors"):
@@ -17,6 +21,16 @@ def load_torch_file(ckpt, safe_load=False):
             sd = pl_sd
     return sd
 
+
+def save_latent(samples, filename_prefix):
+    filename = os.path.join(folder_paths.get_output_directory(), (filename_prefix + "_latent.npy"))
+    np.save(filename, samples)
+
+
+def load_latent(filename):
+    filename = os.path.join(folder_paths.get_output_directory(), filename)
+    return torch.from_numpy(np.load(filename))
+
 def transformers_convert(sd, prefix_from, prefix_to, number):
     resblock_to_replace = {
         "ln_1": "layer_norm1",
diff --git a/nodes.py b/nodes.py
index 0a9513bed..111ef8712 100644
--- a/nodes.py
+++ b/nodes.py
@@ -595,6 +595,35 @@ class LatentUpscale:
         s["samples"] = comfy.utils.common_upscale(samples["samples"], width // 8, height // 8, upscale_method, crop)
         return (s,)
 
+class SaveLatent:
+   @classmethod
+   def INPUT_TYPES(s):
+       return {"required": { "samples": ("LATENT",),
+                             "filename_prefix": ("STRING", {"default": "ComfyUI"})}}
+   RETURN_TYPES = ("LATENT",)
+   FUNCTION = "save"
+
+   CATEGORY = "latent"
+
+   def save(self, samples, filename_prefix):
+       s = samples.copy()
+       comfy.utils.save_latent(samples["samples"], filename_prefix)
+       return (samples,)
+
+class LoadLatent:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "filename": ("STRING", {"default": "ComfyUI_latent.npy"})}}
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "load"
+
+    CATEGORY = "latent"
+
+    def load(self, filename):
+        derp = ({"samples": comfy.utils.load_latent(filename)},)
+        return derp
+
 class LatentRotate:
     @classmethod
     def INPUT_TYPES(s):
@@ -1093,6 +1122,18 @@ class ImagePadForOutpaint:
 
         return (new_image, mask)
 
+class FrameCounter:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "frame": ("INT", {"default": 0})}}
+
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "frame_counter"
+
+    CATEGORY = "operations"
+
+    def frame_counter(self, frame):
+        return (frame,)
 
 NODE_CLASS_MAPPINGS = {
     "KSampler": KSampler,
@@ -1121,6 +1162,8 @@ NODE_CLASS_MAPPINGS = {
     "LatentRotate": LatentRotate,
     "LatentFlip": LatentFlip,
     "LatentCrop": LatentCrop,
+    "SaveLatent": SaveLatent,
+    "LoadLatent": LoadLatent,
     "LoraLoader": LoraLoader,
     "CLIPLoader": CLIPLoader,
     "CLIPVisionEncode": CLIPVisionEncode,
@@ -1140,6 +1183,7 @@ NODE_CLASS_MAPPINGS = {
 
     "CheckpointLoader": CheckpointLoader,
     "DiffusersLoader": DiffusersLoader,
+    "FrameCounter": FrameCounter,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -1176,6 +1220,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "EmptyLatentImage": "Empty Latent Image",
     "LatentUpscale": "Upscale Latent",
     "LatentComposite": "Latent Composite",
+    "SaveLatent": "Save Latent",
+    "LoadLatent": "Load Latent",
     # Image
     "SaveImage": "Save Image",
     "PreviewImage": "Preview Image",
@@ -1188,6 +1234,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     # _for_testing
     "VAEDecodeTiled": "VAE Decode (Tiled)",
     "VAEEncodeTiled": "VAE Encode (Tiled)",
+    # operations
+    "FrameCounter": "Frame Counter",
 }
 
 def load_custom_node(module_path):