Add InsertLatentToIndexed, SaveStringKJ

__init__.py

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

nodes/image_nodes.py

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

nodes/nodes.py

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