diff --git a/ai_diffusion/api.py b/ai_diffusion/api.py index 761fb9f05..c155159ed 100644 --- a/ai_diffusion/api.py +++ b/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 diff --git a/ai_diffusion/comfy_workflow.py b/ai_diffusion/comfy_workflow.py index f8f44cddc..47f4ee285 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) @@ -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) diff --git a/ai_diffusion/custom_workflow.py b/ai_diffusion/custom_workflow.py index 6b02545c5..8ca99e254 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 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: 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/workflow.py b/ai_diffusion/workflow.py index 4b62b0573..787da9473 100644 --- a/ai_diffusion/workflow.py +++ b/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": diff --git a/tests/test_custom_workflow.py b/tests/test_custom_workflow.py index 0c93e4376..5f3fb68a3 100644 --- a/tests/test_custom_workflow.py +++ b/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