From eca2ea5da99968c0cfe7a41ea4445472269a46fd Mon Sep 17 00:00:00 2001
From: Mel Massadian <mel@melmassadian.com>
Date: Mon, 28 Jul 2025 22:18:18 +0200
Subject: [PATCH] =?UTF-8?q?chore:=20=F0=9F=A7=B9=20move=20save=5Ftensors?=
 =?UTF-8?q?=20to=20new=20file?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 nodes/dev.py | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)
 create mode 100644 nodes/dev.py

diff --git a/nodes/dev.py b/nodes/dev.py
new file mode 100644
index 0000000..469d4c3
--- /dev/null
+++ b/nodes/dev.py
@@ -0,0 +1,70 @@
+import folder_paths
+import torch
+
+
+class MTB_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: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        latent: torch.Tensor | None = 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__ = [MTB_SaveTensors]