diff --git a/README.md b/README.md index 89cda79..5b44fd0 100644 --- a/README.md +++ b/README.md @@ -1 +1,2 @@ -# ComfyUI-northTools \ No newline at end of file +# ComfyUI-northTools + diff --git a/__init__.py b/__init__.py index fb87470..5081f2f 100644 --- a/__init__.py +++ b/__init__.py @@ -9,7 +9,10 @@ from .nodes.from_dir_image_list import LoadImagesFromDirList from .nodes.extract_metadata_by_key import ExtractMetadataByKey from .nodes.sum_integers import SumIntegers from .nodes.from_dir_with_index_batch import LoadImagesFromDirByIndexBatch +from .nodes.from_dir_with_index_list import LoadImagesFromDirByIndexList from .nodes.boolean_list_to_indexes import BooleanIndexesToString +from .nodes.concat_history_string import ConcatHistoryString +from .nodes.image_to_true import ImageToTrue NODE_CLASS_MAPPINGS = { @@ -17,7 +20,10 @@ NODE_CLASS_MAPPINGS = { "ExtractMetadataByKey": ExtractMetadataByKey, "SumIntegers": SumIntegers, "LoadImagesFromDirByIndexBatch": LoadImagesFromDirByIndexBatch, + "LoadImagesFromDirByIndexList": LoadImagesFromDirByIndexList, "BooleanIndexesToString": BooleanIndexesToString, + "ConcatHistoryString": ConcatHistoryString, + "ImageToTrue": ImageToTrue, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -25,7 +31,10 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ExtractMetadataByKey": "Extract Metadata value by key", "SumIntegers": "Sum Integers", "LoadImagesFromDirByIndexBatch": "Load Image batch from Directory with indexes", - "BooleanIndexesToString": "Boolean list to index string" + "LoadImagesFromDirByIndexList": "Load Image list from Directory with indexes", + "BooleanIndexesToString": "Boolean list to index string", + "ConcatHistoryString": "Concatenate and remember string", + "ImageToTrue": "Image to True", } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/boolean_list_to_indexes.py b/nodes/boolean_list_to_indexes.py index 4d2ba90..d710f4a 100644 --- a/nodes/boolean_list_to_indexes.py +++ b/nodes/boolean_list_to_indexes.py @@ -3,24 +3,28 @@ """ class BooleanIndexesToString: + memory_bools: list[bool] = [] + pending_reset: bool = False + @classmethod def INPUT_TYPES(cls): return { - "required": { - "bool_list_str": ("STRING", {"default": ""}), # flexible format: TrueFalseTrue, 101, True,False,True, etc. - "bool_list": ("BOOLEAN[]", {"default": []}), # explicit list: [True, False, True] + "optional": { + "bool_list_str": ("STRING", {"default": ""}), + "bool_list": ("BOOLEAN[]", {"default": []}), + "memory_bool_list_str": ("STRING", {"default": ""}), + "memory_bool_list": ("BOOLEAN[]", {"default": []}), + "reset": ("BOOLEAN", {"default": False}), } } - RETURN_TYPES = ("STRING",) RETURN_NAMES = ("index_string",) FUNCTION = "get_indexes" CATEGORY = "util" - @staticmethod - def parse_bool_string(s): + def parse_bool_string(s: str) -> list[bool]: s = s.strip().replace(" ", "") if "," in s: tokens = s.split(",") @@ -43,12 +47,40 @@ class BooleanIndexesToString: true_values = {"true", "1"} return [t.lower() in true_values for t in tokens] + def get_indexes( + self, + bool_list_str: str = "", + bool_list=None, + memory_bool_list_str: str = "", + memory_bool_list=None, + reset: bool = False, + ): + if bool_list is None: + bool_list = [] + if memory_bool_list is None: + memory_bool_list = [] - def get_indexes(self, bool_list_str: str, bool_list): - if bool_list and len(bool_list) > 0: - bools = bool_list + if self.__class__.pending_reset: + self.__class__.memory_bools.clear() + self.__class__.pending_reset = False + + if bool_list_str: + new_bools = self.parse_bool_string(bool_list_str) + elif bool_list: + new_bools = list(bool_list) + elif memory_bool_list_str: + new_bools = self.parse_bool_string(memory_bool_list_str) + elif memory_bool_list: + new_bools = list(memory_bool_list) else: - bools = self.parse_bool_string(bool_list_str) - indexes = [str(i) for i, v in enumerate(bools) if v] - return (",".join(indexes),) + new_bools = [] + + self.__class__.memory_bools += new_bools + + self.__class__.pending_reset = reset + + indexes = [str(i) for i, v in enumerate(self.__class__.memory_bools) if v] + output = ",".join(indexes) + return (output,) + diff --git a/nodes/concat_history_string.py b/nodes/concat_history_string.py new file mode 100644 index 0000000..d71e7b6 --- /dev/null +++ b/nodes/concat_history_string.py @@ -0,0 +1,41 @@ +""" + This node receives a boolean and a string, and concatenates them to a string, that is remembered in the various runs. +""" + + +class ConcatHistoryString: + history = [] + + @classmethod + def INPUT_TYPES(cls): + return { + "optional": { + "value_bool": ("BOOLEAN", {"default": None}), + "value_str": ("STRING", {"default": ""}), + "reset": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("concatenated",) + FUNCTION = "concat" + CATEGORY = "util" + + def concat(self, value_bool=None, value_str="", reset=False): + if reset: + self.__class__.history = [] + + input_value = None + if value_bool is not None: + input_value = "True" if value_bool else "False" + elif value_str != "": + input_value = str(value_str) + + if input_value not in (None, ""): + self.__class__.history.append(input_value) + + return (",".join(self.__class__.history),) + + + + diff --git a/nodes/from_dir_image_list.py b/nodes/from_dir_image_list.py index 61c51b8..fcfad02 100644 --- a/nodes/from_dir_image_list.py +++ b/nodes/from_dir_image_list.py @@ -10,7 +10,7 @@ from PIL import Image, ImageOps class LoadImagesFromDirList: @classmethod - def INPUT_TYPES(s): + def INPUT_TYPES(cls): return { "required": { "directory": ("STRING", {"default": ""}), @@ -22,9 +22,9 @@ class LoadImagesFromDirList: } } - RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING") - RETURN_NAMES = ("IMAGE", "MASK", "FILE PATH", "METADATA") - OUTPUT_IS_LIST = (True, True, True, True) + RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING", "INT") + RETURN_NAMES = ("IMAGE", "MASK", "FILE PATH", "METADATA", "INDEX COUNT") + OUTPUT_IS_LIST = (True, True, True, True, True) FUNCTION = "load_images" CATEGORY = "image" @@ -46,18 +46,19 @@ class LoadImagesFromDirList: valid_extensions = ['.jpg', '.jpeg', '.png', '.webp'] dir_files = [f for f in dir_files if any(f.lower().endswith(ext) for ext in valid_extensions)] dir_files = sorted(dir_files) - dir_files = [os.path.join(directory, x) for x in dir_files] - dir_files = dir_files[start_index:] + dir_files_full = [os.path.join(directory, x) for x in dir_files] + dir_files_full = dir_files_full[start_index:] images = [] masks = [] file_paths = [] metadatas = [] + indexes = [] limit_images = image_load_cap > 0 image_count = 0 - for image_path in dir_files: + for idx, image_path in enumerate(dir_files_full): if os.path.isdir(image_path): continue if limit_images and image_count >= image_load_cap: @@ -65,7 +66,6 @@ class LoadImagesFromDirList: i = Image.open(image_path) i = ImageOps.exif_transpose(i) - # Extract metadata BEFORE conversion metadata = i.info.copy() exif = i.getexif() if exif: @@ -85,6 +85,7 @@ class LoadImagesFromDirList: images.append(image) masks.append(mask) file_paths.append(str(image_path)) + indexes.append(idx + start_index) image_count += 1 - return (images, masks, file_paths, metadatas) + return (images, masks, file_paths, metadatas, indexes) diff --git a/nodes/from_dir_with_index_batch.py b/nodes/from_dir_with_index_batch.py index 9f219a6..655e51a 100644 --- a/nodes/from_dir_with_index_batch.py +++ b/nodes/from_dir_with_index_batch.py @@ -13,25 +13,21 @@ class LoadImagesFromDirByIndexBatch: return { "required": { "directory": ("STRING", {"default": ""}), - "indices": ("STRING", {"default": ""}), }, + "optional": { + "indices": ("STRING", {"default": ""}), + } } - RETURN_TYPES = ("IMAGE", "MASK", "INT") FUNCTION = "load_images" - - CATEGORY = "image" - @classmethod def IS_CHANGED(cls, **kwargs): return hash(frozenset(kwargs.items())) - def _parse_indices(self, indices_str, max_len): - # Parse a string like "0,3,4" into unique, sorted, safe integer indexes within [0,max_len-1] indices = [] for idx in indices_str.split(","): idx = idx.strip() @@ -41,23 +37,21 @@ class LoadImagesFromDirByIndexBatch: indices.append(i) return sorted(set(indices)) - - def load_images(self, directory: str, indices: str): + def load_images(self, directory: str, indices: str = ""): if not os.path.isdir(directory): - raise FileNotFoundError(f"Directory '{directory}' cannot be found.") + return (None, None, 0) dir_files = os.listdir(directory) valid_extensions = ['.jpg', '.jpeg', '.png', '.webp'] dir_files = [f for f in dir_files if any(f.lower().endswith(ext) for ext in valid_extensions)] dir_files = sorted(dir_files) - if not dir_files: - raise FileNotFoundError(f"No valid image files found in '{directory}'.") + if not indices.strip(): + raise StopIteration("Waiting for valid indices string.") index_list = self._parse_indices(indices, len(dir_files)) if not index_list: - raise ValueError("No valid indexes found in input or indexes out of range.") + raise StopIteration("Waiting for valid indices string.") selected_files = [os.path.join(directory, dir_files[i]) for i in index_list] images = [] masks = [] - has_non_empty_mask = False for path in selected_files: i = Image.open(path) i = ImageOps.exif_transpose(i) @@ -67,12 +61,10 @@ class LoadImagesFromDirByIndexBatch: if 'A' in i.getbands(): mask_np = np.array(i.getchannel('A')).astype(np.float32) / 255.0 mask_tensor = 1. - torch.from_numpy(mask_np) - has_non_empty_mask = True else: - # Default mask, matching previous logic mask_tensor = torch.zeros((image_tensor.shape[2], image_tensor.shape[3]), dtype=torch.float32) images.append(image_tensor) masks.append(mask_tensor) image_batch = torch.cat(images, dim=0) if len(images) > 1 else images[0] mask_batch = torch.stack(masks, dim=0) if len(masks) > 1 else masks[0] - return (image_batch, mask_batch, len(images)) + return (image_batch, mask_batch, len(images)) \ No newline at end of file diff --git a/nodes/from_dir_with_index_list.py b/nodes/from_dir_with_index_list.py new file mode 100644 index 0000000..2f06a35 --- /dev/null +++ b/nodes/from_dir_with_index_list.py @@ -0,0 +1,91 @@ +import os +import torch +import numpy as np +from PIL import Image, ImageOps + +""" + Load an image list from directory using an index string to get positions: "0,2,5,...n" +""" + +class LoadImagesFromDirByIndexList: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "directory": ("STRING", {"default": ""}), + }, + "optional": { + "indices": ("STRING", {"default": ""}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING", "INT") + RETURN_NAMES = ("IMAGE", "MASK", "FILE PATH", "METADATA", "INDEX COUNT") + OUTPUT_IS_LIST = (True, True, True, True, True) + FUNCTION = "load_images" + CATEGORY = "image" + + @classmethod + def IS_CHANGED(cls, **kwargs): + return hash(frozenset(kwargs.items())) + + def _parse_indices(self, indices_str, max_len): + indices = [] + for idx in indices_str.split(","): + idx = idx.strip() + if idx.isdigit(): + i = int(idx) + if 0 <= i < max_len: + indices.append(i) + return sorted(set(indices)) + + def load_images(self, directory: str, indices: str = ""): + if not os.path.isdir(directory): + raise FileNotFoundError(f"Directory '{directory}' cannot be found.") + + dir_files = os.listdir(directory) + valid_extensions = ['.jpg', '.jpeg', '.png', '.webp'] + dir_files = [f for f in dir_files if any(f.lower().endswith(ext) for ext in valid_extensions)] + dir_files = sorted(dir_files) + if not dir_files: + raise FileNotFoundError(f"No valid image files in directory '{directory}'.") + + if not indices.strip(): + raise StopIteration("Waiting for valid indices string.") + + index_list = self._parse_indices(indices, len(dir_files)) + if not index_list: + raise StopIteration("Waiting for valid indices string.") + + images = [] + masks = [] + file_paths = [] + metadatas = [] + indexes = [] + + for idx in index_list: + image_path = os.path.join(directory, dir_files[idx]) + i = Image.open(image_path) + i = ImageOps.exif_transpose(i) + metadata = i.info.copy() + exif = i.getexif() + if exif: + metadata['exif'] = dict(exif) + metadatas.append(str(metadata)) + + image = i.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + else: + mask = torch.zeros((image.shape[2], image.shape[3]), dtype=torch.float32) + + images.append(image) + masks.append(mask) + file_paths.append(str(image_path)) + indexes.append(idx) + + return (images, masks, file_paths, metadatas, indexes) \ No newline at end of file diff --git a/nodes/image_to_true.py b/nodes/image_to_true.py new file mode 100644 index 0000000..5fb03a3 --- /dev/null +++ b/nodes/image_to_true.py @@ -0,0 +1,19 @@ +""" + This node simply outputs true when receives an image, used for automation +""" + +class ImageToTrue: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", {}), + } + } + + RETURN_TYPES = ("BOOLEAN",) + FUNCTION = "process" + + def process(self, image): + return (True,) +