nodes_controlnet.py 867 Bytes
Newer Older
1
from comfy.cldm.control_types import UNION_CONTROLNET_TYPES
2
3
4
5
6

class SetUnionControlNetType:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"control_net": ("CONTROL_NET", ),
7
                             "type": (["auto"] + list(UNION_CONTROLNET_TYPES.keys()),)
8
9
                             }}

10
    CATEGORY = "conditioning/controlnet"
11
12
13
14
15
16
    RETURN_TYPES = ("CONTROL_NET",)

    FUNCTION = "set_controlnet_type"

    def set_controlnet_type(self, control_net, type):
        control_net = control_net.copy()
17
        type_number = UNION_CONTROLNET_TYPES.get(type, -1)
18
19
20
21
22
23
24
25
26
27
        if type_number >= 0:
            control_net.set_extra_arg("control_type", [type_number])
        else:
            control_net.set_extra_arg("control_type", [])

        return (control_net,)

NODE_CLASS_MAPPINGS = {
    "SetUnionControlNetType": SetUnionControlNetType,
}