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]