From 71904a4eca8a87861dfe76c1897e922495a9494b Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 24 Sep 2024 10:25:58 +0300
Subject: [PATCH] Add node to save image along with .txt file for captions
 (SaveImageKJ)

---
 __init__.py          |  2 ++
 nodes/image_nodes.py | 73 +++++++++++++++++++++++++++++++++++++++++++-
 nodes/nodes.py       | 15 +++++++++
 3 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index 08e968b..74fccf0 100644
--- a/__init__.py
+++ b/__init__.py
@@ -7,6 +7,7 @@ from .nodes.intrinsic_lora_nodes import *
 from .nodes.mask_nodes import *
 NODE_CONFIG = {
     #constants
+    "BOOLConstant": {"class": BOOLConstant, "name": "BOOL Constant"},
     "INTConstant": {"class": INTConstant, "name": "INT Constant"},
     "FloatConstant": {"class": FloatConstant, "name": "Float Constant"},
     "StringConstant": {"class": StringConstant, "name": "String Constant"},
@@ -72,6 +73,7 @@ NODE_CONFIG = {
     "ReverseImageBatch": {"class": ReverseImageBatch, "name": "Reverse Image Batch"},
     "ReplaceImagesInBatch": {"class": ReplaceImagesInBatch, "name": "Replace Images In Batch"},
     "SaveImageWithAlpha": {"class": SaveImageWithAlpha, "name": "Save Image With Alpha"},
+    "SaveImageKJ": {"class": SaveImageKJ, "name": "Save Image KJ"},
     "SplitImageChannels": {"class": SplitImageChannels, "name": "Split Image Channels"},    
     #batch cropping
     "BatchCropFromMask": {"class": BatchCropFromMask, "name": "Batch Crop From Mask"},
diff --git a/nodes/image_nodes.py b/nodes/image_nodes.py
index 165f3d0..7a03da8 100644
--- a/nodes/image_nodes.py
+++ b/nodes/image_nodes.py
@@ -8,7 +8,7 @@ import math
 import os
 import re
 import json
-import hashlib
+from PIL.PngImagePlugin import PngInfo
 try:
     import cv2
 except:
@@ -1938,3 +1938,74 @@ class ImageGridtoBatch:
        img_tensor = image.view(-1, orig_h, orig_w, C)
        
        return img_tensor,
+
+class SaveImageKJ:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+        self.compress_level = 4
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE", {"tooltip": "The images to save."}),
+                "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
+            },
+            "optional": {
+                "caption_file_extension": ("STRING", {"default": ".txt", "tooltip": "The extension for the caption file."}),
+                "caption": ("STRING", {"forceInput": True, "tooltip": "string to save as .txt file"}), 
+            },
+            "hidden": {
+                "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("filename",)
+    FUNCTION = "save_images"
+
+    OUTPUT_NODE = True
+
+    CATEGORY = "image"
+    DESCRIPTION = "Saves the input images to your ComfyUI output directory."
+
+    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, caption=None, caption_file_extension=".txt"):
+        filename_prefix += self.prefix_append
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
+        results = list()
+        for (batch_number, image) in enumerate(images):
+            i = 255. * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            metadata = None
+            if not args.disable_metadata:
+                metadata = PngInfo()
+                if prompt is not None:
+                    metadata.add_text("prompt", json.dumps(prompt))
+                if extra_pnginfo is not None:
+                    for x in extra_pnginfo:
+                        metadata.add_text(x, json.dumps(extra_pnginfo[x]))
+
+            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
+            base_file_name = f"{filename_with_batch_num}_{counter:05}_"
+            file = f"{base_file_name}.png"
+            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
+            results.append({
+                "filename": file,
+                "subfolder": subfolder,
+                "type": self.type
+            })
+            if caption is not None:
+                txt_file = base_file_name + caption_file_extension
+                file_path = os.path.join(full_output_folder, txt_file)
+                with open(file_path, 'w') as f:
+                    f.write(caption)
+
+            counter += 1
+
+
+
+        return { "ui": { 
+                "images": results },
+                "result": (file,) }
\ No newline at end of file
diff --git a/nodes/nodes.py b/nodes/nodes.py
index 638f73b..e7a92f3 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -19,6 +19,21 @@ class AnyType(str):
     return False
 any = AnyType("*")
 
+class BOOLConstant:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "value": ("BOOLEAN", {"default": True}),
+        },
+        }
+    RETURN_TYPES = ("BOOLEAN",)
+    RETURN_NAMES = ("value",)
+    FUNCTION = "get_value"
+    CATEGORY = "KJNodes/constants"
+
+    def get_value(self, value):
+        return (value,)
+    
 class INTConstant:
     @classmethod
     def INPUT_TYPES(s):