mirror of
https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git
synced 2026-03-31 07:16:59 +08:00
Add node to save image along with .txt file for captions (SaveImageKJ)
This commit is contained in:
parent
13242fc3bf
commit
71904a4eca
@ -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"},
|
||||
|
||||
@ -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,) }
|
||||
@ -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):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user