From 0e70221a05a64dedd07a0845ceebab8fb20eb1e5 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Fri, 17 Jul 2026 22:20:27 -0700 Subject: [PATCH] Add modern ControlNet inpainting support --- adv_control/control.py | 15 +- adv_control/nodes.py | 3 +- adv_control/nodes_main.py | 63 +++++++- examples/qwen_image_inpainting/README.md | 84 +++++++++++ .../qwen_image_inpainting.json | 1 + tests/test_modern_control.py | 138 ++++++++++++++++++ 6 files changed, 295 insertions(+), 9 deletions(-) create mode 100644 examples/qwen_image_inpainting/README.md create mode 100644 examples/qwen_image_inpainting/qwen_image_inpainting.json create mode 100644 tests/test_modern_control.py diff --git a/adv_control/control.py b/adv_control/control.py index 556b363..d6a5c60 100644 --- a/adv_control/control.py +++ b/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) diff --git a/adv_control/nodes.py b/adv_control/nodes.py index 565fbe6..2b2414f 100644 --- a/adv_control/nodes.py +++ b/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, diff --git a/adv_control/nodes_main.py b/adv_control/nodes_main.py index c705e69..a5bc2aa 100644 --- a/adv_control/nodes_main.py +++ b/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: diff --git a/examples/qwen_image_inpainting/README.md b/examples/qwen_image_inpainting/README.md new file mode 100644 index 0000000..2f3d885 --- /dev/null +++ b/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. diff --git a/examples/qwen_image_inpainting/qwen_image_inpainting.json b/examples/qwen_image_inpainting/qwen_image_inpainting.json new file mode 100644 index 0000000..a2ce6fa --- /dev/null +++ b/examples/qwen_image_inpainting/qwen_image_inpainting.json @@ -0,0 +1 @@ +{"id":"91f6bbe2-ed41-4fd6-bac7-71d5b5864ecb","revision":0,"last_node_id":225,"last_link_id":400,"nodes":[{"id":38,"type":"CLIPLoader","pos":[-90,70],"size":[380,106],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"CLIP","type":"CLIP","slot_index":0,"links":[74,75]}],"properties":{"Node name for S&R":"CLIPLoader","cnr_id":"comfy-core","ver":"0.3.51","models":[{"name":"qwen_2.5_vl_7b_fp8_scaled.safetensors","url":"https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors","directory":"text_encoders"}]},"widgets_values":["qwen_2.5_vl_7b_fp8_scaled.safetensors","qwen_image","default"]},{"id":37,"type":"UNETLoader","pos":[-90,-60],"size":[380,82],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"MODEL","type":"MODEL","slot_index":0,"links":[145]}],"properties":{"Node name for S&R":"UNETLoader","cnr_id":"comfy-core","ver":"0.3.51","models":[{"name":"qwen_image_fp8_e4m3fn.safetensors","url":"https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/diffusion_models/qwen_image_fp8_e4m3fn.safetensors","directory":"diffusion_models"}]},"widgets_values":["qwen_image_fp8_e4m3fn.safetensors","default"]},{"id":7,"type":"CLIPTextEncode","pos":[330,320],"size":[460,140],"flags":{},"order":18,"mode":0,"inputs":[{"name":"clip","type":"CLIP","link":75}],"outputs":[{"name":"CONDITIONING","type":"CONDITIONING","slot_index":0,"links":[191]}],"properties":{"Node name for S&R":"CLIPTextEncode"},"widgets_values":[" "]},{"id":84,"type":"ACN_ControlNetLoaderAdvanced","pos":[-90,340],"size":[380,58],"flags":{},"order":2,"mode":0,"inputs":[{"name":"_tk_opt","shape":7,"type":"TIMESTEP_KEYFRAME","link":null}],"outputs":[{"name":"CONTROL_NET","type":"CONTROL_NET","links":[192]}],"properties":{"Node name for S&R":"ACN_ControlNetLoaderAdvanced","cnr_id":"comfyui-advanced-controlnet","models":[{"name":"Qwen-Image-InstantX-ControlNet-Inpainting.safetensors","url":"https://huggingface.co/Comfy-Org/Qwen-Image-InstantX-ControlNets/resolve/main/split_files/controlnet/Qwen-Image-InstantX-ControlNet-Inpainting.safetensors","directory":"controlnet"}]},"widgets_values":["Qwen-Image-InstantX-ControlNet-Inpainting.safetensors"]},{"id":39,"type":"VAELoader","pos":[-90,230],"size":[380,58],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","slot_index":0,"links":[76,144,193]}],"properties":{"Node name for S&R":"VAELoader","cnr_id":"comfy-core","ver":"0.3.51","models":[{"name":"qwen_image_vae.safetensors","url":"https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/vae/qwen_image_vae.safetensors","directory":"vae"}]},"widgets_values":["qwen_image_vae.safetensors"]},{"id":66,"type":"ModelSamplingAuraFlow","pos":[860,-100],"size":[310,58],"flags":{},"order":26,"mode":0,"inputs":[{"name":"model","type":"MODEL","link":149}],"outputs":[{"name":"MODEL","type":"MODEL","links":[156]}],"properties":{"Node name for S&R":"ModelSamplingAuraFlow"},"widgets_values":[3.1]},{"id":108,"type":"ACN_AdvancedControlNetInpaintingApply","pos":[430,560],"size":[360,320],"flags":{},"order":30,"mode":0,"inputs":[{"name":"positive","type":"CONDITIONING","link":190},{"name":"negative","type":"CONDITIONING","link":191},{"name":"control_net","type":"CONTROL_NET","link":192},{"name":"vae","type":"VAE","link":193},{"name":"image","type":"IMAGE","link":294},{"name":"inpaint_mask","type":"MASK","link":220},{"name":"effect_mask_optional","shape":7,"type":"MASK","link":null},{"name":"timestep_kf","shape":7,"type":"TIMESTEP_KEYFRAME","link":null},{"name":"latent_kf_override","shape":7,"type":"LATENT_KEYFRAME","link":null},{"name":"weights_override","shape":7,"type":"CONTROL_NET_WEIGHTS","link":null}],"outputs":[{"name":"positive","type":"CONDITIONING","links":[188]},{"name":"negative","type":"CONDITIONING","links":[189]}],"properties":{"Node name for S&R":"ACN_AdvancedControlNetInpaintingApply"},"widgets_values":[1,0,1]},{"id":76,"type":"VAEEncode","pos":[430,900],"size":[140,46],"flags":{"collapsed":true},"order":28,"mode":0,"inputs":[{"name":"pixels","type":"IMAGE","link":293},{"name":"vae","type":"VAE","link":144}],"outputs":[{"name":"LATENT","type":"LATENT","links":[208]}],"properties":{"Node name for S&R":"VAEEncode"},"widgets_values":[]},{"id":122,"type":"SetLatentNoiseMask","pos":[600,900],"size":[230,50],"flags":{"collapsed":true},"order":32,"mode":0,"inputs":[{"name":"samples","type":"LATENT","link":208},{"name":"mask","type":"MASK","link":219}],"outputs":[{"name":"LATENT","type":"LATENT","links":[210]}],"properties":{"Node name for S&R":"SetLatentNoiseMask"},"widgets_values":[]},{"id":172,"type":"ImageScaleToMaxDimension","pos":[-20,1060],"size":[281,82],"flags":{},"order":23,"mode":0,"inputs":[{"name":"image","type":"IMAGE","link":292}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[293,294]}],"properties":{"Node name for S&R":"ImageScaleToMaxDimension"},"widgets_values":["area",1536]},{"id":8,"type":"VAEDecode","pos":[900,880],"size":[250,46],"flags":{},"order":39,"mode":0,"inputs":[{"name":"samples","type":"LATENT","link":128},{"name":"vae","type":"VAE","link":76}],"outputs":[{"name":"IMAGE","type":"IMAGE","slot_index":0,"links":[110]}],"properties":{"Node name for S&R":"VAEDecode"},"widgets_values":[]},{"id":71,"type":"LoadImage","pos":[-50,580],"size":[310,326],"flags":{},"order":11,"mode":0,"inputs":[],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[292]},{"name":"MASK","type":"MASK","links":[]}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["acn_qwen_inpaint_source.png","image"]},{"id":6,"type":"CLIPTextEncode","pos":[330,110],"size":[460,164],"flags":{},"order":17,"mode":0,"inputs":[{"name":"clip","type":"CLIP","link":74}],"outputs":[{"name":"CONDITIONING","type":"CONDITIONING","slot_index":0,"links":[190]}],"properties":{"Node name for S&R":"CLIPTextEncode"},"widgets_values":["The Queen, on a throne, surrounded by Knights, HD, Realistic, Octane Render, Unreal engine"]},{"id":80,"type":"LoraLoaderModelOnly","pos":[350,-70],"size":[430,82],"flags":{},"order":19,"mode":4,"inputs":[{"name":"model","type":"MODEL","link":145}],"outputs":[{"name":"MODEL","type":"MODEL","links":[149]}],"properties":{"Node name for S&R":"LoraLoaderModelOnly","cnr_id":"comfy-core","ver":"0.3.51","models":[{"name":"Qwen-Image-Lightning-4steps-V1.0.safetensors","url":"https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Lightning-4steps-V1.0.safetensors","directory":"loras"}]},"widgets_values":["Qwen-Image-Lightning-4steps-V1.0.safetensors",1]},{"id":3,"type":"KSampler","pos":[860,20],"size":[310,430],"flags":{},"order":36,"mode":0,"inputs":[{"name":"model","type":"MODEL","link":156},{"name":"positive","type":"CONDITIONING","link":188},{"name":"negative","type":"CONDITIONING","link":189},{"name":"latent_image","type":"LATENT","link":210}],"outputs":[{"name":"LATENT","type":"LATENT","slot_index":0,"links":[128]}],"properties":{"Node name for S&R":"KSampler"},"widgets_values":[134554158057228,"fixed",20,2.5,"euler","simple",1]},{"id":60,"type":"SaveImage","pos":[1280,-100],"size":[500,600],"flags":{},"order":41,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":110}],"outputs":[],"properties":{"Node name for S&R":"SaveImage"},"widgets_values":["acn_qwen_inpainting"]},{"id":224,"type":"LoadImage","pos":[-50,940],"size":[310,326],"flags":{},"order":12,"mode":0,"inputs":[],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[400]},{"name":"MASK","type":"MASK","links":[]}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["acn_qwen_inpaint_mask.png","image"]},{"id":225,"type":"ImageToMask","pos":[300,1000],"size":[220,58],"flags":{},"order":25,"mode":0,"inputs":[{"name":"image","type":"IMAGE","link":400}],"outputs":[{"name":"MASK","type":"MASK","links":[219,220]}],"properties":{"Node name for S&R":"ImageToMask"},"widgets_values":["red"]}],"links":[[74,38,0,6,0,"CLIP"],[75,38,0,7,0,"CLIP"],[76,39,0,8,1,"VAE"],[110,8,0,60,0,"IMAGE"],[128,3,0,8,0,"LATENT"],[144,39,0,76,1,"VAE"],[145,37,0,80,0,"MODEL"],[149,80,0,66,0,"MODEL"],[156,66,0,3,0,"MODEL"],[188,108,0,3,1,"CONDITIONING"],[189,108,1,3,2,"CONDITIONING"],[190,6,0,108,0,"CONDITIONING"],[191,7,0,108,1,"CONDITIONING"],[192,84,0,108,2,"CONTROL_NET"],[193,39,0,108,3,"VAE"],[208,76,0,122,0,"LATENT"],[210,122,0,3,3,"LATENT"],[292,71,0,172,0,"IMAGE"],[293,172,0,76,0,"IMAGE"],[294,172,0,108,4,"IMAGE"],[400,224,0,225,0,"IMAGE"],[219,225,0,122,1,"MASK"],[220,225,0,108,5,"MASK"]],"groups":[{"id":1,"bounding":[-120,-180,1900,1500],"color":"#3f789e","font_size":24,"flags":{}}],"definitions":{"subgraphs":[]},"config":{},"extra":{"ds":{"scale":0.5,"offset":[350,150]},"frontendVersion":"1.32.1"},"version":0.4} diff --git a/tests/test_modern_control.py b/tests/test_modern_control.py new file mode 100644 index 0000000..b4b173a --- /dev/null +++ b/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()