Add InsertLatentToIndexed, SaveStringKJ

This commit is contained in:
kijai 2025-02-11 21:15:20 +02:00
parent bfe72cc964
commit a0eebe8c9d
3 changed files with 99 additions and 1 deletions

View File

@ -75,6 +75,7 @@ NODE_CONFIG = {
"ImageResizeKJ": {"class": ImageResizeKJ, "name": "Resize Image"},
"ImageUpscaleWithModelBatched": {"class": ImageUpscaleWithModelBatched, "name": "Image Upscale With Model Batched"},
"InsertImagesToBatchIndexed": {"class": InsertImagesToBatchIndexed, "name": "Insert Images To Batch Indexed"},
"InsertLatentToIndexed": {"class": InsertLatentToIndex, "name": "Insert Latent To Index"},
"LoadAndResizeImage": {"class": LoadAndResizeImage, "name": "Load & Resize Image"},
"LoadImagesFromFolderKJ": {"class": LoadImagesFromFolderKJ, "name": "Load Images From Folder (KJ)"},
"MergeImageChannels": {"class": MergeImageChannels, "name": "Merge Image Channels"},
@ -105,6 +106,7 @@ NODE_CONFIG = {
"CustomSigmas": {"class": CustomSigmas, "name": "Custom Sigmas"},
#utility
"WidgetToString": {"class": WidgetToString, "name": "Widget To String"},
"SaveStringKJ": {"class": SaveStringKJ, "name": "Save String KJ"},
"DummyOut": {"class": DummyOut, "name": "Dummy Out"},
"GetLatentsFromBatchIndexed": {"class": GetLatentsFromBatchIndexed, "name": "Get Latents From Batch Indexed"},
"ScaleBatchPromptSchedule": {"class": ScaleBatchPromptSchedule, "name": "Scale Batch Prompt Schedule"},

View File

@ -1652,6 +1652,54 @@ Returns a range of latents from a batch.
return ({"samples": chosen_latents,},)
class InsertLatentToIndex:
RETURN_TYPES = ("LATENT", )
FUNCTION = "insert"
CATEGORY = "KJNodes/latents"
DESCRIPTION = """
Inserts a latent at the specified index into the original latent batch.
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"source": ("LATENT",),
"destination": ("LATENT",),
"index": ("INT", {"default": 0,"min": -1, "max": 4096, "step": 1}),
},
}
def insert(self, source, destination, index):
samples_destination = destination["samples"]
samples_source = source["samples"].to(samples_destination)
if len(samples_source.shape) == 4:
B, C, H, W = samples_source.shape
num_latents = B
elif len(samples_source.shape) == 5:
B, C, T, H, W = samples_source.shape
num_latents = T
if index >= num_latents or index < 0:
raise ValueError(f"Index {index} out of bounds for tensor with {num_latents} latents")
if len(samples_source.shape) == 4:
joined_latents = torch.cat([
samples_destination[:index],
samples_source,
samples_destination[index+1:]
], dim=0)
else:
joined_latents = torch.cat([
samples_destination[:, :, :index],
samples_source,
samples_destination[:, :, index+1:]
], dim=2)
return ({"samples": joined_latents,},)
class GetImagesFromBatchIndexed:
RETURN_TYPES = ("IMAGE",)
@ -2524,6 +2572,54 @@ class SaveImageKJ:
"images": results },
"result": (file,) }
class SaveStringKJ:
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": {
"string": ("STRING", {"forceInput": True, "tooltip": "string to save as .txt file"}),
"filename_prefix": ("STRING", {"default": "text", "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."}),
"output_folder": ("STRING", {"default": "output", "tooltip": "The folder to save the images to."}),
},
"optional": {
"file_extension": ("STRING", {"default": ".txt", "tooltip": "The extension for the caption file."}),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("filename",)
FUNCTION = "save_string"
OUTPUT_NODE = True
CATEGORY = "KJNodes/misc"
DESCRIPTION = "Saves the input string to your ComfyUI output directory."
def save_string(self, string, output_folder, filename_prefix="text", 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)
if output_folder != "output":
if not os.path.exists(output_folder):
os.makedirs(output_folder, exist_ok=True)
full_output_folder = output_folder
base_file_name = f"{filename_prefix}_{counter:05}_"
results = list()
txt_file = base_file_name + file_extension
file_path = os.path.join(full_output_folder, txt_file)
with open(file_path, 'w') as f:
f.write(string)
return results,
to_pil_image = T.ToPILImage()
class FastPreview:

View File

@ -2281,7 +2281,7 @@ class ImageNoiseAugmentation:
RETURN_TYPES = ("IMAGE",)
FUNCTION = "add_noise"
CATEGORY = "KJNodes/experimental"
CATEGORY = "KJNodes/image"
def add_noise(self, image, noise_aug_strength, seed):
torch.manual_seed(seed)