feat: ✨ add back Save Tensors

nodes/debug.py

@@ -2,6 +2,9 @@
 from ..log import log
 import io, base64
 import torch
+import folder_paths
+from typing import Optional
+from pathlib import Path
 
 
 class Debug:
@@ -54,4 +57,65 @@ def do_debug(self, **kwargs):
         return output
 
 
-__nodes__ = [Debug]
+class SaveTensors:
+    """Save torch tensors (image, mask or latent) to disk, useful to debug things outside comfy"""
+
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "mtb/debug"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "filename_prefix": ("STRING", {"default": "ComfyPickle"}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "mask": ("MASK",),
+                "latent": ("LATENT",),
+            },
+        }
+
+    FUNCTION = "save"
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+    CATEGORY = "mtb/debug"
+
+    def save(
+        self,
+        filename_prefix,
+        image: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        latent: Optional[torch.Tensor] = None,
+    ):
+        (
+            full_output_folder,
+            filename,
+            counter,
+            subfolder,
+            filename_prefix,
+        ) = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+        full_output_folder = Path(full_output_folder)
+        if image is not None:
+            image_file = f"{filename}_image_{counter:05}.pt"
+            torch.save(image, full_output_folder / image_file)
+            # np.save(full_output_folder/ image_file, image.cpu().numpy())
+
+        if mask is not None:
+            mask_file = f"{filename}_mask_{counter:05}.pt"
+            torch.save(mask, full_output_folder / mask_file)
+            # np.save(full_output_folder/ mask_file, mask.cpu().numpy())
+
+        if latent is not None:
+            # for latent we must use pickle
+            latent_file = f"{filename}_latent_{counter:05}.pt"
+            torch.save(latent, full_output_folder / latent_file)
+            # pickle.dump(latent, open(full_output_folder/ latent_file, "wb"))
+
+            # np.save(full_output_folder/ latent_file, latent[""].cpu().numpy())
+
+        return f"{filename_prefix}_{counter:05}"
+
+
+__nodes__ = [Debug, SaveTensors]