diff --git a/__init__.py b/__init__.py index 50576e9..08e968b 100644 --- a/__init__.py +++ b/__init__.py @@ -103,6 +103,7 @@ NODE_CONFIG = { "EmptyLatentImagePresets": {"class": EmptyLatentImagePresets, "name": "Empty Latent Image Presets"}, "EmptyLatentImageCustomPresets": {"class": EmptyLatentImageCustomPresets, "name": "Empty Latent Image Custom Presets"}, "ModelPassThrough": {"class": ModelPassThrough, "name": "ModelPass"}, + "SetShakkerLabsUnionControlNetType": {"class": SetShakkerLabsUnionControlNetType, "name": "Set Shakker Labs Union ControlNet Type"}, #audioscheduler stuff "NormalizedAmplitudeToMask": {"class": NormalizedAmplitudeToMask}, "NormalizedAmplitudeToFloatList": {"class": NormalizedAmplitudeToFloatList}, diff --git a/nodes/nodes.py b/nodes/nodes.py index e2eaf94..638f73b 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -2023,4 +2023,36 @@ class CustomControlNetWeightsFluxFromList: weights = ControlWeights.controlnet(weights_input=list_of_floats, uncond_multiplier=uncond_multiplier, extras=cn_extras) print(weights.weights_input) - return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) \ No newline at end of file + return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) + +SHAKKERLABS_UNION_CONTROLNET_TYPES = { + "canny": 0, + "tile": 1, + "depth": 2, + "blur": 3, + "pose": 4, + "gray": 5, + "low quality": 6, +} + +class SetShakkerLabsUnionControlNetType: + @classmethod + def INPUT_TYPES(s): + return {"required": {"control_net": ("CONTROL_NET", ), + "type": (["auto"] + list(SHAKKERLABS_UNION_CONTROLNET_TYPES.keys()),) + }} + + CATEGORY = "conditioning/controlnet" + RETURN_TYPES = ("CONTROL_NET",) + + FUNCTION = "set_controlnet_type" + + def set_controlnet_type(self, control_net, type): + control_net = control_net.copy() + type_number = SHAKKERLABS_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,) \ No newline at end of file