Add modern ControlNet inpainting support

adv_control/control.py

@@ -64,22 +64,22 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase):
 
         # make cond_hint appropriate dimensions
         # TODO: change this to not require cond_hint upscaling every step when self.sub_idxs are present
-        if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] * self.real_compression_ratio != self.cond_hint.shape[2] or x_noisy.shape[3] * self.real_compression_ratio != self.cond_hint.shape[3]:
+        if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[-2] * self.real_compression_ratio != self.cond_hint.shape[-2] or x_noisy.shape[-1] * self.real_compression_ratio != self.cond_hint.shape[-1]:
             if self.cond_hint is not None:
                 del self.cond_hint
             self.cond_hint = None
             self.real_compression_ratio = self.compression_ratio
             compression_ratio = self.compression_ratio
             if self.vae is not None and self.mult_by_ratio_when_vae:
-                compression_ratio *= self.vae.downscale_ratio
+                compression_ratio *= self.vae.spacial_compression_encode()
             # if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling
             if self.sub_idxs is not None:
                 actual_cond_hint_orig = self.cond_hint_original
                 if self.cond_hint_original.size(0) < self.full_latent_length:
                     actual_cond_hint_orig = extend_to_batch_size(tensor=actual_cond_hint_orig, batch_size=self.full_latent_length)
-                self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center")
+                self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[-1] * compression_ratio, x_noisy.shape[-2] * compression_ratio, self.upscale_algorithm, "center")
             else:
-                self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center")
+                self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[-1] * compression_ratio, x_noisy.shape[-2] * compression_ratio, self.upscale_algorithm, "center")
             self.cond_hint = self.preprocess_image(self.cond_hint)
             if self.vae is not None:
                 loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
@@ -93,7 +93,10 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase):
                 to_concat = []
                 for c in self.extra_concat_orig:
                     c = c.to(self.cond_hint.device)
-                    c = comfy.utils.common_upscale(c, self.cond_hint.shape[3], self.cond_hint.shape[2], self.upscale_algorithm, "center")
+                    c = comfy.utils.common_upscale(c, self.cond_hint.shape[-1], self.cond_hint.shape[-2], self.upscale_algorithm, "center")
+                    if c.ndim < self.cond_hint.ndim:
+                        c = c.unsqueeze(2)
+                        c = comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[2], dim=2)
                     to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0]))
                 self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1)
 
@@ -123,7 +126,7 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase):
         return super().pre_run_advanced(*args, **kwargs)
 
     def apply_advanced_strengths_and_masks(self, x: Tensor, batched_number: int, flux_shape=None):
-        if self.is_flux:
+        if self.is_flux or x.ndim == 3:
             flux_shape = self.x_noisy_shape
         return super().apply_advanced_strengths_and_masks(x, batched_number, flux_shape)
 

adv_control/nodes.py

@@ -1,7 +1,7 @@
 from comfy_api.latest import ComfyExtension, io
 
 from .nodes_main import (ControlNetLoaderAdvanced, DiffControlNetLoaderAdvanced, AnimaLLLiteLoaderAdvanced,
-                         AdvancedControlNetApply, AdvancedControlNetApplySingle)
+                         AdvancedControlNetApply, AdvancedControlNetInpaintingApply, AdvancedControlNetApplySingle)
 from .nodes_weight import (DefaultWeights, ScaledSoftMaskedUniversalWeights, ScaledSoftUniversalWeights,
                            SoftControlNetWeightsSD15, CustomControlNetWeightsSD15, CustomControlNetWeightsFlux,
                            CustomControlNetWeightsAnima, SoftT2IAdapterWeights, CustomT2IAdapterWeights, ExtrasMiddleMultNode,
@@ -32,6 +32,7 @@ class AdvancedControlNetExtension(ComfyExtension):
             LatentKeyframeBatchedGroupNode,
             LatentKeyframeGroupNode,
             AdvancedControlNetApply,
+            AdvancedControlNetInpaintingApply,
             AdvancedControlNetApplySingle,
             ControlNetLoaderAdvanced,
             DiffControlNetLoaderAdvanced,

adv_control/nodes_main.py

@@ -2,6 +2,7 @@ from comfy_api.latest import io
 from torch import Tensor
 
 import folder_paths
+import comfy.utils
 
 from .control import load_controlnet, convert_to_advanced, is_advanced_controlnet, is_sd3_advanced_controlnet
 from .control_lllite import load_anima_lllite
@@ -112,9 +113,11 @@ class AdvancedControlNetApply(io.ComfyNode):
     def execute(cls, positive, negative, control_net, image, strength, start_percent, end_percent,
                          mask_optional: Tensor=None, vae_optional=None,
                          timestep_kf: TimestepKeyframeGroup=None, latent_kf_override: LatentKeyframeGroup=None,
-                         weights_override: ControlWeights=None, control_apply_to_uncond=False):
+                         weights_override: ControlWeights=None, control_apply_to_uncond=False, extra_concat=None):
         if strength == 0:
             return io.NodeOutput(positive, negative)
+        if extra_concat is None:
+            extra_concat = []
 
         control_hint = image.movedim(-1,1)
         cnets = {}
@@ -134,7 +137,7 @@ class AdvancedControlNetApply(io.ComfyNode):
                         if control_net is None:
                             raise Exception("Passed in control_net is None; something must have went wrong when loading it from a Load ControlNet node.")
                         # copy, convert to advanced if needed, and set cond
-                        c_net = convert_to_advanced(control_net.copy()).set_cond_hint(control_hint, strength, (start_percent, end_percent), vae_optional)
+                        c_net = convert_to_advanced(control_net.copy()).set_cond_hint(control_hint, strength, (start_percent, end_percent), vae_optional, extra_concat)
                         if is_advanced_controlnet(c_net):
                             # disarm node check
                             c_net.disarm()
@@ -183,6 +186,62 @@ class AdvancedControlNetApply(io.ComfyNode):
         return io.NodeOutput(out[0], out[1])
     
 
+class AdvancedControlNetInpaintingApply(io.ComfyNode):
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        return io.Schema(
+            node_id='ACN_AdvancedControlNetInpaintingApply',
+            display_name='Apply Advanced ControlNet Inpainting 🛂🅐🅒🅝',
+            category='Adv-ControlNet 🛂🅐🅒🅝',
+            inputs=[
+                io.Conditioning.Input('positive'),
+                io.Conditioning.Input('negative'),
+                io.ControlNet.Input('control_net'),
+                io.Vae.Input('vae'),
+                io.Image.Input('image'),
+                io.Mask.Input('inpaint_mask'),
+                io.Float.Input('strength', default=1.0, max=10.0, min=0.0, step=0.01),
+                io.Float.Input('start_percent', default=0.0, max=1.0, min=0.0, step=0.001),
+                io.Float.Input('end_percent', default=1.0, max=1.0, min=0.0, step=0.001),
+                io.Mask.Input('effect_mask_optional', optional=True),
+                io.Custom('TIMESTEP_KEYFRAME').Input('timestep_kf', optional=True),
+                io.Custom('LATENT_KEYFRAME').Input('latent_kf_override', optional=True),
+                io.Custom('CONTROL_NET_WEIGHTS').Input('weights_override', optional=True)
+            ],
+            outputs=[
+                io.Conditioning.Output('positive', is_output_list=False),
+                io.Conditioning.Output('negative', is_output_list=False)
+            ]
+        )
+
+    @classmethod
+    def execute(cls, positive, negative, control_net, vae, image, inpaint_mask, strength, start_percent, end_percent,
+                effect_mask_optional: Tensor=None, timestep_kf: TimestepKeyframeGroup=None,
+                latent_kf_override: LatentKeyframeGroup=None, weights_override: ControlWeights=None):
+        if not getattr(control_net, "concat_mask", False):
+            raise ValueError("The provided ControlNet does not use an inpaint source mask; use Apply Advanced ControlNet instead.")
+
+        source_mask = 1.0 - inpaint_mask.reshape((-1, 1, inpaint_mask.shape[-2], inpaint_mask.shape[-1]))
+        mask_apply = comfy.utils.common_upscale(source_mask, image.shape[2], image.shape[1], "bilinear", "center").round()
+        image = image * mask_apply.movedim(1, -1).repeat(1, 1, 1, image.shape[3])
+
+        return AdvancedControlNetApply.execute(
+            positive=positive,
+            negative=negative,
+            control_net=control_net,
+            image=image,
+            strength=strength,
+            start_percent=start_percent,
+            end_percent=end_percent,
+            mask_optional=effect_mask_optional,
+            vae_optional=vae,
+            timestep_kf=timestep_kf,
+            latent_kf_override=latent_kf_override,
+            weights_override=weights_override,
+            extra_concat=[source_mask]
+        )
+
+
 class AdvancedControlNetApplySingle(io.ComfyNode):
     @classmethod
     def define_schema(cls) -> io.Schema:

examples/qwen_image_inpainting/README.md

@@ -0,0 +1,84 @@
+# Qwen Image ControlNet inpainting
+
+This reviewer example adapts the active inpainting branch of ComfyUI's official
+Qwen Image workflow. It uses **Load Advanced ControlNet Model** and **Apply
+Advanced ControlNet Inpainting**, while retaining the official Qwen base
+pipeline and the bypassed optional Lightning LoRA. Node titles are left at
+their ComfyUI defaults; the workflow stores no node title overrides.
+
+## Inputs and models
+
+Download the official inputs to `ComfyUI/input` with these exact names:
+
+- [`acn_qwen_inpaint_source.png`](https://huggingface.co/InstantX/Qwen-Image-ControlNet-Inpainting/resolve/main/assets/images/image1.png)
+- [`acn_qwen_inpaint_mask.png`](https://huggingface.co/InstantX/Qwen-Image-ControlNet-Inpainting/resolve/main/assets/masks/mask1.png)
+
+The model author's repository is
+[`InstantX/Qwen-Image-ControlNet-Inpainting`](https://huggingface.co/InstantX/Qwen-Image-ControlNet-Inpainting).
+Download every model below to the listed folder under `ComfyUI/models`:
+
+| File and exact download | Folder |
+| --- | --- |
+| [`qwen_image_fp8_e4m3fn.safetensors`](https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/diffusion_models/qwen_image_fp8_e4m3fn.safetensors) | `diffusion_models` |
+| [`qwen_2.5_vl_7b_fp8_scaled.safetensors`](https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors) | `text_encoders` |
+| [`qwen_image_vae.safetensors`](https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/vae/qwen_image_vae.safetensors) | `vae` |
+| [`Qwen-Image-InstantX-ControlNet-Inpainting.safetensors`](https://huggingface.co/Comfy-Org/Qwen-Image-InstantX-ControlNets/resolve/main/split_files/controlnet/Qwen-Image-InstantX-ControlNet-Inpainting.safetensors) | `controlnet` |
+| [`Qwen-Image-Lightning-4steps-V1.0.safetensors`](https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Lightning-4steps-V1.0.safetensors) | `loras` (optional and bypassed) |
+
+## Run
+
+1. Download the two inputs and five model files to the folders above.
+2. Load `qwen_image_inpainting.json` in ComfyUI.
+3. Queue the workflow unchanged.
+
+For command-line input reproduction:
+
+```sh
+curl -L https://huggingface.co/InstantX/Qwen-Image-ControlNet-Inpainting/resolve/main/assets/images/image1.png -o ComfyUI/input/acn_qwen_inpaint_source.png
+curl -L https://huggingface.co/InstantX/Qwen-Image-ControlNet-Inpainting/resolve/main/assets/masks/mask1.png -o ComfyUI/input/acn_qwen_inpaint_mask.png
+```
+
+The unchanged example uses seed `134554158057228` (fixed), 20 steps, CFG 2.5,
+Euler, the simple scheduler, denoise 1.0, model shift 3.1, control strength 1.0,
+and control start/end 0.0/1.0. Its prompt is `The Queen, on a throne,
+surrounded by Knights, HD, Realistic, Octane Render, Unreal engine`; the
+negative prompt is one space. The source is scaled with area interpolation to a
+maximum dimension of 1536. The optional 4-step LoRA remains bypassed; enabling
+it requires changing the sampler settings appropriately.
+
+The two native **Load Image** nodes are intentionally separate. **Image To
+Mask** reads the red channel of the mask PNG. That source `inpaint_mask` defines
+the region supplied to the inpainting ControlNet and the latent noise mask. It
+is not the Advanced-ControlNet effect mask. `effect_mask_optional` is left
+unconnected and independently limits where control is injected. The Apply node
+also exposes unconnected timestep keyframe, latent keyframe, and weights ports
+for focused reviewer experiments.
+
+## Measured validation evidence
+
+These results were measured with fixed inputs and settings; they are recorded
+here rather than inferred from the example image:
+
+- A fresh isolated vanilla-versus-Advanced run had latent maximum/mean absolute
+  differences `0/0`, pixel maximum/mean differences `0/0`, and 0 changed
+  pixels.
+- An all-one effect mask exactly equaled unmasked Advanced output at latent and
+  pixel level. An all-zero effect mask exactly equaled no ControlNet at latent
+  and pixel level.
+- For right-half token-mask injection, relative to full control the left latent
+  mean delta was `0` and the right was `0.1332103`; relative to no control the
+  left was `0` and the right was `0.1835042`.
+- In a per-latent batch, the sample with strength 0 exactly equaled no control
+  at latent and pixel level.
+- Soft weights, timestep scheduling, and two-control stacking each executed
+  successfully.
+- The existing Anima real workflow rerun retained exact before/after latent and
+  pixel equality.
+
+Frontend and API validation for a missing source mask names `inpaint_mask`.
+Supplying an incompatible model produces this exact error:
+`The provided ControlNet does not use an inpaint source mask; use Apply Advanced
+ControlNet instead.`
+
+Workflow and result screenshots are linked from the PR instead of stored here
+to avoid repository growth.

tests/test_modern_control.py

@@ -0,0 +1,138 @@
+import os
+import sys
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch, sentinel
+
+comfyui_path = os.environ.get("COMFYUI_PATH")
+if comfyui_path:
+    sys.path.insert(0, comfyui_path)
+
+import torch
+
+from adv_control.control import ControlNetAdvanced
+from adv_control.nodes_main import AdvancedControlNetApply, AdvancedControlNetInpaintingApply
+
+
+class StopControlModel(Exception):
+    pass
+
+
+class ControlModel:
+    dtype = torch.float32
+
+    def __init__(self):
+        self.hint = None
+
+    def __call__(self, x, hint, timesteps, context, **kwargs):
+        self.hint = hint
+        raise StopControlModel
+
+
+class VideoVAE:
+    downscale_ratio = (4, 8, 8)
+
+    def __init__(self):
+        self.encoded_shape = None
+
+    def spacial_compression_encode(self):
+        return 8
+
+    def encode(self, image):
+        self.encoded_shape = image.shape
+        return torch.ones((image.shape[0], 4, 2, 2, 2))
+
+
+class ModernControlPreprocessingTests(unittest.TestCase):
+    def test_effect_mask_is_resized_to_qwen_tokens(self):
+        control = ControlNetAdvanced(ControlModel(), None)
+        control.x_noisy_shape = (1, 16, 4, 6)
+        control.mask_cond_hint = torch.tensor(
+            [[[[0.0, 0.0, 0.0, 1.0, 1.0, 1.0]] * 4]]
+        )
+        control.tk_mask_cond_hint = None
+        control.weights = SimpleNamespace(has_uncond_multiplier=False, has_uncond_mask=False)
+        control.latent_keyframes = None
+        control._current_timestep_keyframe = SimpleNamespace(strength=1.0)
+
+        output = torch.ones((1, 6, 4))
+        control.apply_advanced_strengths_and_masks(output, batched_number=1)
+
+        expected = torch.tensor(
+            [[[0.0] * 4, [0.5] * 4, [1.0] * 4, [0.0] * 4, [0.5] * 4, [1.0] * 4]]
+        )
+        torch.testing.assert_close(output, expected)
+
+    def test_vae_compression_and_source_mask_match_5d_hint(self):
+        control_model = ControlModel()
+        vae = VideoVAE()
+        control = ControlNetAdvanced(control_model, None, compression_ratio=1, latent_format=SimpleNamespace(process_in=lambda value: value))
+        control.real_compression_ratio = 1
+        control.cond_hint_original = torch.ones((1, 3, 16, 16))
+        control.cond_hint = None
+        control.vae = vae
+        control.extra_concat_orig = [torch.zeros((1, 1, 16, 16))]
+        control.sub_idxs = None
+        control.model_sampling_current = SimpleNamespace(timestep=lambda value: value, calculate_input=lambda timestep, value: value)
+        control.prepare_mask_cond_hint = lambda **kwargs: None
+
+        with self.assertRaises(StopControlModel):
+            control.sliding_get_control(
+                torch.ones((1, 4, 2, 2, 2)),
+                torch.ones(1),
+                {"c_crossattn": torch.ones((1, 1, 1))},
+                1,
+                {},
+            )
+
+        self.assertEqual(tuple(vae.encoded_shape), (1, 16, 16, 3))
+        self.assertEqual(tuple(control_model.hint.shape), (1, 5, 2, 2, 2))
+
+
+class AdvancedInpaintingApplyTests(unittest.TestCase):
+    def test_source_mask_and_effect_mask_stay_independent(self):
+        image = torch.ones((1, 2, 2, 3))
+        inpaint_mask = torch.tensor([[[1.0, 0.0], [1.0, 0.0]]])
+        effect_mask = torch.full((1, 2, 2), 0.25)
+        control_net = SimpleNamespace(concat_mask=True)
+
+        with patch.object(AdvancedControlNetApply, "execute", return_value=sentinel.output) as apply:
+            result = AdvancedControlNetInpaintingApply.execute(
+                positive=sentinel.positive,
+                negative=sentinel.negative,
+                control_net=control_net,
+                vae=sentinel.vae,
+                image=image,
+                inpaint_mask=inpaint_mask,
+                strength=1.0,
+                start_percent=0.0,
+                end_percent=1.0,
+                effect_mask_optional=effect_mask,
+            )
+
+        self.assertIs(result, sentinel.output)
+        inputs = apply.call_args.kwargs
+        torch.testing.assert_close(inputs["mask_optional"], effect_mask)
+        torch.testing.assert_close(inputs["extra_concat"][0], 1.0 - inpaint_mask.unsqueeze(1))
+        torch.testing.assert_close(
+            inputs["image"],
+            torch.tensor([[[[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]], [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]]]),
+        )
+
+    def test_non_inpaint_control_has_readable_error(self):
+        with self.assertRaisesRegex(ValueError, "does not use an inpaint source mask"):
+            AdvancedControlNetInpaintingApply.execute(
+                positive=[],
+                negative=[],
+                control_net=SimpleNamespace(concat_mask=False),
+                vae=sentinel.vae,
+                image=torch.ones((1, 2, 2, 3)),
+                inpaint_mask=torch.zeros((1, 2, 2)),
+                strength=1.0,
+                start_percent=0.0,
+                end_percent=1.0,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()