338 lines
13 KiB
Python
338 lines
13 KiB
Python
from typing import List
|
|
|
|
import folder_paths
|
|
from nodes import ControlNetLoader, ControlNetApply, ControlNetApplyAdvanced
|
|
|
|
from .preprocessor import preprocessors, apply_preprocessor
|
|
from .advanced import comfy_load_controlnet
|
|
|
|
|
|
def load_controlnet(control_net_name, control_net_override="None", timestep_keyframe=None):
|
|
if control_net_override != "None":
|
|
if control_net_override not in folder_paths.get_filename_list("controlnet"):
|
|
print(f"Warning: Not found ControlNet model {control_net_override}. Use {control_net_name} instead.")
|
|
else:
|
|
control_net_name = control_net_override
|
|
|
|
if control_net_name == "None":
|
|
return None
|
|
|
|
return comfy_load_controlnet(control_net_name, timestep_keyframe=timestep_keyframe)
|
|
|
|
|
|
def detect_controlnet(preprocessor: str, sd_version: str):
|
|
controlnets = folder_paths.get_filename_list("controlnet")
|
|
controlnets = filter(lambda x: sd_version in x, controlnets)
|
|
if sd_version == "sdxl":
|
|
controlnets = filter(lambda x: "t2i" not in x, controlnets)
|
|
controlnets = filter(lambda x: "lllite" not in x, controlnets)
|
|
|
|
control_net_name = "None"
|
|
if preprocessor in {"canny", "scribble", "mlsd"}:
|
|
control_net_name = next((c for c in controlnets if preprocessor in c), "None")
|
|
if preprocessor in {"scribble", "scribble_hed"}:
|
|
control_net_name = next((c for c in controlnets if "scribble" in c), "None")
|
|
if preprocessor in {"lineart", "lineart_coarse"}:
|
|
control_net_name = next((c for c in controlnets if "lineart." in c), "None")
|
|
if preprocessor in {"lineart_anime", "lineart_manga"}:
|
|
control_net_name = next((c for c in controlnets if "lineart_anime" in c), "None")
|
|
if preprocessor in {"hed", "hed_safe", "pidi", "pidi_safe"}:
|
|
control_net_name = next((c for c in controlnets if "softedge" in c), "None")
|
|
if preprocessor in {"pose", "openpose", "dwpose"}:
|
|
control_net_name = next((c for c in controlnets if "openpose" in c), "None")
|
|
if preprocessor in {"normalmap_bae", "normalmap_midas"}:
|
|
control_net_name = next((c for c in controlnets if "normalbae" in c), "None")
|
|
if preprocessor in {"depth", "depth_midas", "depth_zoe"}:
|
|
control_net_name = next((c for c in controlnets if "depth" in c), "None")
|
|
if preprocessor in {"seg_ofcoco", "seg_ofade20k", "seg_ufade20k"}:
|
|
control_net_name = next((c for c in controlnets if "seg" in c), "None")
|
|
|
|
if preprocessor in {"tile"}:
|
|
control_net_name = next((c for c in controlnets if "tile" in c), "None")
|
|
|
|
return control_net_name
|
|
|
|
|
|
class AVControlNetLoader(ControlNetLoader):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {"control_net_name": (folder_paths.get_filename_list("controlnet"),)},
|
|
"optional": {
|
|
"control_net_override": ("STRING", {"default": "None"}),
|
|
"timestep_keyframe": ("TIMESTEP_KEYFRAME",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONTROL_NET",)
|
|
FUNCTION = "load_controlnet"
|
|
CATEGORY = "Art Venture/Loaders"
|
|
|
|
def load_controlnet(self, control_net_name, control_net_override="None", timestep_keyframe=None):
|
|
return load_controlnet(control_net_name, control_net_override, timestep_keyframe=timestep_keyframe)
|
|
|
|
|
|
class AV_ControlNetPreprocessor:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"preprocessor": (["None"] + preprocessors,),
|
|
"sd_version": (["sd15", "sdxl"],),
|
|
},
|
|
"optional": {
|
|
"resolution": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64}),
|
|
"preprocessor_override": ("STRING", {"default": "None"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("IMAGE", "CNET_NAME")
|
|
FUNCTION = "detect_controlnet"
|
|
CATEGORY = "Art Venture/Loaders"
|
|
|
|
def detect_controlnet(self, image, preprocessor, sd_version, resolution=512, preprocessor_override="None"):
|
|
if preprocessor_override != "None":
|
|
if preprocessor_override not in preprocessors:
|
|
print(
|
|
f"Warning: Not found ControlNet preprocessor {preprocessor_override}. Use {preprocessor} instead."
|
|
)
|
|
else:
|
|
preprocessor = preprocessor_override
|
|
|
|
image = apply_preprocessor(image, preprocessor, resolution=resolution)
|
|
control_net_name = detect_controlnet(preprocessor, sd_version)
|
|
|
|
return (image, control_net_name)
|
|
|
|
|
|
class AVControlNetEfficientStacker:
|
|
controlnets = folder_paths.get_filename_list("controlnet")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"control_net_name": (["None", "Auto: sd15", "Auto: sdxl", "Auto: sdxl_t2i"] + s.controlnets,),
|
|
"image": ("IMAGE",),
|
|
"strength": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01},
|
|
),
|
|
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"preprocessor": (["None"] + preprocessors,),
|
|
},
|
|
"optional": {
|
|
"cnet_stack": ("CONTROL_NET_STACK",),
|
|
"control_net_override": ("STRING", {"default": "None"}),
|
|
"timestep_keyframe": ("TIMESTEP_KEYFRAME",),
|
|
"resolution": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64}),
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONTROL_NET_STACK",)
|
|
RETURN_NAMES = ("CNET_STACK",)
|
|
FUNCTION = "control_net_stacker"
|
|
CATEGORY = "Art Venture/Loaders"
|
|
|
|
def control_net_stacker(
|
|
self,
|
|
control_net_name: str,
|
|
image,
|
|
strength,
|
|
start_percent,
|
|
end_percent,
|
|
preprocessor,
|
|
cnet_stack=None,
|
|
control_net_override="None",
|
|
timestep_keyframe=None,
|
|
resolution=512,
|
|
enabled=True,
|
|
):
|
|
if not enabled:
|
|
return (cnet_stack,)
|
|
|
|
# If control_net_stack is None, initialize as an empty list
|
|
if cnet_stack is None:
|
|
cnet_stack = []
|
|
|
|
if control_net_name.startswith("Auto: "):
|
|
assert preprocessor != "None", "preprocessor must be set when using Auto mode"
|
|
|
|
sd_version = control_net_name[len("Auto: ") :]
|
|
control_net_name = detect_controlnet(preprocessor, sd_version)
|
|
|
|
control_net = load_controlnet(control_net_name, control_net_override, timestep_keyframe=timestep_keyframe)
|
|
|
|
# Extend the control_net_stack with the new tuple
|
|
if control_net is not None:
|
|
image = apply_preprocessor(image, preprocessor, resolution=resolution)
|
|
cnet_stack.extend([(control_net, image, strength, start_percent, end_percent)])
|
|
|
|
return (cnet_stack,)
|
|
|
|
|
|
class AVControlNetEfficientStackerSimple(AVControlNetEfficientStacker):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"control_net_name": (["None", "Auto: sd15", "Auto: sdxl", "Auto: sdxl_t2i"] + s.controlnets,),
|
|
"image": ("IMAGE",),
|
|
"strength": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01},
|
|
),
|
|
"preprocessor": (["None"] + preprocessors,),
|
|
},
|
|
"optional": {
|
|
"cnet_stack": ("CONTROL_NET_STACK",),
|
|
"control_net_override": ("STRING", {"default": "None"}),
|
|
"timestep_keyframe": ("TIMESTEP_KEYFRAME",),
|
|
"resolution": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64}),
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
FUNCTION = "control_net_stacker_simple"
|
|
|
|
def control_net_stacker_simple(
|
|
self,
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
return self.control_net_stacker(*args, start_percent=0.0, end_percent=1.0, **kwargs)
|
|
|
|
|
|
class AVControlNetEfficientLoader(ControlNetApply):
|
|
controlnets = folder_paths.get_filename_list("controlnet")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"control_net_name": (["None"] + s.controlnets,),
|
|
"conditioning": ("CONDITIONING",),
|
|
"image": ("IMAGE",),
|
|
"strength": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01},
|
|
),
|
|
"preprocessor": (["None"] + preprocessors,),
|
|
},
|
|
"optional": {
|
|
"control_net_override": ("STRING", {"default": "None"}),
|
|
"timestep_keyframe": ("TIMESTEP_KEYFRAME",),
|
|
"resolution": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64}),
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "load_controlnet"
|
|
CATEGORY = "Art Venture/Loaders"
|
|
|
|
def load_controlnet(
|
|
self,
|
|
control_net_name,
|
|
conditioning,
|
|
image,
|
|
strength,
|
|
preprocessor,
|
|
control_net_override="None",
|
|
timestep_keyframe=None,
|
|
resolution=512,
|
|
enabled=True,
|
|
):
|
|
if not enabled:
|
|
return (conditioning,)
|
|
|
|
control_net = load_controlnet(control_net_name, control_net_override, timestep_keyframe=timestep_keyframe)
|
|
if control_net is None:
|
|
return (conditioning,)
|
|
|
|
image = apply_preprocessor(image, preprocessor, resolution=resolution)
|
|
|
|
return super().apply_controlnet(conditioning, control_net, image, strength)
|
|
|
|
|
|
class AVControlNetEfficientLoaderAdvanced(ControlNetApplyAdvanced):
|
|
controlnets = folder_paths.get_filename_list("controlnet")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"control_net_name": (["None"] + s.controlnets,),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"image": ("IMAGE",),
|
|
"strength": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01},
|
|
),
|
|
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"preprocessor": (["None"] + preprocessors,),
|
|
},
|
|
"optional": {
|
|
"control_net_override": ("STRING", {"default": "None"}),
|
|
"timestep_keyframe": ("TIMESTEP_KEYFRAME",),
|
|
"resolution": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64}),
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("positive", "negative")
|
|
FUNCTION = "load_controlnet"
|
|
CATEGORY = "Art Venture/Loaders"
|
|
|
|
def load_controlnet(
|
|
self,
|
|
control_net_name,
|
|
positive,
|
|
negative,
|
|
image,
|
|
strength,
|
|
start_percent,
|
|
end_percent,
|
|
preprocessor,
|
|
control_net_override="None",
|
|
timestep_keyframe=None,
|
|
resolution=512,
|
|
enabled=True,
|
|
):
|
|
if not enabled:
|
|
return (positive, negative)
|
|
|
|
control_net = load_controlnet(control_net_name, control_net_override, timestep_keyframe=timestep_keyframe)
|
|
if control_net is None:
|
|
return (positive, negative)
|
|
|
|
image = apply_preprocessor(image, preprocessor, resolution=resolution)
|
|
|
|
return super().apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AV_ControlNetLoader": AVControlNetLoader,
|
|
"AV_ControlNetEfficientLoader": AVControlNetEfficientLoader,
|
|
"AV_ControlNetEfficientLoaderAdvanced": AVControlNetEfficientLoaderAdvanced,
|
|
"AV_ControlNetEfficientStacker": AVControlNetEfficientStacker,
|
|
"AV_ControlNetEfficientStackerSimple": AVControlNetEfficientStackerSimple,
|
|
"AV_ControlNetPreprocessor": AV_ControlNetPreprocessor,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AV_ControlNetLoader": "ControlNet Loader",
|
|
"AV_ControlNetEfficientLoader": "ControlNet Loader",
|
|
"AV_ControlNetEfficientLoaderAdvanced": "ControlNet Loader Adv.",
|
|
"AV_ControlNetEfficientStacker": "ControlNet Stacker Adv.",
|
|
"AV_ControlNetEfficientStackerSimple": "ControlNet Stacker",
|
|
"AV_ControlNetPreprocessor": "ControlNet Preprocessor",
|
|
}
|