diff --git a/README.md b/README.md index fef8e12..a3351ee 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,8 @@ Clone this repo into `custom_nodes` directory of ComfyUI location ## Usage +There are two nodes, one is `ImageSelector` under `image` category, another is `LatentSelecotr` under `latent` category. One is images and one for latent images. + Input: a list of selected indexes, start with 1 (not 0, sorry), seperated by comma. For example: diff --git a/__init__.py b/__init__.py index 9ba12e4..2f132f0 100644 --- a/__init__.py +++ b/__init__.py @@ -32,7 +32,7 @@ class ImageSelector: return { "required": { "images": ("IMAGE", ), - "selected_images": ("STRING", { + "selected_indexes": ("STRING", { "multiline": False, "default": "1,2,3" }), @@ -48,9 +48,9 @@ class ImageSelector: CATEGORY = "image" - def run(self, images: tg.Sequence[tg.Any], selected_images: tg.Text): + def run(self, images: tg.Sequence[tg.Any], selected_indexes: tg.Text): res_images: tg.List[tg.Any] = [] - for s in selected_images.strip().split(','): + for s in selected_indexes.strip().split(','): try: x: int = int(s.strip()) - 1 if x < len(images): @@ -64,6 +64,63 @@ class ImageSelector: return (images, ) +class LatentSelector: + """ + Select some of the latent images and pipe through + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + """ + Input: list of index of selected image, seperated by comma + Indexes start with 1 for simplicity + """ + return { + "required": { + "latent_image": ("LATENT", ), + "selected_indexes": ("STRING", { + "multiline": False, + "default": "1,2,3" + }), + }, + } + + RETURN_TYPES = ("LATENT", ) + #RETURN_NAMES = ("image_output_name",) + + FUNCTION = "run" + + OUTPUT_NODE = False + + CATEGORY = "latent" + + def run(self, latent_image: tg.Sequence[tg.Any], + selected_indexes: tg.Text): + samples = latent_image['samples'] + shape = samples.shape + len_first_dim = shape[0] + + selected_index: tg.List[int] = [] + for s in selected_indexes.strip().split(','): + try: + x: int = int(s.strip()) - 1 + if x < len_first_dim: + selected_index.append(x) + except: + pass + + if selected_index: + return ({'samples': samples[selected_index, :, :, :]}, ) + + return (latent_image, ) + + # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique -NODE_CLASS_MAPPINGS = {"ImageSelector": ImageSelector} +NODE_CLASS_MAPPINGS = { + "ImageSelector": ImageSelector, + "LatentSelector": LatentSelector +}