chore: 🧹 move save_tensors to new file

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]