Fixed datacaption saver

FL_ImageCaptionSaver.py

@@ -5,6 +5,7 @@ import numpy as np
 
 from comfy.utils import ProgressBar
 
+
 class FL_ImageCaptionSaver:
     @classmethod
     def INPUT_TYPES(cls):
@@ -19,7 +20,7 @@ class FL_ImageCaptionSaver:
 
     RETURN_TYPES = ("STRING",)
     FUNCTION = "save_images_with_captions"
-    CATEGORY = "🏵️Fill Nodes/utility"
+    CATEGORY = "🏵️Fill Nodes/Captioning"
     OUTPUT_NODE = True
 
     def sanitize_text(self, text):
@@ -45,17 +46,17 @@ class FL_ImageCaptionSaver:
 
             # Convert tensor to numpy array
             image_np = image_tensor.cpu().numpy()
-            
+
             # Ensure the image is in the correct shape (height, width, channels)
             if image_np.shape[0] == 1:  # If the first dimension is 1, squeeze it
                 image_np = np.squeeze(image_np, axis=0)
-            
+
             # If the image is grayscale (2D), convert to RGB
             if len(image_np.shape) == 2:
                 image_np = np.stack((image_np,) * 3, axis=-1)
             elif image_np.shape[2] == 1:  # If it's (height, width, 1)
                 image_np = np.repeat(image_np, 3, axis=2)
-            
+
             # Ensure values are in 0-255 range
             image_np = (image_np * 255).clip(0, 255).astype(np.uint8)
 
@@ -72,6 +73,3 @@ class FL_ImageCaptionSaver:
             pbar.update_absolute(i)
 
         return (f"Saved {len(images)} images and sanitized captions in '{folder_name}'",)
-
-NODE_CLASS_MAPPINGS = {"FL_ImageCaptionSaver": FL_ImageCaptionSaver}
-NODE_DISPLAY_NAME_MAPPINGS = {"FL_ImageCaptionSaver": "FL Image Caption Saver"}

__init__.py

@@ -7,6 +7,7 @@ from .FL_KohyaSSAdvConfig import FL_KohyaSSAdvConfig
 from .FL_KohyaSSTrain import FL_KohyaSSTrain
 from .FL_LoadImagesFromDirectoryPath import FL_LoadImagesFromDirectoryPath
 from .FL_Kohya_EasyTrain import FL_Kohya_EasyTrain
+from .FL_ImageCaptionSaver import FL_ImageCaptionSaver
 #==============================================================================
 #==============================================================================
 # from .FL_SliderLoraInitWorkspace import FL_SliderLoraInitWorkspace
@@ -22,6 +23,7 @@ NODE_CLASS_MAPPINGS = {
     "FL_KohyaSSTrain": FL_KohyaSSTrain,
     "FL_LoadImagesFromDirectoryPath": FL_LoadImagesFromDirectoryPath,
     "FL_Kohya_EasyTrain": FL_Kohya_EasyTrain,
+    "FL_ImageCaptionSaver": FL_ImageCaptionSaver,
 #==============================================================================
 #==============================================================================
     # "FL_SliderLoraInitWorkspace": FL_SliderLoraInitWorkspace,
@@ -37,6 +39,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "FL_KohyaSSTrain": "FL Kohya Train",
     "FL_LoadImagesFromDirectoryPath": "FL Kohya Data Loader",
     "FL_Kohya_EasyTrain": "FL Kohya Easy Train",
+    "FL_ImageCaptionSaver": "FL Image Caption Saver"
 #==============================================================================
 #==============================================================================
     # "FL_SliderLoraInitWorkspace": "FL Slider LoRA Init Workspace",