From c9c9f8a3d63392170d044356f8bbb39233724978 Mon Sep 17 00:00:00 2001 From: Pauan Date: Fri, 27 Mar 2026 20:47:34 -0700 Subject: [PATCH 1/6] Adding in new Krita Group Layer node --- ai_diffusion/api.py | 6 +++++ ai_diffusion/comfy_workflow.py | 8 ++++++ ai_diffusion/custom_workflow.py | 43 +++++++++++++++++++++++------- ai_diffusion/layer.py | 8 ++++++ ai_diffusion/ui/custom_workflow.py | 2 ++ ai_diffusion/workflow.py | 7 +++++ tests/test_custom_workflow.py | 2 ++ 7 files changed, 66 insertions(+), 10 deletions(-) diff --git a/ai_diffusion/api.py b/ai_diffusion/api.py index 761fb9f05..20e89ddc2 100644 --- a/ai_diffusion/api.py +++ b/ai_diffusion/api.py @@ -174,6 +174,12 @@ class CustomStyleInput: negative_prompt: str +@dataclass +class CustomLayerInput: + images: ImageCollection + names: list[str] + + @dataclass class CustomWorkflowInput: workflow: dict diff --git a/ai_diffusion/comfy_workflow.py b/ai_diffusion/comfy_workflow.py index f8f44cddc..7c7245088 100644 --- a/ai_diffusion/comfy_workflow.py +++ b/ai_diffusion/comfy_workflow.py @@ -1200,6 +1200,14 @@ def load_image_and_mask(self, images: Image | ImageCollection): assert result is not None return result + def send_list_str(self, list: list[str]): + output = self.add("ETN_ListEmpty", 1) + + for item in list: + output = self.add("ETN_ListAppend", 1, list=output, item=item) + + return self.add("ETN_DataList", 1, list=output) + def send_image(self, image: Output): if self._run_mode is ComfyRunMode.runtime: return self.add("ETN_ReturnImage", 1, images=image) diff --git a/ai_diffusion/custom_workflow.py b/ai_diffusion/custom_workflow.py index 6b02545c5..1a48f1bee 100644 --- a/ai_diffusion/custom_workflow.py +++ b/ai_diffusion/custom_workflow.py @@ -21,11 +21,11 @@ ) from . import eventloop -from .api import CustomStyleInput, InpaintContext, WorkflowInput +from .api import CustomStyleInput, CustomLayerInput, InpaintContext, WorkflowInput from .client import ClientModels, ClientOutput, JobInfoOutput, OutputBatchMode, TextOutput from .comfy_workflow import ComfyNode, ComfyWorkflow from .connection import Connection, ConnectionState -from .image import Bounds, Image, Mask +from .image import Bounds, Image, Mask, ImageCollection from .jobs import Job, JobKind, JobParams, JobQueue from .localization import translate as _ from .properties import ObservableProperties, Property @@ -259,14 +259,15 @@ def __getitem__(self, index: int): class ParamKind(Enum): image_layer = 0 mask_layer = 1 - number_int = 2 - number_float = 3 - toggle = 4 - text = 5 - prompt_positive = 6 - prompt_negative = 7 - choice = 8 - style = 9 + group_layer = 2 + number_int = 3 + number_float = 4 + toggle = 5 + text = 6 + prompt_positive = 7 + prompt_negative = 8 + choice = 9 + style = 10 class CustomParam(NamedTuple): @@ -327,6 +328,9 @@ def workflow_parameters(w: ComfyWorkflow): case ("ETN_KritaMaskLayer", _): name = node.input("name", "Mask") yield CustomParam(ParamKind.mask_layer, name) + case ("ETN_KritaGroupLayer", _): + name = node.input("name", "Group") + yield CustomParam(ParamKind.group_layer, name) case ("ETN_Parameter", "number (integer)"): name = node.input("name", "Parameter") default = node.input("default", 0) @@ -565,6 +569,7 @@ def collect_parameters( params[md.name] = layer.get_pixel_frames(bounds) else: params[md.name] = layer.get_pixels(bounds) + elif md.kind is ParamKind.mask_layer: if param is None and len(layers.masks) > 0: param = layers.masks[0].id @@ -575,6 +580,24 @@ def collect_parameters( params[md.name] = layer.get_mask_frames(bounds) else: params[md.name] = layer.get_mask(bounds) + + elif md.kind is ParamKind.group_layer: + if param is None and len(layers.images) > 0: + param = layers.images[0].id + layer = layers.find(QUuid(param)) + if layer is None: + raise ValueError(f"Input layer for parameter {md.name} not found") + + children = list(layer.get_child_images(bounds)) + + names = [name for (name, image) in children] + images = [image for (name, image) in children] + + params[md.name] = CustomLayerInput( + ImageCollection(images), + names, + ) + elif md.kind is ParamKind.style: style = Styles.list().find(str(param)) if style is None: diff --git a/ai_diffusion/layer.py b/ai_diffusion/layer.py index b9c921a33..d56939691 100644 --- a/ai_diffusion/layer.py +++ b/ai_diffusion/layer.py @@ -213,6 +213,14 @@ def get_pixel_frames(self, bounds: Bounds | None = None): def get_mask_frames(self, bounds: Bounds | None = None): return self._get_frames(self.get_mask, bounds) + def get_child_images(self, bounds: Bounds | None = None): + if self.type is LayerType.group: + for child in self.child_layers: + yield from child.get_child_images(bounds) + + elif self.type.is_image: + yield (self.name, self.get_pixels(bounds)) + def move_to_top(self): parent = self._node.parentNode() if acquire_elements(parent.childNodes())[-1] == self._node: diff --git a/ai_diffusion/ui/custom_workflow.py b/ai_diffusion/ui/custom_workflow.py index 2bf3af183..1415a32f3 100644 --- a/ai_diffusion/ui/custom_workflow.py +++ b/ai_diffusion/ui/custom_workflow.py @@ -429,6 +429,8 @@ def _create_param_widget(param: CustomParam, parent: "WorkflowParamsWidget") -> return LayerSelect("image", parent) case ParamKind.mask_layer: return LayerSelect("mask", parent) + case ParamKind.group_layer: + return LayerSelect("image", parent) case ParamKind.number_int: return IntParamWidget(param, parent) case ParamKind.number_float: diff --git a/ai_diffusion/workflow.py b/ai_diffusion/workflow.py index 4b62b0573..aa39ac704 100644 --- a/ai_diffusion/workflow.py +++ b/ai_diffusion/workflow.py @@ -11,6 +11,7 @@ CheckpointInput, ConditioningInput, ControlInput, + CustomLayerInput, CustomStyleInput, CustomWorkflowInput, ExtentInput, @@ -1463,6 +1464,12 @@ def get_param(node: ComfyNode, expected_type: type | tuple[type, type] | None = outputs[node.output(1)] = mask case "ETN_KritaMaskLayer": outputs[node.output(0)] = w.load_mask(get_param(node, (Image, ImageCollection))) + case "ETN_KritaGroupLayer": + layer: CustomLayerInput = get_param(node, CustomLayerInput) + img, mask = w.load_image_and_mask(layer.images) + outputs[node.output(0)] = img + outputs[node.output(1)] = mask + outputs[node.output(2)] = w.send_list_str(layer.names) case "ETN_KritaStyle": style: CustomStyleInput = get_param(node, CustomStyleInput) model, clip, vae = load_checkpoint_with_lora(w, style.models, models) diff --git a/tests/test_custom_workflow.py b/tests/test_custom_workflow.py index 0c93e4376..ba6629f47 100644 --- a/tests/test_custom_workflow.py +++ b/tests/test_custom_workflow.py @@ -343,6 +343,7 @@ def test_parameters(): w.add("ChoiceNodeV3", 1, choice_param=choice_param_v3) w.add("ETN_KritaImageLayer", 1, name="image") w.add("ETN_KritaMaskLayer", 1, name="mask") + w.add("ETN_KritaGroupLayer", 3, name="group") w.add("ETN_KritaStyle", 9, name="style", sampler_preset="live") # type: ignore assert list(workflow_parameters(w)) == [ @@ -357,6 +358,7 @@ def test_parameters(): CustomParam(ParamKind.choice, "choice_v3", "c", choices=["a", "b", "c"]), CustomParam(ParamKind.image_layer, "image"), CustomParam(ParamKind.mask_layer, "mask"), + CustomParam(ParamKind.group_layer, "group"), CustomParam(ParamKind.style, "style", "live"), ] From 7e29dd62e6e4cda35da0421dae2aff4880f22314 Mon Sep 17 00:00:00 2001 From: Pauan Date: Fri, 27 Mar 2026 21:03:35 -0700 Subject: [PATCH 2/6] Fixing ruff errors --- ai_diffusion/custom_workflow.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ai_diffusion/custom_workflow.py b/ai_diffusion/custom_workflow.py index 1a48f1bee..7ec6f64af 100644 --- a/ai_diffusion/custom_workflow.py +++ b/ai_diffusion/custom_workflow.py @@ -21,11 +21,11 @@ ) from . import eventloop -from .api import CustomStyleInput, CustomLayerInput, InpaintContext, WorkflowInput +from .api import CustomLayerInput, CustomStyleInput, InpaintContext, WorkflowInput from .client import ClientModels, ClientOutput, JobInfoOutput, OutputBatchMode, TextOutput from .comfy_workflow import ComfyNode, ComfyWorkflow from .connection import Connection, ConnectionState -from .image import Bounds, Image, Mask, ImageCollection +from .image import Bounds, Image, ImageCollection, Mask from .jobs import Job, JobKind, JobParams, JobQueue from .localization import translate as _ from .properties import ObservableProperties, Property From 1684611e96986ba28fa518a9aef8d2d51c26652b Mon Sep 17 00:00:00 2001 From: Pauan Date: Fri, 27 Mar 2026 21:49:44 -0700 Subject: [PATCH 3/6] The Krita Group Layer node now returns a data list instead of a batch --- ai_diffusion/comfy_workflow.py | 7 +++++++ ai_diffusion/workflow.py | 4 ++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/ai_diffusion/comfy_workflow.py b/ai_diffusion/comfy_workflow.py index 7c7245088..a9b024a43 100644 --- a/ai_diffusion/comfy_workflow.py +++ b/ai_diffusion/comfy_workflow.py @@ -1033,6 +1033,9 @@ def invert_image(self, image: Output): def batch_image(self, batch: Output, image: Output): return self.add("ImageBatch", 1, image1=batch, image2=image) + def unbatch_image(self, image: Output): + return self.add("RebatchImages", 1, images=image, batch_size=1) + def image_batch_element(self, batch: Output, index: int): return self.add("ImageFromBatch", 1, image=batch, batch_index=index, length=1) @@ -1113,6 +1116,10 @@ def batch_mask(self, batch: Output, mask: Output): image = self.mask_to_image(mask) return self.image_to_mask(self.batch_image(image_batch, image)) + def unbatch_mask(self, mask: Output): + image = self.mask_to_image(mask) + return self.image_to_mask(self.unbatch_image(image)) + def mask_batch_element(self, mask_batch: Output, index: int): image_batch = self.mask_to_image(mask_batch) image = self.image_batch_element(image_batch, index) diff --git a/ai_diffusion/workflow.py b/ai_diffusion/workflow.py index aa39ac704..c5a1e7e93 100644 --- a/ai_diffusion/workflow.py +++ b/ai_diffusion/workflow.py @@ -1467,8 +1467,8 @@ def get_param(node: ComfyNode, expected_type: type | tuple[type, type] | None = case "ETN_KritaGroupLayer": layer: CustomLayerInput = get_param(node, CustomLayerInput) img, mask = w.load_image_and_mask(layer.images) - outputs[node.output(0)] = img - outputs[node.output(1)] = mask + outputs[node.output(0)] = w.unbatch_image(img) + outputs[node.output(1)] = w.unbatch_mask(mask) outputs[node.output(2)] = w.send_list_str(layer.names) case "ETN_KritaStyle": style: CustomStyleInput = get_param(node, CustomStyleInput) From b7df402f5634c9e55e656f7d1a3f95624e13075c Mon Sep 17 00:00:00 2001 From: Pauan Date: Thu, 2 Apr 2026 22:13:09 -0700 Subject: [PATCH 4/6] Removing Layer Group node and merging its behavior into Layer Image node --- ai_diffusion/api.py | 3 +- ai_diffusion/custom_workflow.py | 68 ++++++++++++++++-------------- ai_diffusion/ui/custom_workflow.py | 2 - ai_diffusion/workflow.py | 13 +++--- tests/test_custom_workflow.py | 4 +- 5 files changed, 45 insertions(+), 45 deletions(-) diff --git a/ai_diffusion/api.py b/ai_diffusion/api.py index 20e89ddc2..c155159ed 100644 --- a/ai_diffusion/api.py +++ b/ai_diffusion/api.py @@ -176,8 +176,9 @@ class CustomStyleInput: @dataclass class CustomLayerInput: - images: ImageCollection + images: Image | ImageCollection names: list[str] + is_batch: bool @dataclass diff --git a/ai_diffusion/custom_workflow.py b/ai_diffusion/custom_workflow.py index 7ec6f64af..8ca99e254 100644 --- a/ai_diffusion/custom_workflow.py +++ b/ai_diffusion/custom_workflow.py @@ -259,15 +259,14 @@ def __getitem__(self, index: int): class ParamKind(Enum): image_layer = 0 mask_layer = 1 - group_layer = 2 - number_int = 3 - number_float = 4 - toggle = 5 - text = 6 - prompt_positive = 7 - prompt_negative = 8 - choice = 9 - style = 10 + number_int = 2 + number_float = 3 + toggle = 4 + text = 5 + prompt_positive = 6 + prompt_negative = 7 + choice = 8 + style = 9 class CustomParam(NamedTuple): @@ -324,13 +323,10 @@ def workflow_parameters(w: ComfyWorkflow): yield CustomParam(ParamKind.style, name, node.input("sampler_preset", "auto")) case ("ETN_KritaImageLayer", _): name = node.input("name", "Image") - yield CustomParam(ParamKind.image_layer, name) + yield CustomParam(ParamKind.image_layer, name, node.input("group_mode", "flatten")) case ("ETN_KritaMaskLayer", _): name = node.input("name", "Mask") yield CustomParam(ParamKind.mask_layer, name) - case ("ETN_KritaGroupLayer", _): - name = node.input("name", "Group") - yield CustomParam(ParamKind.group_layer, name) case ("ETN_Parameter", "number (integer)"): name = node.input("name", "Parameter") default = node.input("default", 0) @@ -565,10 +561,35 @@ def collect_parameters( layer = layers.find(QUuid(param)) if layer is None: raise ValueError(f"Input layer for parameter {md.name} not found") + if is_animation and layer.is_animated: - params[md.name] = layer.get_pixel_frames(bounds) + params[md.name] = CustomLayerInput( + layer.get_pixel_frames(bounds), + [layer.name], + True, + ) + + elif md.default == "flatten": + params[md.name] = CustomLayerInput( + layer.get_pixels(bounds), + [layer.name], + True, + ) + + elif md.default == "all children": + children = list(layer.get_child_images(bounds)) + + names = [name for (name, image) in children] + images = [image for (name, image) in children] + + params[md.name] = CustomLayerInput( + ImageCollection(images), + names, + False, + ) + else: - params[md.name] = layer.get_pixels(bounds) + raise ValueError(f"Unknown group_mode {md.default}") elif md.kind is ParamKind.mask_layer: if param is None and len(layers.masks) > 0: @@ -581,23 +602,6 @@ def collect_parameters( else: params[md.name] = layer.get_mask(bounds) - elif md.kind is ParamKind.group_layer: - if param is None and len(layers.images) > 0: - param = layers.images[0].id - layer = layers.find(QUuid(param)) - if layer is None: - raise ValueError(f"Input layer for parameter {md.name} not found") - - children = list(layer.get_child_images(bounds)) - - names = [name for (name, image) in children] - images = [image for (name, image) in children] - - params[md.name] = CustomLayerInput( - ImageCollection(images), - names, - ) - elif md.kind is ParamKind.style: style = Styles.list().find(str(param)) if style is None: diff --git a/ai_diffusion/ui/custom_workflow.py b/ai_diffusion/ui/custom_workflow.py index 1415a32f3..2bf3af183 100644 --- a/ai_diffusion/ui/custom_workflow.py +++ b/ai_diffusion/ui/custom_workflow.py @@ -429,8 +429,6 @@ def _create_param_widget(param: CustomParam, parent: "WorkflowParamsWidget") -> return LayerSelect("image", parent) case ParamKind.mask_layer: return LayerSelect("mask", parent) - case ParamKind.group_layer: - return LayerSelect("image", parent) case ParamKind.number_int: return IntParamWidget(param, parent) case ParamKind.number_float: diff --git a/ai_diffusion/workflow.py b/ai_diffusion/workflow.py index c5a1e7e93..ef98fff77 100644 --- a/ai_diffusion/workflow.py +++ b/ai_diffusion/workflow.py @@ -1459,17 +1459,16 @@ def get_param(node: ComfyNode, expected_type: type | tuple[type, type] | None = case "ETN_Parameter": outputs[node.output(0)] = get_param(node) case "ETN_KritaImageLayer": - img, mask = w.load_image_and_mask(get_param(node, (Image, ImageCollection))) + layer: CustomLayerInput = get_param(node, CustomLayerInput) + img, mask = w.load_image_and_mask(layer.images) + if not layer.is_batch: + img = w.unbatch_image(img) + mask = w.unbatch_mask(mask) outputs[node.output(0)] = img outputs[node.output(1)] = mask + outputs[node.output(2)] = w.send_list_str(layer.names) case "ETN_KritaMaskLayer": outputs[node.output(0)] = w.load_mask(get_param(node, (Image, ImageCollection))) - case "ETN_KritaGroupLayer": - layer: CustomLayerInput = get_param(node, CustomLayerInput) - img, mask = w.load_image_and_mask(layer.images) - outputs[node.output(0)] = w.unbatch_image(img) - outputs[node.output(1)] = w.unbatch_mask(mask) - outputs[node.output(2)] = w.send_list_str(layer.names) case "ETN_KritaStyle": style: CustomStyleInput = get_param(node, CustomStyleInput) model, clip, vae = load_checkpoint_with_lora(w, style.models, models) diff --git a/tests/test_custom_workflow.py b/tests/test_custom_workflow.py index ba6629f47..5f3fb68a3 100644 --- a/tests/test_custom_workflow.py +++ b/tests/test_custom_workflow.py @@ -341,9 +341,8 @@ def test_parameters(): w.add("ChoiceNode", 1, choice_param=choice_param) choice_param_v3 = w.add("ETN_Parameter", 1, name="choice_v3", type="choice", default="c") w.add("ChoiceNodeV3", 1, choice_param=choice_param_v3) - w.add("ETN_KritaImageLayer", 1, name="image") + w.add("ETN_KritaImageLayer", 3, name="image") w.add("ETN_KritaMaskLayer", 1, name="mask") - w.add("ETN_KritaGroupLayer", 3, name="group") w.add("ETN_KritaStyle", 9, name="style", sampler_preset="live") # type: ignore assert list(workflow_parameters(w)) == [ @@ -358,7 +357,6 @@ def test_parameters(): CustomParam(ParamKind.choice, "choice_v3", "c", choices=["a", "b", "c"]), CustomParam(ParamKind.image_layer, "image"), CustomParam(ParamKind.mask_layer, "mask"), - CustomParam(ParamKind.group_layer, "group"), CustomParam(ParamKind.style, "style", "live"), ] From 19380dbb9ad504af4d1fa363f699e48092db9596 Mon Sep 17 00:00:00 2001 From: Pauan Date: Thu, 2 Apr 2026 22:54:55 -0700 Subject: [PATCH 5/6] Now supports sending lists of any type --- ai_diffusion/comfy_workflow.py | 2 +- ai_diffusion/workflow.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ai_diffusion/comfy_workflow.py b/ai_diffusion/comfy_workflow.py index a9b024a43..b292e80bd 100644 --- a/ai_diffusion/comfy_workflow.py +++ b/ai_diffusion/comfy_workflow.py @@ -1207,7 +1207,7 @@ def load_image_and_mask(self, images: Image | ImageCollection): assert result is not None return result - def send_list_str(self, list: list[str]): + def send_list(self, list: list[Input]): output = self.add("ETN_ListEmpty", 1) for item in list: diff --git a/ai_diffusion/workflow.py b/ai_diffusion/workflow.py index ef98fff77..787da9473 100644 --- a/ai_diffusion/workflow.py +++ b/ai_diffusion/workflow.py @@ -1466,7 +1466,7 @@ def get_param(node: ComfyNode, expected_type: type | tuple[type, type] | None = mask = w.unbatch_mask(mask) outputs[node.output(0)] = img outputs[node.output(1)] = mask - outputs[node.output(2)] = w.send_list_str(layer.names) + outputs[node.output(2)] = w.send_list(layer.names) case "ETN_KritaMaskLayer": outputs[node.output(0)] = w.load_mask(get_param(node, (Image, ImageCollection))) case "ETN_KritaStyle": From 3f5c60876351cc09369fd56022b93a4094b821c0 Mon Sep 17 00:00:00 2001 From: Pauan Date: Fri, 3 Apr 2026 08:41:45 -0700 Subject: [PATCH 6/6] Minor optimization for send_list --- ai_diffusion/comfy_workflow.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ai_diffusion/comfy_workflow.py b/ai_diffusion/comfy_workflow.py index b292e80bd..47f4ee285 100644 --- a/ai_diffusion/comfy_workflow.py +++ b/ai_diffusion/comfy_workflow.py @@ -1208,6 +1208,9 @@ def load_image_and_mask(self, images: Image | ImageCollection): return result def send_list(self, list: list[Input]): + if len(list) == 1: + return list[0] + output = self.add("ETN_ListEmpty", 1) for item in list: