WIP rewrite branch for new ComfyUI features.
This commit is contained in:
Alistair Buxton
@@ -1 +1,2 @@
|
||||
__pycache__
|
||||
.idea
|
||||
|
||||
+11
-10
@@ -1,16 +1,17 @@
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
|
||||
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 .registry import set_pack_options, get_nodes
|
||||
|
||||
set_pack_options('jobiter', 'Job Iterator')
|
||||
|
||||
from . import sequence, paths, job, image, debug
|
||||
|
||||
from . import nodes
|
||||
|
||||
class JobIterator(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return get_nodes()
|
||||
|
||||
async def comfy_entrypoint() -> JobIterator: # ComfyUI calls this to load your extension and its nodes.
|
||||
return JobIterator()
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 685 KiB After Width: | Height: | Size: 1.1 MiB |
@@ -1,180 +0,0 @@
|
||||
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 JobToList:
|
||||
"""Converts a job into a list."""
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"job": ("JOB",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("ATTRIBUTES",)
|
||||
RETURN_NAMES = ("attributes",)
|
||||
OUTPUT_IS_LIST = (True,)
|
||||
FUNCTION = "go"
|
||||
CATEGORY = "ali1234/job"
|
||||
|
||||
def go(self, job):
|
||||
return (job,)
|
||||
@@ -0,0 +1 @@
|
||||
from . import sequence, job, loop
|
||||
@@ -0,0 +1,80 @@
|
||||
import math
|
||||
|
||||
from ..registry import register_node
|
||||
from ..types import *
|
||||
|
||||
|
||||
class NamedSeq:
|
||||
def nth(self, n):
|
||||
raise NotImplementedError
|
||||
|
||||
def __getitem__(self, n):
|
||||
return dict(self.nth(n))
|
||||
|
||||
|
||||
class SingleSeq(NamedSeq):
|
||||
def __init__(self, name, seq):
|
||||
self._name = name
|
||||
self._seq = seq
|
||||
|
||||
def nth(self, n):
|
||||
n = n % len(self)
|
||||
yield (self._name, self._seq[n])
|
||||
|
||||
def __len__(self):
|
||||
return len(self._seq)
|
||||
|
||||
|
||||
class CombinedSeq(NamedSeq):
|
||||
def __init__(self, *args):
|
||||
self._seqs = args
|
||||
|
||||
|
||||
class ProductSeq(CombinedSeq):
|
||||
def nth(self, n):
|
||||
n = n % len(self)
|
||||
indices = []
|
||||
for s in reversed(self._seqs):
|
||||
indices.append(n % len(s))
|
||||
n //= len(s)
|
||||
|
||||
for i, s in zip(reversed(indices), self._seqs):
|
||||
yield from s.nth(i)
|
||||
|
||||
def __len__(self):
|
||||
return math.prod(len(s) for s in self._seqs)
|
||||
|
||||
|
||||
class ZipSeq(CombinedSeq):
|
||||
def nth(self, n):
|
||||
n = n % len(self)
|
||||
for s in self._seqs:
|
||||
yield from s.nth(n)
|
||||
|
||||
def __len__(self):
|
||||
return min(len(s) for s in self._seqs)
|
||||
|
||||
|
||||
@register_node(display_name="Make Job")
|
||||
def MakeJob(sequence: Sequence(), name: String() = "") -> (Job(), ):
|
||||
return SingleSeq(name, sequence),
|
||||
|
||||
|
||||
combine_modes = {
|
||||
'zip': ZipSeq,
|
||||
'product': ProductSeq,
|
||||
}
|
||||
|
||||
@register_node(display_name="Combine Jobs")
|
||||
def CombineJobs(jobs: Autogrow(input=Job(), prefix='job', min=1), method: Combo(choices=combine_modes) = 'product') -> (Job(), ):
|
||||
return method(*jobs),
|
||||
|
||||
|
||||
@register_node(display_name="Enumerate Job")
|
||||
def EnumerateJob(job: Job(), name: String() = "") -> (Job(), ):
|
||||
return ZipSeq(SingleSeq(name, range(len(job))), job),
|
||||
|
||||
|
||||
@register_node(display_name="Get Value From Mapping")
|
||||
def GetValueFromMapping(mapping: Mapping(), name: String() = "") -> (Any(), ):
|
||||
return mapping[name],
|
||||
+216
@@ -0,0 +1,216 @@
|
||||
"""
|
||||
Generic sequence looping via ComfyUI graph expansion (expand), supporting:
|
||||
- multiple independent loops in one workflow
|
||||
- nested loops
|
||||
- accumulation driven ONLY by LoopEnd.value (always collects; if value unconnected -> None)
|
||||
- strict "loop region" definition by forward reachability from the LoopStart
|
||||
- validation: no edges may "escape" the loop body to outside nodes except from
|
||||
the LoopEnd node itself (enables safe nesting and composition)
|
||||
|
||||
UX / wiring:
|
||||
- Users wire the upstream sequence ONLY into LoopStart.
|
||||
- Between LoopStart and LoopEnd there are only TWO wires:
|
||||
1) flow (raw_link pairing, used for region discovery / nesting)
|
||||
2) data (loop-carried state dict: index, accum, sequence)
|
||||
|
||||
Important implementation detail:
|
||||
- The expanded clone must NOT rely on preserving external links as raw tuples.
|
||||
Therefore LoopStart stores the upstream sequence in `data["seq"]`, and LoopEnd
|
||||
reads it from there. This guarantees the recursive LoopEnd sees the real
|
||||
sequence via an internal connection (Start.data -> End.data) inside the loop.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple, TypedDict
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
from comfy_execution.graph_utils import GraphBuilder, is_link
|
||||
|
||||
from ..types import * # SequenceType, FlowType, etc.
|
||||
from ..registry import register_node_class, make_unique_name, make_category
|
||||
|
||||
|
||||
def _build_children_adjacency(prompt: Dict[str, Any]) -> Dict[str, List[str]]:
|
||||
"""Build parent->children adjacency from prompt's input links."""
|
||||
children: Dict[str, List[str]] = {}
|
||||
for child_id, node in prompt.items():
|
||||
for _, v in node.get("inputs", {}).items():
|
||||
if is_link(v):
|
||||
parent_id = str(v[0])
|
||||
children.setdefault(parent_id, []).append(str(child_id))
|
||||
return children
|
||||
|
||||
|
||||
def _reachable_region(children: Dict[str, List[str]], start_id: str, end_id: str) -> Set[str]:
|
||||
"""Forward reachability from start_id, stopping recursion at end_id."""
|
||||
inside: Set[str] = set()
|
||||
stack: List[str] = [str(start_id)]
|
||||
while stack:
|
||||
nid = stack.pop()
|
||||
if nid in inside:
|
||||
continue
|
||||
inside.add(nid)
|
||||
if nid == str(end_id):
|
||||
continue
|
||||
for c in children.get(nid, []):
|
||||
stack.append(c)
|
||||
return inside
|
||||
|
||||
|
||||
def _find_escape_edges(prompt: Dict[str, Any], inside: Set[str], end_id: str) -> List[Tuple[str, str, str]]:
|
||||
"""
|
||||
Illegal edges are any inside->outside link, excluding those originating from end_id.
|
||||
Each entry: (parent_id, child_id, child_input_name)
|
||||
"""
|
||||
escapes: List[Tuple[str, str, str]] = []
|
||||
end_id = str(end_id)
|
||||
|
||||
for child_id, node in prompt.items():
|
||||
child_id = str(child_id)
|
||||
for input_name, v in node.get("inputs", {}).items():
|
||||
if not is_link(v):
|
||||
continue
|
||||
parent_id = str(v[0])
|
||||
if parent_id in inside and child_id not in inside and parent_id != end_id:
|
||||
escapes.append((parent_id, child_id, str(input_name)))
|
||||
|
||||
return escapes
|
||||
|
||||
|
||||
def _format_escape_error(escapes: List[Tuple[str, str, str]], start_id: str, end_id: str) -> str:
|
||||
lines = [
|
||||
"Loop body has illegal connections escaping to nodes outside the loop region.",
|
||||
f"LoopStart id: {start_id}",
|
||||
f"LoopEnd id: {end_id}",
|
||||
"",
|
||||
"Illegal edges (producer -> consumer[input]):",
|
||||
]
|
||||
for parent_id, child_id, input_name in escapes:
|
||||
lines.append(f" - {parent_id} -> {child_id}[{input_name}]")
|
||||
lines += [
|
||||
"",
|
||||
"Fix: route values out of the loop via the LoopEnd node's outputs, or move the consumer node inside the loop.",
|
||||
"Note: outside->inside edges are allowed (loop-invariant dependencies).",
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@register_node_class
|
||||
class LoopStart(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id=make_unique_name(cls.__name__),
|
||||
display_name="Loop Start",
|
||||
category=make_category("Loop"),
|
||||
inputs=[
|
||||
io.MultiType.Input("sequence", types=[SequenceType, JobType]),
|
||||
SequenceType.Input("accum", optional=True),
|
||||
io.Int.Input("first_index", default=0),
|
||||
],
|
||||
outputs=[
|
||||
FlowType.Output("flow"),
|
||||
SequenceType.Output("accum"),
|
||||
io.Int.Output("index"),
|
||||
io.AnyType.Output("item"),
|
||||
],
|
||||
hidden=[io.Hidden.unique_id],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, sequence, accum=None, first_index=0):
|
||||
if accum is None:
|
||||
accum = []
|
||||
flow = {'sequence': sequence, 'index': first_index, 'accum': accum, 'start_id': cls.hidden.unique_id}
|
||||
return io.NodeOutput(flow, accum, first_index, sequence[first_index])
|
||||
|
||||
|
||||
@register_node_class
|
||||
class LoopEnd(io.ComfyNode):
|
||||
DISPLAY_NAME = "Loop End"
|
||||
CATEGORY = None
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id=make_unique_name(cls.__name__),
|
||||
display_name=cls.DISPLAY_NAME,
|
||||
category=make_category("Loop"),
|
||||
enable_expand=True,
|
||||
inputs=[
|
||||
FlowType.Input("flow"),
|
||||
io.AnyType.Input("result", optional=True),
|
||||
io.Combo.Input('mode', options=("append", "replace"), default="append")
|
||||
],
|
||||
outputs=[
|
||||
SequenceType.Output("result"),
|
||||
],
|
||||
hidden=[io.Hidden.dynprompt, io.Hidden.unique_id],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, flow, mode, result=None):
|
||||
flow = dict(flow)
|
||||
if mode == "append":
|
||||
flow['accum'] = flow['accum'] + [result]
|
||||
elif mode == "replace":
|
||||
flow['accum'] = result
|
||||
flow['index'] = flow['index'] + 1
|
||||
|
||||
# Terminate: stop expanding.
|
||||
if flow['index'] >= len(flow['sequence']):
|
||||
return io.NodeOutput(flow['accum'])
|
||||
|
||||
# Pairing: find connected start node id via flow
|
||||
start_id = flow['start_id']
|
||||
end_id = flow.setdefault('end_id', cls.hidden.unique_id)
|
||||
|
||||
# Region discovery/validation on the ORIGINAL prompt graph
|
||||
prompt = flow.setdefault('prompt', cls.hidden.dynprompt.get_original_prompt())
|
||||
children = _build_children_adjacency(prompt)
|
||||
inside = _reachable_region(children, start_id=start_id, end_id=end_id)
|
||||
|
||||
if end_id not in inside:
|
||||
raise RuntimeError(
|
||||
"LoopEnd is not reachable from LoopStart. "
|
||||
"Ensure your loop body connects from LoopStart outputs to LoopEnd inputs."
|
||||
)
|
||||
|
||||
escapes = _find_escape_edges(prompt, inside=inside, end_id=end_id)
|
||||
if escapes:
|
||||
raise ValueError(_format_escape_error(escapes, start_id=start_id, end_id=end_id))
|
||||
|
||||
graph = GraphBuilder()
|
||||
|
||||
for node_id in inside:
|
||||
original_node = cls.hidden.dynprompt.get_node(node_id)
|
||||
class_type = original_node["class_type"]
|
||||
n = graph.node(class_type, node_id)
|
||||
n.set_override_display_id(node_id)
|
||||
|
||||
# Wire inputs (internal links rewired to clones; external links/values kept as-is)
|
||||
for node_id in inside:
|
||||
original_node = cls.hidden.dynprompt.get_node(node_id)
|
||||
node = graph.lookup_node(node_id)
|
||||
|
||||
for k, v in original_node.get("inputs", {}).items():
|
||||
if node_id != start_id:
|
||||
if is_link(v) and v[0] in inside:
|
||||
node.set_input(k, graph.lookup_node(v[0]).out(v[1]))
|
||||
else:
|
||||
node.set_input(k, v)
|
||||
|
||||
new_start = graph.lookup_node(start_id)
|
||||
new_start.set_input("sequence", flow['sequence'])
|
||||
new_start.set_input("first_index", flow['index'])
|
||||
new_start.set_input("accum", flow['accum'])
|
||||
|
||||
new_end = graph.lookup_node(end_id)
|
||||
new_end.set_input("flow", flow)
|
||||
|
||||
return {
|
||||
"result": (new_end.out(0), ), # results, data
|
||||
"expand": graph.finalize(),
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
import ast
|
||||
import itertools
|
||||
import torch
|
||||
|
||||
from ..registry import register_node
|
||||
from ..types import *
|
||||
|
||||
# Sequence creation.
|
||||
|
||||
@register_node(category="Sequence")
|
||||
def Range(start: Int() = 0, stop: Int() = 10, step: Int() = 1) -> (Sequence(), ):
|
||||
return range(start, stop, step),
|
||||
|
||||
|
||||
@register_node(category="Sequence")
|
||||
def Literal(literal: String(multiline = True) = "") -> (Any(), ):
|
||||
return ast.literal_eval(literal),
|
||||
|
||||
|
||||
@register_node(category="Sequence")
|
||||
def Combinatorics(sequence: Sequence(), min: Int() = 3, max: Int() = 4, replacement: Bool() = False, permutations: Bool() = False) -> (Sequence(), ):
|
||||
"""Sequence combinatorics."""
|
||||
comb = itertools.combinations_with_replacement if replacement else itertools.combinations
|
||||
combs = itertools.chain(*(comb(sequence, n) for n in range(min, max+1)))
|
||||
if permutations:
|
||||
combs = itertools.chain(*(itertools.permutations(x) for x in combs))
|
||||
return list(combs),
|
||||
|
||||
|
||||
reorder_modes = {'reverse': reversed, 'sort': sorted}
|
||||
|
||||
@register_node(category="Sequence")
|
||||
def Reorder(sequence: Sequence(), mode: Combo(choices=reorder_modes) = 'reverse') -> (Sequence(), ):
|
||||
return list(mode(sequence)),
|
||||
|
||||
|
||||
format_modes = {
|
||||
'mapping': lambda x, y: x.format(**y),
|
||||
'iterable': lambda x, y: x.format(*y),
|
||||
'single': lambda x, y: x.format(y)
|
||||
}
|
||||
|
||||
@register_node(category="Text")
|
||||
def Format(vars: Any(), string: String(multiline=True) = "", mode: Combo(choices=format_modes) = 'mapping') -> (String(), ):
|
||||
"""Performs string replacement using the standard Python format() method."""
|
||||
return mode(string, vars),
|
||||
|
||||
|
||||
@register_node(category="Text")
|
||||
def Join(sequence: Sequence(), sep: String() = ", ") -> (String(), ):
|
||||
"""Joins a sequence of strings into a single string."""
|
||||
return sep.join(str(x) for x in sequence),
|
||||
|
||||
@register_node(display_name="Join Tensor Sequence", category="Utils")
|
||||
def JoinTensorSequence(sequence: Sequence(), dim: Int(min=0, max=10) = 0) -> (Any(), ):
|
||||
return torch.cat(sequence, dim=dim),
|
||||
|
||||
|
||||
@register_node(display_name="Split Image Batch", category="Utils")
|
||||
def SplitImageBatch(images: Image(), threshold: Float(min=0.001, max=1, step=0.001, round=0.001) = 0.01) -> (Sequence(), ):
|
||||
c = torch.argwhere(images.diff(dim=0).pow(2).mean(dim=(1, 2, 3)) > threshold).flatten() + 1
|
||||
print(c)
|
||||
last = 0
|
||||
output = []
|
||||
for frame in c:
|
||||
print(frame)
|
||||
output.append(images[last:frame])
|
||||
last = frame
|
||||
output.append(images[last:])
|
||||
return output,
|
||||
|
||||
|
||||
@register_node(display_name="Mapped Join", category="Text")
|
||||
def MappedJoin(sequence: Sequence(), sep: String() = ", ") -> (Sequence(), ):
|
||||
"""Joins a sequence of iterables of strings into a sequence of single strings."""
|
||||
t = (sep.join(str(x) for x in s) for s in sequence)
|
||||
return list(t),
|
||||
|
||||
|
||||
@register_node(display_name="Make Sequence")
|
||||
def MakeSequence(items: AutogrowMatch(prefix='item', min=1)) -> (Sequence(), ):
|
||||
"""Turns any inputs into a sequence."""
|
||||
return items,
|
||||
|
||||
|
||||
@register_node(display_name="Get Item From Sequence", category="Sequence")
|
||||
def GetFromSequence(sequence: MultiType(types=(Sequence(), Job())), index: Int() = 0) -> (Any(), ):
|
||||
"""Gets the nth item from a sequence. Also works on jobs."""
|
||||
return sequence[index],
|
||||
+80
@@ -0,0 +1,80 @@
|
||||
import inspect
|
||||
from functools import wraps
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
from .types import Autogrow, AutogrowMatch, Combo
|
||||
|
||||
PACK_BASE_CATEGORY = None
|
||||
PACK_UID = None
|
||||
PACK_NODES = []
|
||||
|
||||
|
||||
def set_pack_options(uid: str, category: str = None):
|
||||
global PACK_BASE_CATEGORY, PACK_UID
|
||||
PACK_BASE_CATEGORY = category
|
||||
PACK_UID = uid
|
||||
|
||||
|
||||
def make_category(category):
|
||||
cat_list = []
|
||||
if PACK_BASE_CATEGORY is not None:
|
||||
cat_list.append(PACK_BASE_CATEGORY)
|
||||
if category is not None:
|
||||
cat_list.append(category)
|
||||
return '/'.join(cat_list)
|
||||
|
||||
|
||||
def make_unique_name(name, version=0):
|
||||
return f'{PACK_UID}_{version}_{name}'
|
||||
|
||||
|
||||
def get_nodes():
|
||||
return PACK_NODES
|
||||
|
||||
|
||||
def register_node_class(cls):
|
||||
PACK_NODES.append(cls)
|
||||
return cls
|
||||
|
||||
|
||||
def register_node(category=None, version=0, display_name=None, output=False):
|
||||
def decorator(f):
|
||||
if PACK_UID is None:
|
||||
raise Exception("PACK_UID is not set. Call set_pack_options in __init__.py to set it.")
|
||||
|
||||
unique_name = make_unique_name(f.__name__, version)
|
||||
|
||||
sig = inspect.signature(f)
|
||||
node_attrs = {}
|
||||
|
||||
final_category = make_category(category)
|
||||
if not final_category:
|
||||
print(f"WARNING: No category specified for {f.__name__} and no base category. It won't be shown in menus.")
|
||||
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
node_id=unique_name,
|
||||
display_name=display_name or f.__name__,
|
||||
category=final_category,
|
||||
inputs=list(v.annotation.Input(k, v.default) for k, v in sig.parameters.items()),
|
||||
outputs=list(v.Output() for v in sig.return_annotation),
|
||||
)
|
||||
|
||||
node_attrs['define_schema'] = classmethod(define_schema)
|
||||
|
||||
def execute(cls, **kwargs) -> io.NodeOutput:
|
||||
for k, v in sig.parameters.items():
|
||||
if isinstance(v.annotation, Combo):
|
||||
kwargs[k] = v.annotation[kwargs[k]]
|
||||
elif isinstance(v.annotation, Autogrow) or isinstance(v.annotation, AutogrowMatch):
|
||||
kwargs[k] = list(kwargs[k].values())
|
||||
return io.NodeOutput(*f(**kwargs))
|
||||
|
||||
node_attrs['execute'] = classmethod(execute)
|
||||
|
||||
node_class = type(unique_name, (io.ComfyNode,), node_attrs)
|
||||
PACK_NODES.append(node_class)
|
||||
return node_class
|
||||
|
||||
return decorator
|
||||
-159
@@ -1,159 +0,0 @@
|
||||
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], )
|
||||
@@ -0,0 +1,171 @@
|
||||
import typing, inspect
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
|
||||
class ComfyType(BaseModel):
|
||||
TYPE: typing.ClassVar
|
||||
display_name: str = None
|
||||
optional: bool = None
|
||||
|
||||
def opts(self):
|
||||
return self.model_dump(mode='python', exclude_none=True)
|
||||
|
||||
def Input(self, name, default=inspect._empty):
|
||||
if default is inspect._empty:
|
||||
return self.TYPE.Input(id=name, **self.opts())
|
||||
else:
|
||||
return self.TYPE.Input(id=name, default=default, **self.opts())
|
||||
|
||||
def Output(self):
|
||||
return self.TYPE.Output()
|
||||
|
||||
class ComfyWidgetType(ComfyType):
|
||||
"""Base type for ComfyUI types that have options controlling how they are displayed."""
|
||||
socketless: bool = None
|
||||
force_input: bool = None
|
||||
|
||||
|
||||
class Int(ComfyWidgetType):
|
||||
TYPE = io.Int
|
||||
min: int = None
|
||||
max: int = None
|
||||
step: int = None
|
||||
display_mode: io.NumberDisplay = None
|
||||
|
||||
|
||||
class Float(ComfyWidgetType):
|
||||
TYPE = io.Float
|
||||
min: float = None
|
||||
max: float = None
|
||||
step: float = None
|
||||
round: float = None
|
||||
display_mode: io.NumberDisplay = None
|
||||
|
||||
|
||||
class String(ComfyWidgetType):
|
||||
TYPE = io.String
|
||||
multiline: bool = False
|
||||
|
||||
|
||||
class Bool(ComfyWidgetType):
|
||||
TYPE = io.Boolean
|
||||
label_on: str = None
|
||||
label_off: str = None
|
||||
|
||||
|
||||
class Image(ComfyWidgetType):
|
||||
TYPE = io.Image
|
||||
|
||||
|
||||
class Video(ComfyWidgetType):
|
||||
TYPE = io.Video
|
||||
|
||||
|
||||
class Combo(ComfyWidgetType):
|
||||
TYPE = io.Combo
|
||||
choices: typing.Mapping[str, typing.Any]
|
||||
|
||||
def opts(self):
|
||||
return {
|
||||
'options': list(self.choices.keys()),
|
||||
**self.model_dump(mode='python', exclude_none=True, exclude={'choices'})
|
||||
}
|
||||
|
||||
def Input(self, name, default=inspect._empty):
|
||||
if default is inspect._empty:
|
||||
return io.MultiType.Input(id=self.TYPE.Input(id=name, **self.opts()), types=(self.TYPE, io.String))
|
||||
else:
|
||||
return io.MultiType.Input(id=self.TYPE.Input(id=name, default=default, **self.opts()), types=(self.TYPE, io.String))
|
||||
|
||||
def Output(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def __getitem__(self, item):
|
||||
return self.choices[item]
|
||||
|
||||
|
||||
class Any(ComfyType):
|
||||
TYPE = io.AnyType
|
||||
|
||||
|
||||
SequenceType = io.Custom('SEQUENCE')
|
||||
MappingType = io.Custom('MAPPING')
|
||||
JobType = io.Custom('JOB')
|
||||
FlowType = io.Custom('JOB_FLOW')
|
||||
|
||||
|
||||
class Sequence(ComfyType):
|
||||
TYPE = SequenceType
|
||||
|
||||
|
||||
class Mapping(ComfyType):
|
||||
TYPE = MappingType
|
||||
|
||||
|
||||
class Job(ComfyType):
|
||||
TYPE = JobType
|
||||
|
||||
|
||||
class Autogrow(ComfyType):
|
||||
TYPE = io.Autogrow
|
||||
input: ComfyType = Any()
|
||||
prefix: str = None
|
||||
min: int = None
|
||||
max: int = None
|
||||
|
||||
def opts(self):
|
||||
return self.model_dump(mode='python', exclude_none=True, exclude={'input', 'prefix', 'min', 'max', 'match'})
|
||||
|
||||
def autoopts(self):
|
||||
return self.model_dump(mode='python', exclude_none=True, include={'prefix', 'min', 'max'})
|
||||
|
||||
def Input(self, name, default):
|
||||
template = io.Autogrow.TemplatePrefix(self.input.Input(self.prefix), **self.autoopts())
|
||||
return self.TYPE.Input(name, template=template, **self.opts())
|
||||
|
||||
def Output(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class AutogrowMatch(ComfyType):
|
||||
TYPE = io.Autogrow
|
||||
prefix: str = None
|
||||
types: tuple = Any()
|
||||
min: int = None
|
||||
max: int = None
|
||||
|
||||
def opts(self):
|
||||
return self.model_dump(mode='python', exclude_none=True, exclude={'prefix', 'min', 'max', 'types'})
|
||||
|
||||
def autoopts(self):
|
||||
return self.model_dump(mode='python', exclude_none=True, include={'prefix', 'min', 'max'})
|
||||
|
||||
def Input(self, name, default):
|
||||
match_template = template_matchtype = io.MatchType.Template("type")
|
||||
template = io.Autogrow.TemplatePrefix(io.MatchType.Input(self.prefix, template=match_template), **self.autoopts())
|
||||
return self.TYPE.Input(name, template=template, **self.opts())
|
||||
|
||||
def Output(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class List(ComfyType):
|
||||
type: ComfyType = None
|
||||
|
||||
def Output(self):
|
||||
return self.type.TYPE.Output(is_output_list=True)
|
||||
|
||||
|
||||
class MultiType(ComfyType):
|
||||
TYPE = io.MultiType
|
||||
types: tuple = None
|
||||
|
||||
def opts(self):
|
||||
return self.model_dump(mode='python', exclude_none=True, exclude={'types'})
|
||||
|
||||
def Input(self, name, default):
|
||||
return self.TYPE.Input(name, types=tuple(t.TYPE for t in self.types), **self.opts())
|
||||
|
||||
Reference in New Issue
Block a user