Support sending/receiving image tiles (crop regions) in custom workflows #2230

* requires a KritaSelection node which defines the crop bbox
* applies to the selection mask and KritaCanvas
* does not affect KritaImageLayer or KritaMaskLayer
* KritaOutput can have arbitrary offsets

Author: Acly

ai_diffusion/api.py

@@ -123,6 +123,13 @@ class InpaintMode(Enum):
     custom = 6
 
 
+class InpaintContext(Enum):
+    automatic = 0
+    mask_bounds = 1
+    entire_image = 2
+    layer_bounds = 3
+
+
 class FillMode(Enum):
     none = 0
     neutral = 1

ai_diffusion/client.py

@@ -8,7 +8,7 @@
 
 from .api import WorkflowInput
 from .comfy_workflow import ComfyObjectInfo
-from .image import ImageCollection
+from .image import ImageCollection, Point
 from .properties import Property, ObservableProperties
 from .files import FileLibrary, FileFormat
 from .style import Style
@@ -49,6 +49,7 @@ class OutputBatchMode(Enum):
 
 class JobInfoOutput(NamedTuple):
     name: str = ""
+    offset: Point = Point(0, 0)
     batch_mode: OutputBatchMode = OutputBatchMode.default
     resize_canvas: bool = False
 

ai_diffusion/comfy_client.py

@@ -16,15 +16,15 @@
 from .client import filter_supported_styles, loras_to_upload
 from .comfy_workflow import ComfyObjectInfo
 from .files import FileFormat
-from .image import Image, ImageCollection
+from .image import Image, ImageCollection, Point
 from .network import RequestManager, NetworkError
 from .websockets.src import websockets
 from .style import Styles
 from .resources import ControlMode, ResourceId, ResourceKind, Arch
 from .resources import CustomNode, UpscalerName, resource_id
 from .settings import PerformanceSettings, settings
 from .localization import translate as _
-from .util import client_logger as log
+from .util import client_logger as log, parse_enum
 from .workflow import create as create_workflow
 from . import platform_tools, resources, util
 
@@ -906,7 +906,8 @@ def _extract_job_info_output(job_id: str, msg: dict):
             if isinstance(info, dict):
                 result = JobInfoOutput(
                     name=info.get("name", ""),
-                    batch_mode=OutputBatchMode[info.get("batch_mode", "default")],
+                    offset=Point(info.get("offset_x", 0), info.get("offset_y", 0)),
+                    batch_mode=parse_enum(OutputBatchMode, info.get("batch_mode", "default")),
                     resize_canvas=info.get("resize_canvas", False),
                 )
                 return ClientMessage(ClientEvent.output, job_id, result=result)

ai_diffusion/comfy_workflow.py

@@ -1352,6 +1352,12 @@ def inputs(self, node_name: str, category="") -> dict[str, list] | None:
         result.update(inputs.get("optional", {}))
         return result
 
+    def outputs(self, node_name: str) -> list[str]:
+        node = self.nodes.get(node_name)
+        if node is None:
+            return []
+        return node.get("output_name", [])
+
 
 def _convert_ui_workflow(w: dict, node_inputs: ComfyObjectInfo):
     version = w.get("version")

ai_diffusion/custom_workflow.py

@@ -10,15 +10,15 @@
 from PyQt5.QtCore import Qt, QObject, QUuid, QAbstractListModel, QSortFilterProxyModel, QModelIndex
 from PyQt5.QtCore import QMetaObject, QTimer, pyqtSignal
 
-from .api import WorkflowInput
+from .api import WorkflowInput, InpaintContext
 from .client import OutputBatchMode, TextOutput, ClientOutput, JobInfoOutput
 from .comfy_workflow import ComfyWorkflow, ComfyNode
 from .connection import Connection, ConnectionState
-from .image import Bounds, Image
+from .image import Bounds, Image, Mask
 from .jobs import Job, JobParams, JobQueue, JobKind
 from .properties import Property, ObservableProperties
 from .style import Styles
-from .util import base_type_match, user_data_dir, client_logger as log
+from .util import base_type_match, parse_enum, user_data_dir, client_logger as log
 from .ui import theme
 from . import eventloop
 
@@ -531,6 +531,30 @@ def collect_parameters(self, layers: "LayerManager", bounds: Bounds, animation=F
 
         return params
 
+    def prepare_mask(
+        self,
+        selection_node: ComfyNode,
+        mask: Mask | None,
+        mask_bounds: Bounds | None,
+        canvas_bounds: Bounds,
+    ):
+        ctx = selection_node.input("context", "entire_image").replace(" ", "_")
+        pad = selection_node.input("padding", 0)
+        if mask and mask_bounds:
+            match parse_enum(InpaintContext, ctx):
+                case InpaintContext.entire_image:
+                    bounds = canvas_bounds
+                case InpaintContext.automatic:
+                    bounds = Bounds.pad(mask.bounds, pad)
+                case InpaintContext.mask_bounds:
+                    bounds = Bounds.pad(mask_bounds, pad)
+                case _:
+                    raise ValueError(f"Invalid inpaint context: {ctx}")
+            bounds = Bounds.clamp(bounds, canvas_bounds.extent)
+            mask.bounds = mask.bounds.relative_to(bounds)
+            return mask, bounds
+        return None, canvas_bounds
+
     def switch_to_web_workflow(self):
         self._switch_workflow_bind = self._workflows.rowsInserted.connect(self._set_workflow_index)
         self._switch_workflow_timer = QTimer()
@@ -559,6 +583,7 @@ def handle_output(self, job: Job, output: ClientOutput | None):
             self.outputs_changed.emit(self.outputs)
         elif isinstance(output, JobInfoOutput):
             job.params.resize_canvas = output.resize_canvas
+            job.params.bounds = Bounds(*output.offset, *job.params.bounds.extent)
             if output.name:
                 job.params.name = output.name
             match output.batch_mode:

ai_diffusion/image.py

@@ -73,6 +73,10 @@ def from_qsize(qsize: QSize):
     def largest(a: "Extent", b: "Extent"):
         return a if a.width * a.height > b.width * b.height else b
 
+    @staticmethod
+    def min(a: "Extent", b: "Extent"):
+        return Extent(min(a.width, b.width), min(a.height, b.height))
+
     @staticmethod
     def ratio(a: "Extent", b: "Extent"):
         return sqrt(a.pixel_count / b.pixel_count)

ai_diffusion/model.py

@@ -16,7 +16,8 @@
 
 from . import eventloop, workflow, util
 from .api import ConditioningInput, ControlInput, WorkflowKind, WorkflowInput, SamplingInput
-from .api import InpaintMode, InpaintParams, FillMode, ImageInput, CustomWorkflowInput, UpscaleInput
+from .api import FillMode, ImageInput, CustomWorkflowInput, UpscaleInput
+from .api import InpaintMode, InpaintContext, InpaintParams
 from .localization import translate as _
 from .util import clamp, ensure, trim_text, client_logger as log
 from .settings import ApplyBehavior, ApplyRegionBehavior, GenerationFinishedAction, ImageFileFormat
@@ -472,23 +473,28 @@ async def _generate_custom(self, previous_input: WorkflowInput | None):
 
         try:
             wf = ensure(self.custom.graph)
-            bounds = Bounds(0, 0, *self._doc.extent)
-            img_input = ImageInput.from_extent(bounds.extent)
-            img_input.initial_image = self._get_current_image(bounds)
             is_live = self.custom.mode is CustomGenerationMode.live
             is_anim = self.custom.mode is CustomGenerationMode.animation
             seed = self.seed if is_live or self.fixed_seed else workflow.generate_seed()
+            canvas_bounds = Bounds(0, 0, *self._doc.extent)
+            bounds = canvas_bounds
+            mask = None
 
-            if next(wf.find(type="ETN_KritaSelection"), None):
-                mask, _ = self._doc.create_mask_from_selection()
-                if mask:
-                    img_input.hires_mask = mask.to_image(bounds.extent)
+            if selection_node := next(wf.find(type="ETN_KritaSelection"), None):
+                mods = get_selection_modifiers(InpaintMode.fill, self.strength, is_live)
+                mask, select_bounds = self._doc.create_mask_from_selection(mods.padding, 8, 256)
+                mask, bounds = self.custom.prepare_mask(selection_node, mask, select_bounds, bounds)
 
-            params = self.custom.collect_parameters(self.layers, bounds, is_anim)
+            img_input = ImageInput.from_extent(bounds.extent)
+            img_input.initial_image = self._get_current_image(bounds)
+            img_input.hires_mask = mask.to_image(bounds.extent) if mask else None
+
+            params = self.custom.collect_parameters(self.layers, canvas_bounds, is_anim)
             input = WorkflowInput(
                 WorkflowKind.custom,
                 img_input,
                 sampling=SamplingInput("custom", "custom", 1, 1000, seed=seed),
+                inpaint=InpaintParams(InpaintMode.fill, bounds),
                 custom_workflow=CustomWorkflowInput(wf.root, params),
             )
             job_params = JobParams(bounds, self.custom.job_name, metadata=self.custom.params)
@@ -646,7 +652,7 @@ def show_preview(self, job_id: str, index: int, name_prefix="Preview"):
         image = job.results[index]
         bounds = job.params.bounds
         if image.extent != bounds.extent:
-            image = Image.crop(image, Bounds(0, 0, *bounds.extent))
+            image = Image.crop(image, Bounds(0, 0, *Extent.min(bounds.extent, image.extent)))
         if self._layer and self._layer.was_removed:
             self._layer = None  # layer was removed by user
         if self._layer is not None:
@@ -950,13 +956,6 @@ def edit_style(self) -> Style | None:
         return None
 
 
-class InpaintContext(Enum):
-    automatic = 0
-    mask_bounds = 1
-    entire_image = 2
-    layer_bounds = 3
-
-
 class CustomInpaint(QObject, ObservableProperties):
     mode = Property(InpaintMode.automatic, persist=True)
     fill = Property(FillMode.neutral, persist=True)

ai_diffusion/util.py

@@ -16,6 +16,7 @@
 
 T = TypeVar("T")
 R = TypeVar("R")
+E = TypeVar("E", bound=Enum)
 QOBJECT = TypeVar("QOBJECT", bound=QObject)
 
 plugin_dir = dir = Path(__file__).parent
@@ -92,6 +93,15 @@ def ensure(value: Optional[T], msg="") -> T:
     return value
 
 
+def parse_enum(enum_class: type[E], value: str, default: E | None = None) -> E:
+    try:
+        return enum_class[value]
+    except KeyError:
+        if default is not None:
+            return default
+        raise ValueError(f"Invalid value '{value}' for enum {enum_class.__name__}")
+
+
 def maybe(func: Callable[[T], R], value: Optional[T]) -> Optional[R]:
     if value is not None:
         return func(value)

ai_diffusion/workflow.py

@@ -1328,6 +1328,7 @@ def expand_custom(
     w: ComfyWorkflow,
     input: CustomWorkflowInput,
     images: ImageInput,
+    bounds: Bounds,
     seed: int,
     models: ClientModels,
 ):
@@ -1373,6 +1374,8 @@ def get_param(node: ComfyNode, expected_type: type | tuple[type, type] | None =
                     image = ensure(images.initial_image)
                     outputs[node.output(0)] = w.solid_mask(image.extent, 1.0)
                 outputs[node.output(1)] = images.hires_mask is not None
+                outputs[node.output(2)] = bounds.x
+                outputs[node.output(3)] = bounds.y
             case "ETN_Parameter":
                 outputs[node.output(0)] = get_param(node)
             case "ETN_KritaImageLayer":
@@ -1689,7 +1692,14 @@ def create(i: WorkflowInput, models: ClientModels, comfy_mode=ComfyRunMode.serve
         )
     elif i.kind is WorkflowKind.custom:
         seed = ensure(i.sampling).seed
-        return expand_custom(workflow, ensure(i.custom_workflow), ensure(i.images), seed, models)
+        return expand_custom(
+            workflow,
+            ensure(i.custom_workflow),
+            ensure(i.images),
+            ensure(i.inpaint).target_bounds,
+            seed,
+            models,
+        )
     else:
         raise ValueError(f"Unsupported workflow kind: {i.kind}")
 

tests/test_custom_workflow.py

@@ -349,6 +349,43 @@ def test_parameter_order():
     ]
 
 
+def test_prepare_mask():
+    connection_workflows = {"connection1": make_dummy_graph(42)}
+    connection = create_mock_connection(connection_workflows)
+    workflows = WorkflowCollection(connection)
+
+    jobs = JobQueue()
+    workspace = CustomWorkspace(workflows, dummy_generate, jobs)
+
+    mask = Mask.rectangle(Bounds(10, 10, 40, 40), 0)
+    canvas_bounds = Bounds(0, 0, 100, 100)
+    selection_bounds = Bounds(12, 12, 34, 34)
+    selection_node = ComfyNode(0, "ETN_Selection", {"context": "automatic", "padding": 3})
+
+    prepared_mask, bounds = workspace.prepare_mask(
+        selection_node, copy(mask), selection_bounds, canvas_bounds
+    )
+    assert bounds == Bounds(6, 6, 48, 48)  # mask.bounds + padding // multiple of 8
+    assert prepared_mask is not None
+    assert prepared_mask.bounds == Bounds(4, 4, 40, 40)
+
+    selection_node.inputs["context"] = "mask_bounds"
+    prepared_mask, bounds = workspace.prepare_mask(
+        selection_node, copy(mask), selection_bounds, canvas_bounds
+    )
+    assert bounds == Bounds(9, 9, 40, 40)  # selection_bounds + padding // multiple of 8
+    assert prepared_mask is not None
+    assert prepared_mask.bounds == Bounds(1, 1, 40, 40)
+
+    selection_node.inputs["context"] = "entire_image"
+    prepared_mask, bounds = workspace.prepare_mask(
+        selection_node, copy(mask), selection_bounds, canvas_bounds
+    )
+    assert bounds == canvas_bounds
+    assert prepared_mask is not None
+    assert prepared_mask.bounds == mask.bounds
+
+
 def test_text_output():
     connection_workflows = {"connection1": make_dummy_graph(42)}
     connection = create_mock_connection(connection_workflows, ComfyObjectInfo({}))
@@ -490,10 +527,7 @@ def test_expand():
     }
 
     w = ComfyWorkflow()
-    w = workflow.expand_custom(w, input, images, 123, models)
-
-    def find_img_id(image: Image):
-        return next((id for id, img in w.images.items() if img == image), "not-found")
+    w = workflow.expand_custom(w, input, images, Bounds(0, 0, 4, 4), 123, models)
 
     expected = [
         ComfyNode(1, "ETN_LoadImageCache", {"id": img_id(images.initial_image)}),
@@ -554,7 +588,7 @@ def test_expand_animation():
     models = ClientModels()
 
     w = ComfyWorkflow()
-    w = workflow.expand_custom(w, input, images, 123, models)
+    w = workflow.expand_custom(w, input, images, Bounds(0, 0, 4, 4), 123, models)
 
     expected = [
         ComfyNode(1, "ETN_LoadImageCache", {"id": img_id(in_images[0])}),
@@ -582,3 +616,53 @@ def test_expand_animation():
     ]
     for node in expected:
         assert node in w, f"Node {node} not found in\n{json.dumps(w.root, indent=2)}"
+
+
+def test_expand_selection():
+    ext = ComfyWorkflow()
+    select, select_active, off_x, off_y = ext.add(
+        "ETN_KritaSelection", 4, context="automatic", padding=2
+    )
+    canvas, width, height, seed = ext.add("ETN_KritaCanvas", 4)
+    ext.add(
+        "Sink",
+        1,
+        image=canvas,
+        width=width,
+        height=height,
+        mask=select,
+        has_selection=select_active,
+        offset_x=off_x,
+        offset_y=off_y,
+    )
+
+    params = {}
+    input = CustomWorkflowInput(workflow=ext.root, params=params)
+    images = ImageInput.from_extent(Extent(8, 16))
+    images.initial_image = Image.create(Extent(8, 16), Qt.GlobalColor.red)
+    images.hires_mask = Image.create(Extent(8, 16), Qt.GlobalColor.green)
+    bounds = Bounds(2, 3, 8, 16)  # selection from (2,2) to (6,6)
+    models = ClientModels()
+
+    w = ComfyWorkflow()
+    w = workflow.expand_custom(w, input, images, bounds, 123, models)
+
+    expected = [
+        ComfyNode(1, "ETN_LoadImageCache", {"id": img_id(images.hires_mask)}),
+        ComfyNode(2, "ETN_LoadImageCache", {"id": img_id(images.initial_image)}),
+        ComfyNode(
+            3,
+            "Sink",
+            {
+                "image": Output(2, 0),
+                "width": 8,
+                "height": 16,
+                "mask": Output(1, 1),
+                "has_selection": True,
+                "offset_x": 2,
+                "offset_y": 3,
+            },
+        ),
+    ]
+    for node in expected:
+        assert node in w, f"Node {node} not found in\n{json.dumps(w.root, indent=2)}"