Adding in new Krita Group Layer node

ai_diffusion/api.py

@@ -174,6 +174,13 @@ class CustomStyleInput:
     negative_prompt: str
 
 
+@dataclass
+class CustomLayerInput:
+    images: Image | ImageCollection
+    names: list[str]
+    is_batch: bool
+
+
 @dataclass
 class CustomWorkflowInput:
     workflow: dict

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)
@@ -1200,6 +1207,17 @@ def load_image_and_mask(self, images: Image | ImageCollection):
         assert result is not None
         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:
+            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)

ai_diffusion/custom_workflow.py

@@ -21,11 +21,11 @@
 )
 
 from . import eventloop
-from .api import CustomStyleInput, 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
+from .image import Bounds, Image, ImageCollection, Mask
 from .jobs import Job, JobKind, JobParams, JobQueue
 from .localization import translate as _
 from .properties import ObservableProperties, Property
@@ -323,7 +323,7 @@ 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)
@@ -561,10 +561,36 @@ 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:
                     param = layers.masks[0].id
@@ -575,6 +601,7 @@ def collect_parameters(
                     params[md.name] = layer.get_mask_frames(bounds)
                 else:
                     params[md.name] = layer.get_mask(bounds)
+
             elif md.kind is ParamKind.style:
                 style = Styles.list().find(str(param))
                 if style is None:

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:

ai_diffusion/workflow.py

@@ -11,6 +11,7 @@
     CheckpointInput,
     ConditioningInput,
     ControlInput,
+    CustomLayerInput,
     CustomStyleInput,
     CustomWorkflowInput,
     ExtentInput,
@@ -1458,9 +1459,14 @@ 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(layer.names)
             case "ETN_KritaMaskLayer":
                 outputs[node.output(0)] = w.load_mask(get_param(node, (Image, ImageCollection)))
             case "ETN_KritaStyle":

tests/test_custom_workflow.py

@@ -341,7 +341,7 @@ 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_KritaStyle", 9, name="style", sampler_preset="live")  # type: ignore