From e793734e8bdfc80699fb218ba6ea348544d8d21e Mon Sep 17 00:00:00 2001 From: Fillip Date: Mon, 12 Aug 2024 21:25:41 -0700 Subject: [PATCH] Fixed datacaption saver --- FL_ImageCaptionSaver.py | 12 +++++------- __init__.py | 3 +++ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/FL_ImageCaptionSaver.py b/FL_ImageCaptionSaver.py index 26e5ed9..f8ebba6 100644 --- a/FL_ImageCaptionSaver.py +++ b/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"} \ No newline at end of file diff --git a/__init__.py b/__init__.py index 2bf2b94..eec4bd7 100644 --- a/__init__.py +++ b/__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",