Add node to save image along with .txt file for captions (SaveImageKJ)

This commit is contained in:
kijai 2024-09-24 10:25:58 +03:00
parent 13242fc3bf
commit 71904a4eca
3 changed files with 89 additions and 1 deletions

View File

@ -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"},

View File

@ -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,) }

View File

@ -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):