90 lines
3.6 KiB
Python
90 lines
3.6 KiB
Python
from comfy_api.latest import io
|
|
from torch import Tensor
|
|
import math
|
|
|
|
import folder_paths
|
|
|
|
from .control_plusplus import load_controlnetplusplus, PlusPlusInput, PlusPlusInputGroup, PlusPlusImageWrapper
|
|
|
|
class PlusPlusLoaderAdvanced(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id='ACN_ControlNet++LoaderAdvanced',
|
|
display_name='Load ControlNet++ Model (Multi) 🛂🅐🅒🅝',
|
|
category='Adv-ControlNet 🛂🅐🅒🅝/ControlNet++',
|
|
inputs=[
|
|
io.Custom('PLUS_INPUT').Input('plus_input'),
|
|
io.Combo.Input('name', options=folder_paths.get_filename_list("controlnet"))
|
|
],
|
|
outputs=[
|
|
io.ControlNet.Output('CONTROL_NET', is_output_list=False),
|
|
io.Image.Output('IMAGE', is_output_list=False)
|
|
]
|
|
)
|
|
|
|
|
|
@classmethod
|
|
def execute(cls, plus_input: PlusPlusInputGroup, name: str):
|
|
controlnet_path = folder_paths.get_full_path("controlnet", name)
|
|
controlnet = load_controlnetplusplus(controlnet_path)
|
|
controlnet.verify_control_type(name, plus_input)
|
|
controlnet.allow_condhint_latents = True
|
|
return io.NodeOutput(controlnet, PlusPlusImageWrapper(plus_input),)
|
|
|
|
class PlusPlusLoaderSingle(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id='ACN_ControlNet++LoaderSingle',
|
|
display_name='Load ControlNet++ Model (Single) 🛂🅐🅒🅝',
|
|
category='Adv-ControlNet 🛂🅐🅒🅝/ControlNet++',
|
|
inputs=[
|
|
io.Combo.Input('name', options=folder_paths.get_filename_list("controlnet")),
|
|
io.Combo.Input('control_type', options=['openpose', 'depth', 'hed/pidi/scribble/ted', 'canny/lineart/mlsd', 'normal', 'segment', 'tile', 'inpaint/outpaint', 'none'], default='none')
|
|
],
|
|
outputs=[
|
|
io.ControlNet.Output('CONTROL_NET', is_output_list=False)
|
|
]
|
|
)
|
|
|
|
|
|
@classmethod
|
|
def execute(cls, name: str, control_type: str):
|
|
controlnet_path = folder_paths.get_full_path("controlnet", name)
|
|
controlnet = load_controlnetplusplus(controlnet_path)
|
|
controlnet.single_control_type = control_type
|
|
controlnet.verify_control_type(name)
|
|
return io.NodeOutput(controlnet,)
|
|
|
|
class PlusPlusInputNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id='ACN_ControlNet++InputNode',
|
|
display_name='ControlNet++ Input 🛂🅐🅒🅝',
|
|
category='Adv-ControlNet 🛂🅐🅒🅝/ControlNet++',
|
|
inputs=[
|
|
io.Image.Input('image'),
|
|
io.Combo.Input('control_type', options=['openpose', 'depth', 'hed/pidi/scribble/ted', 'canny/lineart/mlsd', 'normal', 'segment', 'tile', 'inpaint/outpaint']),
|
|
io.Custom('PLUS_INPUT').Input('prev_plus_input', optional=True)
|
|
],
|
|
outputs=[
|
|
io.Custom('PLUS_INPUT').Output('PLUS_INPUT', is_output_list=False)
|
|
]
|
|
)
|
|
|
|
|
|
@classmethod
|
|
def execute(cls, image: Tensor, control_type: str, strength=1.0, prev_plus_input: PlusPlusInputGroup=None):
|
|
if prev_plus_input is None:
|
|
prev_plus_input = PlusPlusInputGroup()
|
|
prev_plus_input = prev_plus_input.clone()
|
|
|
|
if math.isclose(strength, 0.0):
|
|
strength = 0.0000001
|
|
pp_input = PlusPlusInput(image, control_type, strength)
|
|
prev_plus_input.add(pp_input)
|
|
|
|
return io.NodeOutput(prev_plus_input,)
|