diff --git a/comfy_api/v3/io.py b/comfy_api/v3/io.py index 4cc758562..78be29321 100644 --- a/comfy_api/v3/io.py +++ b/comfy_api/v3/io.py @@ -1,10 +1,12 @@ from __future__ import annotations -from typing import Any, Literal, TYPE_CHECKING, TypeVar, Callable, Optional, cast, TypedDict, NotRequired +from typing import Any, Literal, TYPE_CHECKING, TypeVar, Callable, Optional, cast, TypedDict +from typing_extensions import NotRequired from enum import Enum from abc import ABC, abstractmethod from dataclasses import dataclass, asdict from collections import Counter from comfy_api.v3.resources import Resources, ResourcesLocal +import copy # used for type hinting import torch from spandrel import ImageModelDescriptor @@ -189,17 +191,19 @@ class WidgetInputV3(InputV3): ''' def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, default: Any=None, - socketless: bool=None, widgetType: str=None, extra_dict=None): + socketless: bool=None, widgetType: str=None, extra_dict=None, force_input: bool=None): super().__init__(id, display_name, optional, tooltip, lazy, extra_dict) self.default = default self.socketless = socketless self.widgetType = widgetType + self.force_input = force_input def as_dict_V1(self): return super().as_dict_V1() | prune_dict({ "default": self.default, "socketless": self.socketless, "widgetType": self.widgetType, + "forceInput": self.force_input, }) def get_io_type_V1(self): @@ -754,44 +758,69 @@ class MultiType: else: return super().as_dict_V1() -class DynamicInput(InputV3): +class DynamicInput(InputV3, ABC): ''' Abstract class for dynamic input registration. ''' - def __init__(self, io_type: str, id: str, display_name: str=None): - super().__init__(io_type, id, display_name) + @abstractmethod + def get_dynamic(self) -> list[InputV3]: + ... -class DynamicOutput(OutputV3): +class DynamicOutput(OutputV3, ABC): ''' Abstract class for dynamic output registration. ''' - def __init__(self, io_type: str, id: str, display_name: str=None): - super().__init__(io_type, id, display_name) + @abstractmethod + def get_dynamic(self) -> list[OutputV3]: + ... -# io_type="COMFY_MULTIGROW_V3" -class AutoGrowDynamicInput(DynamicInput): - ''' - Dynamic Input that adds another template_input each time one is provided. - Additional inputs are forced to have 'optional=True'. - ''' - def __init__(self, id: str, template_input: InputV3, min: int=1, max: int=None): - super().__init__("AutoGrowDynamicInput", id) - self.template_input = template_input - if min is not None: - assert(min >= 1) - if max is not None: - assert(max >= 1) - self.min = min - self.max = max +@comfytype(io_type="COMFY_AUTOGROW_V3") +class AutogrowDynamic: + Type = list[Any] + class Input(DynamicInput): + def __init__(self, id: str, template_input: InputV3, min: int=1, max: int=None, + display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None): + super().__init__(id, display_name, optional, tooltip, lazy, extra_dict) + self.template_input = template_input + if min is not None: + assert(min >= 1) + if max is not None: + assert(max >= 1) + self.min = min + self.max = max + + def get_dynamic(self) -> list[InputV3]: + curr_count = 1 + new_inputs = [] + for i in range(self.min): + new_input = copy.copy(self.template_input) + new_input.id = f"{new_input.id}{curr_count}_${self.id}_ag$" + if new_input.display_name is not None: + new_input.display_name = f"{new_input.display_name}{curr_count}" + new_input.optional = self.optional or new_input.optional + if isinstance(self.template_input, WidgetInputV3): + new_input.force_input = True + new_inputs.append(new_input) + curr_count += 1 + # pretend to expand up to max + for i in range(curr_count-1, self.max): + new_input = copy.copy(self.template_input) + new_input.id = f"{new_input.id}{curr_count}_${self.id}_ag$" + if new_input.display_name is not None: + new_input.display_name = f"{new_input.display_name}{curr_count}" + new_input.optional = True + if isinstance(self.template_input, WidgetInputV3): + new_input.force_input = True + new_inputs.append(new_input) + curr_count += 1 + return new_inputs # io_type="COMFY_COMBODYNAMIC_V3" class ComboDynamicInput(DynamicInput): def __init__(self, id: str): pass -AutoGrowDynamicInput(id="dynamic", template_input=Image.Input(id="image")) - class HiddenHolder: def __init__(self, unique_id: str, prompt: Any, @@ -960,6 +989,11 @@ class classproperty(object): return self.f(owner) +def add_to_dict_v1(i: InputV3, input: dict): + key = "optional" if i.optional else "required" + input.setdefault(key, {})[i.id] = (i.get_io_type_V1(), i.as_dict_V1()) + + class ComfyNodeV3: """Common base class for all V3 nodes.""" @@ -1117,8 +1151,12 @@ class ComfyNodeV3: } if schema.inputs: for i in schema.inputs: - key = "optional" if i.optional else "required" - input.setdefault(key, {})[i.id] = (i.get_io_type_V1(), i.as_dict_V1()) + if isinstance(i, DynamicInput): + dynamic_inputs = i.get_dynamic() + for d in dynamic_inputs: + add_to_dict_v1(d, input) + else: + add_to_dict_v1(i, input) if schema.hidden and include_hidden: for hidden in schema.hidden: input.setdefault("hidden", {})[hidden.name] = (hidden.value,) diff --git a/comfy_api/v3/resources.py b/comfy_api/v3/resources.py index 6ff59d6ae..0c504e73f 100644 --- a/comfy_api/v3/resources.py +++ b/comfy_api/v3/resources.py @@ -41,7 +41,7 @@ class ResourcesLocal(Resources): def __init__(self): super().__init__() self.local_resources: dict[ResourceKey, Any] = {} - + def get(self, key: ResourceKey, default: Any=...) -> Any: cached = self.local_resources.get(key, None) if cached is not None: diff --git a/comfy_extras/nodes_v3_test.py b/comfy_extras/nodes_v3_test.py index 9120d8b8c..ba7f88ee7 100644 --- a/comfy_extras/nodes_v3_test.py +++ b/comfy_extras/nodes_v3_test.py @@ -149,7 +149,37 @@ class V3LoraLoader(io.ComfyNodeV3): return io.NodeOutput(model_lora, clip_lora) +class NInputsTest(io.ComfyNodeV3): + @classmethod + def DEFINE_SCHEMA(cls): + return io.SchemaV3( + node_id="V3_NInputsTest", + display_name="V3 N Inputs Test", + inputs=[ + io.AutogrowDynamic.Input("nmock", template_input=io.Image.Input("image"), min=1, max=3), + io.AutogrowDynamic.Input("nmock2", template_input=io.Int.Input("int"), optional=True, min=1, max=4), + ], + outputs=[ + io.Image.Output("image_out"), + ], + ) + + @classmethod + def execute(cls, nmock, nmock2): + first_image = nmock[0] + all_images = [] + for img in nmock: + if img.shape != first_image.shape: + img = img.movedim(-1,1) + img = comfy.utils.common_upscale(img, first_image.shape[2], first_image.shape[1], "lanczos", "center") + img = img.movedim(1,-1) + all_images.append(img) + combined_image = torch.cat(all_images, dim=0) + return io.NodeOutput(combined_image) + + NODES_LIST: list[type[io.ComfyNodeV3]] = [ V3TestNode, V3LoraLoader, + NInputsTest, ] diff --git a/execution.py b/execution.py index 5ab7f1fe0..3e58b53ef 100644 --- a/execution.py +++ b/execution.py @@ -28,7 +28,7 @@ from comfy_execution.graph import ( ) from comfy_execution.graph_utils import GraphBuilder, is_link from comfy_execution.validation import validate_node_input -from comfy_api.v3.io import NodeOutput, ComfyNodeV3, Hidden, NodeStateLocal, ResourcesLocal +from comfy_api.v3.io import NodeOutput, ComfyNodeV3, Hidden, NodeStateLocal, ResourcesLocal, AutogrowDynamic class ExecutionResult(Enum): @@ -229,6 +229,16 @@ def _map_node_over_list(obj, input_data_all, func, allow_interrupt=False, execut if obj.local_resources is None: obj.local_resources = ResourcesLocal() class_clone.resources = obj.local_resources + # TODO: delete this when done testing mocking dynamic inputs + for si in obj.SCHEMA.inputs: + if isinstance(si, AutogrowDynamic.Input): + add_key = si.id + dynamic_list = [] + real_inputs = {k: v for k, v in inputs.items()} + for d in si.get_dynamic(): + dynamic_list.append(real_inputs.pop(d.id, None)) + dynamic_list = [x for x in dynamic_list if x is not None] + inputs = {**real_inputs, add_key: dynamic_list} results.append(getattr(type(obj), func).__func__(class_clone, **inputs)) # V1 else: