mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-12-16 01:37:04 +08:00
convert nodes_controlnet.py to V3 schema (#10202)
This commit is contained in:
parent
b1467da480
commit
99ce2a1f66
@ -1,20 +1,26 @@
|
|||||||
from comfy.cldm.control_types import UNION_CONTROLNET_TYPES
|
from comfy.cldm.control_types import UNION_CONTROLNET_TYPES
|
||||||
import nodes
|
import nodes
|
||||||
import comfy.utils
|
import comfy.utils
|
||||||
|
from typing_extensions import override
|
||||||
|
from comfy_api.latest import ComfyExtension, io
|
||||||
|
|
||||||
class SetUnionControlNetType:
|
class SetUnionControlNetType(io.ComfyNode):
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(s):
|
def define_schema(cls):
|
||||||
return {"required": {"control_net": ("CONTROL_NET", ),
|
return io.Schema(
|
||||||
"type": (["auto"] + list(UNION_CONTROLNET_TYPES.keys()),)
|
node_id="SetUnionControlNetType",
|
||||||
}}
|
category="conditioning/controlnet",
|
||||||
|
inputs=[
|
||||||
|
io.ControlNet.Input("control_net"),
|
||||||
|
io.Combo.Input("type", options=["auto"] + list(UNION_CONTROLNET_TYPES.keys())),
|
||||||
|
],
|
||||||
|
outputs=[
|
||||||
|
io.ControlNet.Output(),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
CATEGORY = "conditioning/controlnet"
|
@classmethod
|
||||||
RETURN_TYPES = ("CONTROL_NET",)
|
def execute(cls, control_net, type) -> io.NodeOutput:
|
||||||
|
|
||||||
FUNCTION = "set_controlnet_type"
|
|
||||||
|
|
||||||
def set_controlnet_type(self, control_net, type):
|
|
||||||
control_net = control_net.copy()
|
control_net = control_net.copy()
|
||||||
type_number = UNION_CONTROLNET_TYPES.get(type, -1)
|
type_number = UNION_CONTROLNET_TYPES.get(type, -1)
|
||||||
if type_number >= 0:
|
if type_number >= 0:
|
||||||
@ -22,27 +28,36 @@ class SetUnionControlNetType:
|
|||||||
else:
|
else:
|
||||||
control_net.set_extra_arg("control_type", [])
|
control_net.set_extra_arg("control_type", [])
|
||||||
|
|
||||||
return (control_net,)
|
return io.NodeOutput(control_net)
|
||||||
|
|
||||||
class ControlNetInpaintingAliMamaApply(nodes.ControlNetApplyAdvanced):
|
set_controlnet_type = execute # TODO: remove
|
||||||
|
|
||||||
|
|
||||||
|
class ControlNetInpaintingAliMamaApply(io.ComfyNode):
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(s):
|
def define_schema(cls):
|
||||||
return {"required": {"positive": ("CONDITIONING", ),
|
return io.Schema(
|
||||||
"negative": ("CONDITIONING", ),
|
node_id="ControlNetInpaintingAliMamaApply",
|
||||||
"control_net": ("CONTROL_NET", ),
|
category="conditioning/controlnet",
|
||||||
"vae": ("VAE", ),
|
inputs=[
|
||||||
"image": ("IMAGE", ),
|
io.Conditioning.Input("positive"),
|
||||||
"mask": ("MASK", ),
|
io.Conditioning.Input("negative"),
|
||||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
io.ControlNet.Input("control_net"),
|
||||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
io.Vae.Input("vae"),
|
||||||
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001})
|
io.Image.Input("image"),
|
||||||
}}
|
io.Mask.Input("mask"),
|
||||||
|
io.Float.Input("strength", default=1.0, min=0.0, max=10.0, step=0.01),
|
||||||
|
io.Float.Input("start_percent", default=0.0, min=0.0, max=1.0, step=0.001),
|
||||||
|
io.Float.Input("end_percent", default=1.0, min=0.0, max=1.0, step=0.001),
|
||||||
|
],
|
||||||
|
outputs=[
|
||||||
|
io.Conditioning.Output(display_name="positive"),
|
||||||
|
io.Conditioning.Output(display_name="negative"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
FUNCTION = "apply_inpaint_controlnet"
|
@classmethod
|
||||||
|
def execute(cls, positive, negative, control_net, vae, image, mask, strength, start_percent, end_percent) -> io.NodeOutput:
|
||||||
CATEGORY = "conditioning/controlnet"
|
|
||||||
|
|
||||||
def apply_inpaint_controlnet(self, positive, negative, control_net, vae, image, mask, strength, start_percent, end_percent):
|
|
||||||
extra_concat = []
|
extra_concat = []
|
||||||
if control_net.concat_mask:
|
if control_net.concat_mask:
|
||||||
mask = 1.0 - mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
|
mask = 1.0 - mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
|
||||||
@ -50,11 +65,20 @@ class ControlNetInpaintingAliMamaApply(nodes.ControlNetApplyAdvanced):
|
|||||||
image = image * mask_apply.movedim(1, -1).repeat(1, 1, 1, image.shape[3])
|
image = image * mask_apply.movedim(1, -1).repeat(1, 1, 1, image.shape[3])
|
||||||
extra_concat = [mask]
|
extra_concat = [mask]
|
||||||
|
|
||||||
return self.apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent, vae=vae, extra_concat=extra_concat)
|
result = nodes.ControlNetApplyAdvanced().apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent, vae=vae, extra_concat=extra_concat)
|
||||||
|
return io.NodeOutput(result[0], result[1])
|
||||||
|
|
||||||
|
apply_inpaint_controlnet = execute # TODO: remove
|
||||||
|
|
||||||
|
|
||||||
|
class ControlNetExtension(ComfyExtension):
|
||||||
|
@override
|
||||||
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||||
|
return [
|
||||||
|
SetUnionControlNetType,
|
||||||
|
ControlNetInpaintingAliMamaApply,
|
||||||
|
]
|
||||||
|
|
||||||
NODE_CLASS_MAPPINGS = {
|
|
||||||
"SetUnionControlNetType": SetUnionControlNetType,
|
async def comfy_entrypoint() -> ControlNetExtension:
|
||||||
"ControlNetInpaintingAliMamaApply": ControlNetInpaintingAliMamaApply,
|
return ControlNetExtension()
|
||||||
}
|
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user