diff --git a/comfy_extras/v3/nodes_controlnet.py b/comfy_extras/v3/nodes_controlnet.py new file mode 100644 index 000000000..3dad4add6 --- /dev/null +++ b/comfy_extras/v3/nodes_controlnet.py @@ -0,0 +1,222 @@ +from comfy.cldm.control_types import UNION_CONTROLNET_TYPES +import comfy.utils +from comfy_api.v3 import io + + +class ControlNetApplyAdvanced_V3(io.ComfyNodeV3): + @classmethod + def DEFINE_SCHEMA(cls): + return io.SchemaV3( + node_id="ControlNetApplyAdvanced_V3", + display_name="Apply ControlNet _V3", + category="conditioning/controlnet", + inputs=[ + io.Conditioning.Input( + "positive", + display_name="positive", + ), + io.Conditioning.Input( + "negative", + display_name="negative", + ), + io.ControlNet.Input( + "control_net", + display_name="control_net", + ), + io.Image.Input( + "image", + display_name="image", + ), + io.Float.Input( + "strength", + display_name="strength", + default=1.0, + min=0.0, + max=10.0, + step=0.01, + ), + io.Float.Input( + "start_percent", + display_name="start percent", + default=0.0, + min=0.0, + max=1.0, + step=0.001, + ), + io.Float.Input( + "end_percent", + display_name="end percent", + default=1.0, + min=0.0, + max=1.0, + step=0.001, + ), + io.Vae.Input( + "vae", + optional=True, + ), + ], + outputs=[ + io.Conditioning.Output( + "positive_out", + display_name="positive", + ), + io.Conditioning.Output( + "negative_out", + display_name="negative", + ), + ], + ) + + @classmethod + def execute(cls, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None, extra_concat=[]): + if strength == 0: + return (positive, negative) + + control_hint = image.movedim(-1,1) + cnets = {} + + out = [] + for conditioning in [positive, negative]: + c = [] + for t in conditioning: + d = t[1].copy() + + prev_cnet = d.get('control', None) + if prev_cnet in cnets: + c_net = cnets[prev_cnet] + else: + c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae=vae, extra_concat=extra_concat) + c_net.set_previous_controlnet(prev_cnet) + cnets[prev_cnet] = c_net + + d['control'] = c_net + d['control_apply_to_uncond'] = False + n = [t[0], d] + c.append(n) + out.append(c) + return (out[0], out[1]) + + +class SetUnionControlNetType_V3(io.ComfyNodeV3): + @classmethod + def DEFINE_SCHEMA(cls): + return io.SchemaV3( + node_id="SetUnionControlNetType_V3", + category="conditioning/controlnet", + inputs=[ + io.ControlNet.Input( + "control_net", + display_name="control_net", + ), + io.Combo.Input( + "type", + options=["auto"] + list(UNION_CONTROLNET_TYPES.keys()), + ), + ], + outputs=[ + io.ControlNet.Output( + "control_net_out", + display_name="control_net", + ), + ], + ) + + @classmethod + def execute(cls, control_net, type): + control_net = control_net.copy() + type_number = UNION_CONTROLNET_TYPES.get(type, -1) + if type_number >= 0: + control_net.set_extra_arg("control_type", [type_number]) + else: + control_net.set_extra_arg("control_type", []) + + return (control_net,) + + +class ControlNetInpaintingAliMamaApply_V3(ControlNetApplyAdvanced_V3): + @classmethod + def DEFINE_SCHEMA(cls): + return io.SchemaV3( + node_id="ControlNetInpaintingAliMamaApply_V3", + category="conditioning/controlnet", + inputs=[ + io.Conditioning.Input( + "positive", + display_name="positive", + ), + io.Conditioning.Input( + "negative", + display_name="negative", + ), + io.ControlNet.Input( + "control_net", + display_name="control_net", + ), + io.Vae.Input( + "vae", + display_name="vae", + ), + io.Image.Input( + "image", + display_name="image", + ), + io.Mask.Input( + "mask", + display_name="mask", + ), + io.Float.Input( + "strength", + display_name="strength", + default=1.0, + min=0.0, + max=10.0, + step=0.01, + ), + io.Float.Input( + "start_percent", + display_name="start percent", + default=0.0, + min=0.0, + max=1.0, + step=0.001, + ), + io.Float.Input( + "end_percent", + display_name="end percent", + default=1.0, + min=0.0, + max=1.0, + step=0.001, + ), + + ], + outputs=[ + io.Conditioning.Output( + "positive_out", + display_name="positive", + ), + io.Conditioning.Output( + "negative_out", + display_name="negative", + ), + ], + ) + + @classmethod + def execute(cls, positive, negative, control_net, vae, image, mask, strength, start_percent, end_percent): + extra_concat = [] + if control_net.concat_mask: + mask = 1.0 - mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])) + mask_apply = comfy.utils.common_upscale(mask, image.shape[2], image.shape[1], "bilinear", "center").round() + image = image * mask_apply.movedim(1, -1).repeat(1, 1, 1, image.shape[3]) + extra_concat = [mask] + + return super().execute(positive, negative, control_net, image, strength, start_percent, end_percent, vae=vae, extra_concat=extra_concat) + + +NODES_LIST: list[type[io.ComfyNodeV3]] = [ + ControlNetApplyAdvanced_V3, + SetUnionControlNetType_V3, + ControlNetInpaintingAliMamaApply_V3, +] diff --git a/nodes.py b/nodes.py index 4e4dae917..428b32bea 100644 --- a/nodes.py +++ b/nodes.py @@ -2299,6 +2299,7 @@ def init_builtin_extra_nodes(): "nodes_tcfg.py", "nodes_v3_test.py", "nodes_v1_test.py", + "v3/nodes_controlnet.py", "v3/nodes_images.py", "v3/nodes_mask.py", "v3/nodes_webcam.py",