A for loop for ComfyUI
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}


def register_node(c):
    assert not isinstance(c.RETURN_TYPES, str), "Error: string found instead of tuple."
    assert not isinstance(c.RETURN_NAMES, str), "Error: string found instead of tuple."
    NODE_CLASS_MAPPINGS[c.__name__] = c
    NODE_DISPLAY_NAME_MAPPINGS[c.__name__] = c.__name__
    return c


from . import sequence, paths, job, image, debug




__init__.py
ali1234/comfyui-job-iterator
import sys
import os
import textwrap
import code

from . import register_node


@register_node
class Stringify:
    """Convert any input to str/repr."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "x": ("*", ),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("str", "repr")
    FUNCTION = "go"
    CATEGORY = "ali1234/debug"

    def go(self, x):
        return (str(x), repr(x))


class RestoreStdStreams(object):
    # ComfyUI-Manager patches sys.stdout and sys.stder
    # which breaks GNU Readline support and makes the
    # REPL annoying to use. This context manager temporarily
    # puts back the originals.

    def __enter__(self):
        self._stdout, self._stderr = sys.stdout, sys.stderr
        sys.stdout, sys.stderr = sys.__stdout__, sys.__stderr__
        self._stdout.flush(); self._stderr.flush()

    def __exit__(self, exc_type, exc_value, traceback):
        _stdout, _stderr = sys.stdout, sys.stderr
        sys.stdout, sys.stderr = self._stdout, self._stderr
        _stdout.flush(); _stderr.flush()


EOF = 'Ctrl-Z plus Return' if os.sep == '\\' else 'Ctrl-D (i.e. EOF)'
MESSAGE = f"Press {EOF} to continue processing."


class Quitter:
    # Replace exit() and quit() to make sure the user exits with EOF
    # otherwise the whole server gets terminated.

    def __repr__(self):
        return MESSAGE

    def __call__(self):
        print(MESSAGE)


@register_node
class Interact:
    """Opens an interactive REPL whenever the node is evaluated."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {},
            "optional": {x: "*" for x in ('a', 'b', 'c', 'd')},
        }

    RETURN_TYPES = ()
    RETURN_NAMES = ()
    FUNCTION = "interact"
    CATEGORY = "ali1234/debug"
    OUTPUT_NODE = True

    def interact(self, **kwargs):
        if sys.__stdout__.isatty():
            with RestoreStdStreams():
                code.interact(
                    banner=textwrap.dedent(f"""
                        Interactive debugging started.
                        Try `print(a)`.
                        {MESSAGE}
                    """),
                    exitmsg="Resuming workflow...",
                    local={
                        **kwargs,
                        'quit': Quitter(),
                        'exit': Quitter()
                    }
                )
        else:
            # Don't block the server if there is no tty.
            print("Skipping interactive prompt because there is no tty.")
        return ()

debug.py
ali1234/comfyui-job-iterator
import torch
import numpy as np

from . import register_node


@register_node
class JoinImageBatch:
    """Turns an image batch into one big image."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "mode": (("horizontal", "vertical"), {"default": "horizontal"}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "join"
    CATEGORY = "ali1234/image"

    def join(self, images, mode):
        n, h, w, c = images.shape
        image = None
        if mode == "vertical":
            # for vertical we can just reshape
            image = images.reshape(1, n * h, w, c)
        elif mode == "horizontal":
            # for horizontal we have to swap axes
            image = torch.transpose(torch.transpose(images, 1, 2).reshape(1, n * w, h, c), 1, 2)
        return (image,)


@register_node
class JoinImages:
    """Turns joins two images into one big image."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image_a": ("IMAGE",),
                "image_b": ("IMAGE",),
                "mode": (("horizontal", "vertical"), {"default": "horizontal"}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "join"
    CATEGORY = "ali1234/image"

    def join(self, image_a, image_b, mode):
        dim = {'horizontal': 2, 'vertical': 1}[mode]
        return (torch.concat((image_a, image_b), dim), )


@register_node
class SelectImageBatch:
    """Selects one image from an image batch."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "select": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "select"
    CATEGORY = "ali1234/image"

    def select(self, images, select):
        n, h, w, c = images.shape
        if select >= n:
            select = n - 1
        return (images[select].reshape(1, h, w, c),)


@register_node
class GetImageSize:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("INT", "INT", "INT")
    RETURN_NAMES = ("width", "height", "batch_size")
    FUNCTION = "go"
    CATEGORY = "ali1234/image"

    def go(self, images):
        return (images.shape[2], images.shape[1], images.shape[0])





@register_node
class StringToImage:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {"default": "Hello world!"}),
                "width": ("INT", {"default": 384}),
                "height": ("INT", {"default": 16}),
                "colour": ("COLOR", {"default": "white"}),
                "background": ("COLOR", {"default": "black"}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "render"
    CATEGORY = "ali1234/image"

    def render(self, text, width, height, colour, background):
        from PIL import Image, ImageDraw, ImageFont
        font = ImageFont.load_default()
        img = Image.new("RGB", (width, height), background)
        draw = ImageDraw.Draw(img)
        _, _, w, h = draw.textbbox((0, 0), text, font=font)
        draw.text(((width - w) / 2, ((height - h) / 2) - 1), text, font=font, fill=colour)
        tensor = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
        return (tensor,)


@register_node
class ProgressBar:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "progress": ("FLOAT", {'default': 0, 'forceInput': True}),
                "padding": ("INT", {'default': 3}),
                "width": ("INT", {"default": 384}),
                "height": ("INT", {"default": 16}),
                "colour": ("COLOR", {"default": "white"}),
                "background": ("COLOR", {"default": "black"}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "render"
    CATEGORY = "ali1234/image"

    def render(self, progress, padding, width, height, colour, background):
        from PIL import Image, ImageDraw
        img = Image.new("RGB", (width, height), background)
        draw = ImageDraw.Draw(img)
        draw.rectangle((padding, padding, width - padding - 1, height - padding - 1), outline=colour)
        if progress > 0:
            ip = padding + 2
            draw.rectangle((ip, ip, max(ip+1, (width - ip - 1) * progress), height - ip - 1), outline=colour, fill=colour)
        tensor = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
        return (tensor,)

image.py
ali1234/comfyui-job-iterator
import collections
import itertools

from execution import PromptExecutor

from . import register_node


@register_node
class MakeJob:
    """Turns a sequence into a job with one attribute."""

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "sequence": ("SEQUENCE", ),
                "name": ("STRING", {"default": ''}),
            },
        }

    RETURN_TYPES = ("JOB", "INT")
    RETURN_NAMES = ("job", "count")
    FUNCTION = "go"
    CATEGORY = "ali1234/job"

    def merge_dicts(self, *dicts):
        #return collections.ChainMap(*reversed(dicts))
        return dict(itertools.chain.from_iterable(d.items() for d in dicts))

    def go(self, sequence, name):
        result = [{name: value} for value in sequence]
        return (result, len(result))


@register_node
class CombineJobs(MakeJob):
    """Combines multiple jobs."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "a": ("JOB", ),
                "method": (("zip", "product"), {"default": "zip"}),
            },
            "optional": {
                x: ("JOB", ) for x in ('b', 'c', 'd', 'e')
            }
        }

    def go(self, method, **kwargs):
        method = {'product': itertools.product, 'zip': zip}[method]
        result = [self.merge_dicts(*steps) for steps in method(*kwargs.values())]
        return (result, len(result))


@register_node
class EnumerateJob(MakeJob):
    """Combines multiple jobs."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "job": ("JOB", ),
                "name": ("STRING", {"default": ''}),
            },
        }

    def go(self, job, name):
        result = [self.merge_dicts(step, {name: n}) for n, step in enumerate(job)]
        return (result, len(result))


@register_node
class GetJobStep:
    """Gets the job step by number."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "job": ("JOB", ),
                "step": ("INT", {"default": 0}),
                "wrap": (("repeat", "clamp"), {"default": "repeat"}),
            },
        }
    RETURN_TYPES = ("ATTRIBUTES", )
    RETURN_NAMES = ("attributes", )
    FUNCTION = "go"
    CATEGORY = "ali1234/job"

    def go(self, job, step, wrap):
        if wrap == 'repeat':
            while step < 0:
                step += len(job)
            step = step % len(job)
        elif wrap == 'clamp':
            step = max(min(step, len(job)), 0)
        return (job[step], )


@register_node
class FormatAttributes:
    """Applies attributes to a format string."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "attributes": ("ATTRIBUTES",),
                "format": ("STRING", {'default': '', 'multiline': True, "dynamicPrompts": False})
            },
        }

    RETURN_TYPES = ("STRING", )
    RETURN_NAMES = ("string", )
    FUNCTION = "go"
    CATEGORY = "ali1234/job"

    def go(self, attributes, format):
        return (format.format(**attributes), )


class AnyType(str):
    def __ne__(self, __value: object) -> bool:
        return False


@register_node
class GetAttribute:
    """Gets a named attribute from a step."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "attributes": ("ATTRIBUTES", ),
                "name": ("STRING", {"default": ''}),
            },
        }

    RETURN_TYPES = (AnyType("*"), )
    RETURN_NAMES = ("value", )
    FUNCTION = "go"
    CATEGORY = "ali1234/job"

    def go(self, attributes, name):
        return (attributes[name], )


# Dynamically register typed attribute getters to avoid wildcard bug
# https://github.com/comfyanonymous/ComfyUI/pull/770
for t in ('INT', 'FLOAT', 'STRING'):
    register_node(type(
        'GetAttribute'+t.title(),
        (GetAttribute, ),
        {
            'RETURN_TYPES': (t, ),
        }
    ))




@register_node
class JobIterator:
    """Magic node that runs the workflow multiple times until all steps are done."""
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "job": ("JOB",),
                "start_step": ("INT", {"default": 0}),
            },
        }

    RETURN_TYPES = ("ATTRIBUTES", "INT", "INT")
    RETURN_NAMES = ("attributes", "count", "step")
    FUNCTION = "go"
    CATEGORY = "ali1234/job"

    def go(self, job, start_step):
        print(f'JobIterator: {start_step + 1} / {len(job)}')
        return (job[start_step], len(job), start_step)


# finally monkey patch the prompt executor to handle batch prompts.

orig_execute = PromptExecutor.execute


def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
    print("Prompt executor has been patched by Job Iterator!")
    orig_execute(self, prompt, prompt_id, extra_data, execute_outputs)

    job_iterator = None
    for k, v in prompt.items():
        try:
            if v['class_type'] == 'JobIterator':
                job_iterator = k
                break
        except KeyError:
            continue

    if job_iterator is not None:
        steps = self.outputs[job_iterator][1][0]
        while prompt[job_iterator]['inputs']['start_step'] < (steps - 1):
            prompt[job_iterator]['inputs']['start_step'] += 1
            orig_execute(self, prompt, prompt_id, extra_data, execute_outputs)


PromptExecutor.execute = execute

job.py
ali1234/comfyui-job-iterator
import pathlib
from folder_paths import folder_names_and_paths
from . import register_node


folder_types = tuple(sorted(folder_names_and_paths.keys()))
folder_default = "checkpoint" if "checkpoint" in folder_types else folder_types[0]

@register_node
class ModelFinder:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "filenames": ("SEQUENCE", ),
                "model_type": (folder_types, {"default": folder_default}),
                "recursive": ("BOOLEAN", {"default": True}),
                "skip_missing": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ("SEQUENCE", "SEQUENCE")
    RETURN_NAMES = ("paths", "stems")
    FUNCTION = "go"
    CATEGORY = "ali1234/sequence"

    def find_models(self, fn, paths, recursive):
        for p in paths:
            for f in (p.rglob(fn) if recursive else p.glob(fn)):
                yield f.relative_to(p)

    def go(self, filenames, model_type, recursive, skip_missing):
        paths = [pathlib.Path(folder) for folder in folder_names_and_paths[model_type][0]]
        result = []
        for fn in filenames:
            if '..' in fn:
                raise Exception(f'".." is not allowed: {fn}.')
            found = False
            for f in self.find_models(fn, paths, recursive):
                found = True
                result.append(f)
            if not found and not skip_missing:
                raise Exception(f'Could not find file: {fn}')

        return ([str(f) for f in result], [f.stem for f in result])

paths.py
ali1234/comfyui-job-iterator
import ast
import functools
import itertools

from . import register_node


# Sequence creation.


@register_node
class Range:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "start": ("INT", {"default": 0, "min": -9999999, "max": 9999999, "step": 1}),
                "stop": ("INT", {"default": 10, "min": -9999999, "max": 9999999, "step": 1}),
                "step": ("INT", {"default": 1, "min": -9999999, "max": 9999999, "step": 1}),
            },
        }

    RETURN_TYPES = ("SEQUENCE", )
    RETURN_NAMES = ("sequence", )
    FUNCTION = "go"
    CATEGORY = "ali1234/sequence"

    def go(self, start, stop, step):
        return (range(start, stop, step), )


@register_node
class Literal:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "literal": ("STRING", {'default': '', 'multiline': True})
            },
        }

    RETURN_TYPES = ("SEQUENCE", )
    RETURN_NAMES = ("sequence", )
    FUNCTION = "go"
    CATEGORY = "ali1234/sequence"

    def go(self, literal):
        return (ast.literal_eval(literal), )


# Processing of existing sequences.


@register_node
class Reorder:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "sequence": ("SEQUENCE", ),
                "method": (("reverse", "sort", "reverse sort"), {"default": "sort"}),
            },
        }

    RETURN_TYPES = ("SEQUENCE", )
    RETURN_NAMES = ("sequence", )
    FUNCTION = "go"
    CATEGORY = "ali1234/sequence"

    def go(self, sequence, method):
        f = {"reverse": reversed, "sort": sorted, "reverse_sort": lambda x: sorted(x, reverse=True)}[method]
        return (f(sequence), )


@register_node
class Combinations:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "sequence": ("SEQUENCE", ),
                "count": ("INT", {"default": 0, "min": 0, "max": 9999999, "step": 1}),
                "replacement": ("BOOLEAN", {"default": False, "label_on": "Yes", "label_off": "No"}),
            },
        }

    RETURN_TYPES = ("SEQUENCE", )
    RETURN_NAMES = ("sequence", )
    FUNCTION = "go"
    CATEGORY = "ali1234/sequence"

    def go(self, sequence, count, replacement):
        if replacement:
            return (itertools.combinations_with_replacement(sequence, count), )
        else:
            return (itertools.combinations(sequence, count), )


@register_node
class Permutations:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "sequence": ("SEQUENCE", ),
                "count": ("INT", {"default": 0, "min": 0, "max": 9999999, "step": 1}),
            },
        }

    RETURN_TYPES = ("SEQUENCE", )
    RETURN_NAMES = ("sequence", )
    FUNCTION = "go"
    CATEGORY = "ali1234/sequence"

    def go(self, sequence, count):
        return ([x for x in itertools.permutations(sequence, count)], )


@register_node
class Slice:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "sequence": ("SEQUENCE",),
                "start": ("INT", {"default": 0, "min": -9999999, "max": 9999999, "step": 1}),
                "stop": ("INT", {"default": 9999999, "min": -9999999, "max": 9999999, "step": 1}),
                "step": ("INT", {"default": 1, "min": -9999999, "max": 9999999, "step": 1}),
            },
        }

    RETURN_TYPES = ("SEQUENCE", )
    RETURN_NAMES = ("sequence", )
    FUNCTION = "go"
    CATEGORY = "ali1234/sequence"

    def go(self, sequence, start, stop, step):
        return (sequence[start:stop:step], )



@register_node
class Join:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "sequence": ("SEQUENCE", ),
                "join_str": ("STRING", {'default': ', '}),
            },
        }

    RETURN_TYPES = ("SEQUENCE", )
    RETURN_NAMES = ("sequence", )
    FUNCTION = "go"
    CATEGORY = "ali1234/sequence"

    def go(self, sequence, join_str):
        return ([join_str.join(x) for x in sequence], )

sequence.py
ali1234/comfyui-job-iterator