Add custom nodes, Civitai loras (LFS), and vast.ai setup script
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Has been cancelled
Execution Tests / test (macos-latest) (push) Has been cancelled
Execution Tests / test (ubuntu-latest) (push) Has been cancelled
Execution Tests / test (windows-latest) (push) Has been cancelled
Test server launches without errors / test (push) Has been cancelled
Unit Tests / test (macos-latest) (push) Has been cancelled
Unit Tests / test (ubuntu-latest) (push) Has been cancelled
Unit Tests / test (windows-2022) (push) Has been cancelled
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Has been cancelled
Execution Tests / test (macos-latest) (push) Has been cancelled
Execution Tests / test (ubuntu-latest) (push) Has been cancelled
Execution Tests / test (windows-latest) (push) Has been cancelled
Test server launches without errors / test (push) Has been cancelled
Unit Tests / test (macos-latest) (push) Has been cancelled
Unit Tests / test (ubuntu-latest) (push) Has been cancelled
Unit Tests / test (windows-2022) (push) Has been cancelled
Includes 30 custom nodes committed directly, 7 Civitai-exclusive loras stored via Git LFS, and a setup script that installs all dependencies and downloads HuggingFace-hosted models on vast.ai. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,43 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
class AnimeFace_SemSegPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
#This preprocessor is only trained on 512x resolution
|
||||
#https://github.com/siyeong0/Anime-Face-Segmentation/blob/main/predict.py#L25
|
||||
return define_preprocessor_inputs(
|
||||
remove_background_using_abg=INPUT.BOOLEAN(True),
|
||||
resolution=INPUT.RESOLUTION(default=512, min=512, max=512)
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
RETURN_NAMES = ("IMAGE", "ABG_CHARACTER_MASK (MASK)")
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Semantic Segmentation"
|
||||
|
||||
def execute(self, image, remove_background_using_abg=True, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.anime_face_segment import AnimeFaceSegmentor
|
||||
|
||||
model = AnimeFaceSegmentor.from_pretrained().to(model_management.get_torch_device())
|
||||
if remove_background_using_abg:
|
||||
out_image_with_mask = common_annotator_call(model, image, resolution=resolution, remove_background=True)
|
||||
out_image = out_image_with_mask[..., :3]
|
||||
mask = out_image_with_mask[..., 3:]
|
||||
mask = rearrange(mask, "n h w c -> n c h w")
|
||||
else:
|
||||
out_image = common_annotator_call(model, image, resolution=resolution, remove_background=False)
|
||||
N, H, W, C = out_image.shape
|
||||
mask = torch.ones(N, C, H, W)
|
||||
del model
|
||||
return (out_image, mask)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AnimeFace_SemSegPreprocessor": AnimeFace_SemSegPreprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AnimeFace_SemSegPreprocessor": "Anime Face Segmentor"
|
||||
}
|
||||
87
custom_nodes/comfyui_controlnet_aux/node_wrappers/anyline.py
Normal file
87
custom_nodes/comfyui_controlnet_aux/node_wrappers/anyline.py
Normal file
@@ -0,0 +1,87 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import comfy.model_management as model_management
|
||||
import comfy.utils
|
||||
|
||||
# Requires comfyui_controlnet_aux funcsions and classes
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
|
||||
def get_intensity_mask(image_array, lower_bound, upper_bound):
|
||||
mask = image_array[:, :, 0]
|
||||
mask = np.where((mask >= lower_bound) & (mask <= upper_bound), mask, 0)
|
||||
mask = np.expand_dims(mask, 2).repeat(3, axis=2)
|
||||
return mask
|
||||
|
||||
def combine_layers(base_layer, top_layer):
|
||||
mask = top_layer.astype(bool)
|
||||
temp = 1 - (1 - top_layer) * (1 - base_layer)
|
||||
result = base_layer * (~mask) + temp * mask
|
||||
return result
|
||||
|
||||
class AnyLinePreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
merge_with_lineart=INPUT.COMBO(["lineart_standard", "lineart_realisitic", "lineart_anime", "manga_line"], default="lineart_standard"),
|
||||
resolution=INPUT.RESOLUTION(default=1280, step=8),
|
||||
lineart_lower_bound=INPUT.FLOAT(default=0),
|
||||
lineart_upper_bound=INPUT.FLOAT(default=1),
|
||||
object_min_size=INPUT.INT(default=36, min=1),
|
||||
object_connectivity=INPUT.INT(default=1, min=1)
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
|
||||
FUNCTION = "get_anyline"
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def __init__(self):
|
||||
self.device = model_management.get_torch_device()
|
||||
|
||||
def get_anyline(self, image, merge_with_lineart="lineart_standard", resolution=512, lineart_lower_bound=0, lineart_upper_bound=1, object_min_size=36, object_connectivity=1):
|
||||
from custom_controlnet_aux.teed import TEDDetector
|
||||
from skimage import morphology
|
||||
pbar = comfy.utils.ProgressBar(3)
|
||||
|
||||
# Process the image with MTEED model
|
||||
mteed_model = TEDDetector.from_pretrained("TheMistoAI/MistoLine", "MTEED.pth", subfolder="Anyline").to(self.device)
|
||||
mteed_result = common_annotator_call(mteed_model, image, resolution=resolution, show_pbar=False)
|
||||
mteed_result = mteed_result.numpy()
|
||||
del mteed_model
|
||||
pbar.update(1)
|
||||
|
||||
# Process the image with the lineart standard preprocessor
|
||||
if merge_with_lineart == "lineart_standard":
|
||||
from custom_controlnet_aux.lineart_standard import LineartStandardDetector
|
||||
lineart_standard_detector = LineartStandardDetector()
|
||||
lineart_result = common_annotator_call(lineart_standard_detector, image, guassian_sigma=2, intensity_threshold=3, resolution=resolution, show_pbar=False).numpy()
|
||||
del lineart_standard_detector
|
||||
else:
|
||||
from custom_controlnet_aux.lineart import LineartDetector
|
||||
from custom_controlnet_aux.lineart_anime import LineartAnimeDetector
|
||||
from custom_controlnet_aux.manga_line import LineartMangaDetector
|
||||
lineart_detector = dict(lineart_realisitic=LineartDetector, lineart_anime=LineartAnimeDetector, manga_line=LineartMangaDetector)[merge_with_lineart]
|
||||
lineart_detector = lineart_detector.from_pretrained().to(self.device)
|
||||
lineart_result = common_annotator_call(lineart_detector, image, resolution=resolution, show_pbar=False).numpy()
|
||||
del lineart_detector
|
||||
pbar.update(1)
|
||||
|
||||
final_result = []
|
||||
for i in range(len(image)):
|
||||
_lineart_result = get_intensity_mask(lineart_result[i], lower_bound=lineart_lower_bound, upper_bound=lineart_upper_bound)
|
||||
_cleaned = morphology.remove_small_objects(_lineart_result.astype(bool), min_size=object_min_size, connectivity=object_connectivity)
|
||||
_lineart_result = _lineart_result * _cleaned
|
||||
_mteed_result = mteed_result[i]
|
||||
|
||||
# Combine the results
|
||||
final_result.append(torch.from_numpy(combine_layers(_mteed_result, _lineart_result)))
|
||||
pbar.update(1)
|
||||
return (torch.stack(final_result),)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AnyLineArtPreprocessor_aux": AnyLinePreprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AnyLineArtPreprocessor_aux": "AnyLine Lineart"
|
||||
}
|
||||
29
custom_nodes/comfyui_controlnet_aux/node_wrappers/binary.py
Normal file
29
custom_nodes/comfyui_controlnet_aux/node_wrappers/binary.py
Normal file
@@ -0,0 +1,29 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Binary_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
bin_threshold=INPUT.INT(default=100, max=255),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, bin_threshold=100, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.binary import BinaryDetector
|
||||
|
||||
return (common_annotator_call(BinaryDetector(), image, bin_threshold=bin_threshold, resolution=resolution), )
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"BinaryPreprocessor": Binary_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"BinaryPreprocessor": "Binary Lines"
|
||||
}
|
||||
30
custom_nodes/comfyui_controlnet_aux/node_wrappers/canny.py
Normal file
30
custom_nodes/comfyui_controlnet_aux/node_wrappers/canny.py
Normal file
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Canny_Edge_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
low_threshold=INPUT.INT(default=100, max=255),
|
||||
high_threshold=INPUT.INT(default=200, max=255),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, low_threshold=100, high_threshold=200, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.canny import CannyDetector
|
||||
|
||||
return (common_annotator_call(CannyDetector(), image, low_threshold=low_threshold, high_threshold=high_threshold, resolution=resolution), )
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CannyEdgePreprocessor": Canny_Edge_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CannyEdgePreprocessor": "Canny Edge"
|
||||
}
|
||||
26
custom_nodes/comfyui_controlnet_aux/node_wrappers/color.py
Normal file
26
custom_nodes/comfyui_controlnet_aux/node_wrappers/color.py
Normal file
@@ -0,0 +1,26 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Color_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/T2IAdapter-only"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.color import ColorDetector
|
||||
|
||||
return (common_annotator_call(ColorDetector(), image, resolution=resolution), )
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ColorPreprocessor": Color_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ColorPreprocessor": "Color Pallete"
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class DensePose_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
model=INPUT.COMBO(["densepose_r50_fpn_dl.torchscript", "densepose_r101_fpn_dl.torchscript"]),
|
||||
cmap=INPUT.COMBO(["Viridis (MagicAnimate)", "Parula (CivitAI)"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def execute(self, image, model="densepose_r50_fpn_dl.torchscript", cmap="Viridis (MagicAnimate)", resolution=512):
|
||||
from custom_controlnet_aux.densepose import DenseposeDetector
|
||||
model = DenseposeDetector \
|
||||
.from_pretrained(filename=model) \
|
||||
.to(model_management.get_torch_device())
|
||||
return (common_annotator_call(model, image, cmap="viridis" if "Viridis" in cmap else "parula", resolution=resolution), )
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DensePosePreprocessor": DensePose_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DensePosePreprocessor": "DensePose Estimator"
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Depth_Anything_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
ckpt_name=INPUT.COMBO(
|
||||
["depth_anything_vitl14.pth", "depth_anything_vitb14.pth", "depth_anything_vits14.pth"]
|
||||
),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, ckpt_name="depth_anything_vitl14.pth", resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.depth_anything import DepthAnythingDetector
|
||||
|
||||
model = DepthAnythingDetector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
class Zoe_Depth_Anything_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
environment=INPUT.COMBO(["indoor", "outdoor"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, environment="indoor", resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.zoe import ZoeDepthAnythingDetector
|
||||
ckpt_name = "depth_anything_metric_depth_indoor.pt" if environment == "indoor" else "depth_anything_metric_depth_outdoor.pt"
|
||||
model = ZoeDepthAnythingDetector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DepthAnythingPreprocessor": Depth_Anything_Preprocessor,
|
||||
"Zoe_DepthAnythingPreprocessor": Zoe_Depth_Anything_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DepthAnythingPreprocessor": "Depth Anything",
|
||||
"Zoe_DepthAnythingPreprocessor": "Zoe Depth Anything"
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Depth_Anything_V2_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
ckpt_name=INPUT.COMBO(
|
||||
["depth_anything_v2_vitg.pth", "depth_anything_v2_vitl.pth", "depth_anything_v2_vitb.pth", "depth_anything_v2_vits.pth"],
|
||||
default="depth_anything_v2_vitl.pth"
|
||||
),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, ckpt_name="depth_anything_v2_vitl.pth", resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.depth_anything_v2 import DepthAnythingV2Detector
|
||||
|
||||
model = DepthAnythingV2Detector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, max_depth=1)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
""" class Depth_Anything_Metric_V2_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return create_node_input_types(
|
||||
environment=(["indoor", "outdoor"], {"default": "indoor"}),
|
||||
max_depth=("FLOAT", {"min": 0, "max": 100, "default": 20.0, "step": 0.01})
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, environment, resolution=512, max_depth=20.0, **kwargs):
|
||||
from custom_controlnet_aux.depth_anything_v2 import DepthAnythingV2Detector
|
||||
filename = dict(indoor="depth_anything_v2_metric_hypersim_vitl.pth", outdoor="depth_anything_v2_metric_vkitti_vitl.pth")[environment]
|
||||
model = DepthAnythingV2Detector.from_pretrained(filename=filename).to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, max_depth=max_depth)
|
||||
del model
|
||||
return (out, ) """
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DepthAnythingV2Preprocessor": Depth_Anything_V2_Preprocessor,
|
||||
#"Metric_DepthAnythingV2Preprocessor": Depth_Anything_Metric_V2_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DepthAnythingV2Preprocessor": "Depth Anything V2 - Relative",
|
||||
#"Metric_DepthAnythingV2Preprocessor": "Depth Anything V2 - Metric"
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, run_script
|
||||
import comfy.model_management as model_management
|
||||
import sys
|
||||
|
||||
def install_deps():
|
||||
try:
|
||||
import sklearn
|
||||
except:
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', 'scikit-learn'])
|
||||
|
||||
class DiffusionEdge_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
environment=INPUT.COMBO(["indoor", "urban", "natrual"]),
|
||||
patch_batch_size=INPUT.INT(default=4, min=1, max=16),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, environment="indoor", patch_batch_size=4, resolution=512, **kwargs):
|
||||
install_deps()
|
||||
from custom_controlnet_aux.diffusion_edge import DiffusionEdgeDetector
|
||||
|
||||
model = DiffusionEdgeDetector \
|
||||
.from_pretrained(filename = f"diffusion_edge_{environment}.pt") \
|
||||
.to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, patch_batch_size=patch_batch_size)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DiffusionEdge_Preprocessor": DiffusionEdge_Preprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DiffusionEdge_Preprocessor": "Diffusion Edge (batch size ↑ => speed ↑, VRAM ↑)",
|
||||
}
|
||||
31
custom_nodes/comfyui_controlnet_aux/node_wrappers/dsine.py
Normal file
31
custom_nodes/comfyui_controlnet_aux/node_wrappers/dsine.py
Normal file
@@ -0,0 +1,31 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class DSINE_Normal_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
fov=INPUT.FLOAT(max=365.0, default=60.0),
|
||||
iterations=INPUT.INT(min=1, max=20, default=5),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, fov=60.0, iterations=5, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.dsine import DsineDetector
|
||||
|
||||
model = DsineDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, fov=fov, iterations=iterations, resolution=resolution)
|
||||
del model
|
||||
return (out,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DSINE-NormalMapPreprocessor": DSINE_Normal_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DSINE-NormalMapPreprocessor": "DSINE Normal Map"
|
||||
}
|
||||
166
custom_nodes/comfyui_controlnet_aux/node_wrappers/dwpose.py
Normal file
166
custom_nodes/comfyui_controlnet_aux/node_wrappers/dwpose.py
Normal file
@@ -0,0 +1,166 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import numpy as np
|
||||
import warnings
|
||||
from ..src.custom_controlnet_aux.dwpose import DwposeDetector, AnimalposeDetector
|
||||
import os
|
||||
import json
|
||||
|
||||
DWPOSE_MODEL_NAME = "yzd-v/DWPose"
|
||||
#Trigger startup caching for onnxruntime
|
||||
GPU_PROVIDERS = ["CUDAExecutionProvider", "DirectMLExecutionProvider", "OpenVINOExecutionProvider", "ROCMExecutionProvider", "CoreMLExecutionProvider"]
|
||||
def check_ort_gpu():
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
for provider in GPU_PROVIDERS:
|
||||
if provider in ort.get_available_providers():
|
||||
return True
|
||||
return False
|
||||
except:
|
||||
return False
|
||||
|
||||
if not os.environ.get("DWPOSE_ONNXRT_CHECKED"):
|
||||
if check_ort_gpu():
|
||||
print("DWPose: Onnxruntime with acceleration providers detected")
|
||||
else:
|
||||
warnings.warn("DWPose: Onnxruntime not found or doesn't come with acceleration providers, switch to OpenCV with CPU device. DWPose might run very slowly")
|
||||
os.environ['AUX_ORT_PROVIDERS'] = ''
|
||||
os.environ["DWPOSE_ONNXRT_CHECKED"] = '1'
|
||||
|
||||
class DWPose_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
detect_hand=INPUT.COMBO(["enable", "disable"]),
|
||||
detect_body=INPUT.COMBO(["enable", "disable"]),
|
||||
detect_face=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION(),
|
||||
bbox_detector=INPUT.COMBO(
|
||||
["None"] + ["yolox_l.torchscript.pt", "yolox_l.onnx", "yolo_nas_l_fp16.onnx", "yolo_nas_m_fp16.onnx", "yolo_nas_s_fp16.onnx"],
|
||||
default="yolox_l.onnx"
|
||||
),
|
||||
pose_estimator=INPUT.COMBO(
|
||||
["dw-ll_ucoco_384_bs5.torchscript.pt", "dw-ll_ucoco_384.onnx", "dw-ll_ucoco.onnx"],
|
||||
default="dw-ll_ucoco_384_bs5.torchscript.pt"
|
||||
),
|
||||
scale_stick_for_xinsr_cn=INPUT.COMBO(["disable", "enable"])
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
|
||||
FUNCTION = "estimate_pose"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def estimate_pose(self, image, detect_hand="enable", detect_body="enable", detect_face="enable", resolution=512, bbox_detector="yolox_l.onnx", pose_estimator="dw-ll_ucoco_384.onnx", scale_stick_for_xinsr_cn="disable", **kwargs):
|
||||
if bbox_detector == "None":
|
||||
yolo_repo = DWPOSE_MODEL_NAME
|
||||
elif bbox_detector == "yolox_l.onnx":
|
||||
yolo_repo = DWPOSE_MODEL_NAME
|
||||
elif "yolox" in bbox_detector:
|
||||
yolo_repo = "hr16/yolox-onnx"
|
||||
elif "yolo_nas" in bbox_detector:
|
||||
yolo_repo = "hr16/yolo-nas-fp16"
|
||||
else:
|
||||
raise NotImplementedError(f"Download mechanism for {bbox_detector}")
|
||||
|
||||
if pose_estimator == "dw-ll_ucoco_384.onnx":
|
||||
pose_repo = DWPOSE_MODEL_NAME
|
||||
elif pose_estimator.endswith(".onnx"):
|
||||
pose_repo = "hr16/UnJIT-DWPose"
|
||||
elif pose_estimator.endswith(".torchscript.pt"):
|
||||
pose_repo = "hr16/DWPose-TorchScript-BatchSize5"
|
||||
else:
|
||||
raise NotImplementedError(f"Download mechanism for {pose_estimator}")
|
||||
|
||||
model = DwposeDetector.from_pretrained(
|
||||
pose_repo,
|
||||
yolo_repo,
|
||||
det_filename=(None if bbox_detector == "None" else bbox_detector), pose_filename=pose_estimator,
|
||||
torchscript_device=model_management.get_torch_device()
|
||||
)
|
||||
detect_hand = detect_hand == "enable"
|
||||
detect_body = detect_body == "enable"
|
||||
detect_face = detect_face == "enable"
|
||||
scale_stick_for_xinsr_cn = scale_stick_for_xinsr_cn == "enable"
|
||||
self.openpose_dicts = []
|
||||
def func(image, **kwargs):
|
||||
pose_img, openpose_dict = model(image, **kwargs)
|
||||
self.openpose_dicts.append(openpose_dict)
|
||||
return pose_img
|
||||
|
||||
out = common_annotator_call(func, image, include_hand=detect_hand, include_face=detect_face, include_body=detect_body, image_and_json=True, resolution=resolution, xinsr_stick_scaling=scale_stick_for_xinsr_cn)
|
||||
del model
|
||||
return {
|
||||
'ui': { "openpose_json": [json.dumps(self.openpose_dicts, indent=4)] },
|
||||
"result": (out, self.openpose_dicts)
|
||||
}
|
||||
|
||||
class AnimalPose_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
bbox_detector = INPUT.COMBO(
|
||||
["None"] + ["yolox_l.torchscript.pt", "yolox_l.onnx", "yolo_nas_l_fp16.onnx", "yolo_nas_m_fp16.onnx", "yolo_nas_s_fp16.onnx"],
|
||||
default="yolox_l.torchscript.pt"
|
||||
),
|
||||
pose_estimator = INPUT.COMBO(
|
||||
["rtmpose-m_ap10k_256_bs5.torchscript.pt", "rtmpose-m_ap10k_256.onnx"],
|
||||
default="rtmpose-m_ap10k_256_bs5.torchscript.pt"
|
||||
),
|
||||
resolution = INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
|
||||
FUNCTION = "estimate_pose"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def estimate_pose(self, image, resolution=512, bbox_detector="yolox_l.onnx", pose_estimator="rtmpose-m_ap10k_256.onnx", **kwargs):
|
||||
if bbox_detector == "None":
|
||||
yolo_repo = DWPOSE_MODEL_NAME
|
||||
elif bbox_detector == "yolox_l.onnx":
|
||||
yolo_repo = DWPOSE_MODEL_NAME
|
||||
elif "yolox" in bbox_detector:
|
||||
yolo_repo = "hr16/yolox-onnx"
|
||||
elif "yolo_nas" in bbox_detector:
|
||||
yolo_repo = "hr16/yolo-nas-fp16"
|
||||
else:
|
||||
raise NotImplementedError(f"Download mechanism for {bbox_detector}")
|
||||
|
||||
if pose_estimator == "dw-ll_ucoco_384.onnx":
|
||||
pose_repo = DWPOSE_MODEL_NAME
|
||||
elif pose_estimator.endswith(".onnx"):
|
||||
pose_repo = "hr16/UnJIT-DWPose"
|
||||
elif pose_estimator.endswith(".torchscript.pt"):
|
||||
pose_repo = "hr16/DWPose-TorchScript-BatchSize5"
|
||||
else:
|
||||
raise NotImplementedError(f"Download mechanism for {pose_estimator}")
|
||||
|
||||
model = AnimalposeDetector.from_pretrained(
|
||||
pose_repo,
|
||||
yolo_repo,
|
||||
det_filename=(None if bbox_detector == "None" else bbox_detector), pose_filename=pose_estimator,
|
||||
torchscript_device=model_management.get_torch_device()
|
||||
)
|
||||
|
||||
self.openpose_dicts = []
|
||||
def func(image, **kwargs):
|
||||
pose_img, openpose_dict = model(image, **kwargs)
|
||||
self.openpose_dicts.append(openpose_dict)
|
||||
return pose_img
|
||||
|
||||
out = common_annotator_call(func, image, image_and_json=True, resolution=resolution)
|
||||
del model
|
||||
return {
|
||||
'ui': { "openpose_json": [json.dumps(self.openpose_dicts, indent=4)] },
|
||||
"result": (out, self.openpose_dicts)
|
||||
}
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DWPreprocessor": DWPose_Preprocessor,
|
||||
"AnimalPosePreprocessor": AnimalPose_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DWPreprocessor": "DWPose Estimator",
|
||||
"AnimalPosePreprocessor": "AnimalPose Estimator (AP10K)"
|
||||
}
|
||||
53
custom_nodes/comfyui_controlnet_aux/node_wrappers/hed.py
Normal file
53
custom_nodes/comfyui_controlnet_aux/node_wrappers/hed.py
Normal file
@@ -0,0 +1,53 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class HED_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.hed import HEDdetector
|
||||
|
||||
model = HEDdetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, safe = kwargs["safe"] == "enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
class Fake_Scribble_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.hed import HEDdetector
|
||||
|
||||
model = HEDdetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, scribble=True, safe=kwargs["safe"]=="enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"HEDPreprocessor": HED_Preprocessor,
|
||||
"FakeScribblePreprocessor": Fake_Scribble_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"HEDPreprocessor": "HED Soft-Edge Lines",
|
||||
"FakeScribblePreprocessor": "Fake Scribble Lines (aka scribble_hed)"
|
||||
}
|
||||
32
custom_nodes/comfyui_controlnet_aux/node_wrappers/inpaint.py
Normal file
32
custom_nodes/comfyui_controlnet_aux/node_wrappers/inpaint.py
Normal file
@@ -0,0 +1,32 @@
|
||||
import torch
|
||||
from ..utils import INPUT
|
||||
|
||||
class InpaintPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return dict(
|
||||
required=dict(image=INPUT.IMAGE(), mask=INPUT.MASK()),
|
||||
optional=dict(black_pixel_for_xinsir_cn=INPUT.BOOLEAN(False))
|
||||
)
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "preprocess"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/others"
|
||||
|
||||
def preprocess(self, image, mask, black_pixel_for_xinsir_cn=False):
|
||||
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(image.shape[1], image.shape[2]), mode="bilinear")
|
||||
mask = mask.movedim(1,-1).expand((-1,-1,-1,3))
|
||||
image = image.clone()
|
||||
if black_pixel_for_xinsir_cn:
|
||||
masked_pixel = 0.0
|
||||
else:
|
||||
masked_pixel = -1.0
|
||||
image[mask > 0.5] = masked_pixel
|
||||
return (image,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"InpaintPreprocessor": InpaintPreprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"InpaintPreprocessor": "Inpaint Preprocessor"
|
||||
}
|
||||
32
custom_nodes/comfyui_controlnet_aux/node_wrappers/leres.py
Normal file
32
custom_nodes/comfyui_controlnet_aux/node_wrappers/leres.py
Normal file
@@ -0,0 +1,32 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class LERES_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
rm_nearest=INPUT.FLOAT(max=100.0),
|
||||
rm_background=INPUT.FLOAT(max=100.0),
|
||||
boost=INPUT.COMBO(["disable", "enable"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, rm_nearest=0, rm_background=0, resolution=512, boost="disable", **kwargs):
|
||||
from custom_controlnet_aux.leres import LeresDetector
|
||||
|
||||
model = LeresDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, thr_a=rm_nearest, thr_b=rm_background, boost=boost == "enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LeReS-DepthMapPreprocessor": LERES_Depth_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LeReS-DepthMapPreprocessor": "LeReS Depth Map (enable boost for leres++)"
|
||||
}
|
||||
30
custom_nodes/comfyui_controlnet_aux/node_wrappers/lineart.py
Normal file
30
custom_nodes/comfyui_controlnet_aux/node_wrappers/lineart.py
Normal file
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class LineArt_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
coarse=INPUT.COMBO((["disable", "enable"])),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.lineart import LineartDetector
|
||||
|
||||
model = LineartDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, coarse = kwargs["coarse"] == "enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LineArtPreprocessor": LineArt_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LineArtPreprocessor": "Realistic Lineart"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class AnimeLineArt_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.lineart_anime import LineartAnimeDetector
|
||||
|
||||
model = LineartAnimeDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AnimeLineArtPreprocessor": AnimeLineArt_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AnimeLineArtPreprocessor": "Anime Lineart"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Lineart_Standard_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
guassian_sigma=INPUT.FLOAT(default=6.0, max=100.0),
|
||||
intensity_threshold=INPUT.INT(default=8, max=16),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, guassian_sigma=6, intensity_threshold=8, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.lineart_standard import LineartStandardDetector
|
||||
return (common_annotator_call(LineartStandardDetector(), image, guassian_sigma=guassian_sigma, intensity_threshold=intensity_threshold, resolution=resolution), )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LineartStandardPreprocessor": Lineart_Standard_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LineartStandardPreprocessor": "Standard Lineart"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Manga2Anime_LineArt_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.manga_line import LineartMangaDetector
|
||||
|
||||
model = LineartMangaDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Manga2Anime_LineArt_Preprocessor": Manga2Anime_LineArt_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Manga2Anime_LineArt_Preprocessor": "Manga Lineart (aka lineart_anime_denoise)"
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, run_script
|
||||
import comfy.model_management as model_management
|
||||
import os, sys
|
||||
import subprocess, threading
|
||||
|
||||
def install_deps():
|
||||
try:
|
||||
import mediapipe
|
||||
except ImportError:
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', 'mediapipe'])
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', '--upgrade', 'protobuf'])
|
||||
|
||||
class Media_Pipe_Face_Mesh_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
max_faces=INPUT.INT(default=10, min=1, max=50), #Which image has more than 50 detectable faces?
|
||||
min_confidence=INPUT.FLOAT(default=0.5, min=0.1),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "detect"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def detect(self, image, max_faces=10, min_confidence=0.5, resolution=512):
|
||||
#Ref: https://github.com/Fannovel16/comfy_controlnet_preprocessors/issues/70#issuecomment-1677967369
|
||||
install_deps()
|
||||
from custom_controlnet_aux.mediapipe_face import MediapipeFaceDetector
|
||||
return (common_annotator_call(MediapipeFaceDetector(), image, max_faces=max_faces, min_confidence=min_confidence, resolution=resolution), )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MediaPipe-FaceMeshPreprocessor": Media_Pipe_Face_Mesh_Preprocessor
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MediaPipe-FaceMeshPreprocessor": "MediaPipe Face Mesh"
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, MAX_RESOLUTION, run_script
|
||||
import comfy.model_management as model_management
|
||||
import numpy as np
|
||||
import torch
|
||||
from einops import rearrange
|
||||
import os, sys
|
||||
import subprocess, threading
|
||||
import scipy.ndimage
|
||||
import cv2
|
||||
import torch.nn.functional as F
|
||||
|
||||
def install_deps():
|
||||
try:
|
||||
import mediapipe
|
||||
except ImportError:
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', 'mediapipe'])
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', '--upgrade', 'protobuf'])
|
||||
|
||||
try:
|
||||
import trimesh
|
||||
except ImportError:
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', 'trimesh[easy]'])
|
||||
|
||||
#Sauce: https://github.com/comfyanonymous/ComfyUI/blob/8c6493578b3dda233e9b9a953feeaf1e6ca434ad/comfy_extras/nodes_mask.py#L309
|
||||
def expand_mask(mask, expand, tapered_corners):
|
||||
c = 0 if tapered_corners else 1
|
||||
kernel = np.array([[c, 1, c],
|
||||
[1, 1, 1],
|
||||
[c, 1, c]])
|
||||
mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
|
||||
out = []
|
||||
for m in mask:
|
||||
output = m.numpy()
|
||||
for _ in range(abs(expand)):
|
||||
if expand < 0:
|
||||
output = scipy.ndimage.grey_erosion(output, footprint=kernel)
|
||||
else:
|
||||
output = scipy.ndimage.grey_dilation(output, footprint=kernel)
|
||||
output = torch.from_numpy(output)
|
||||
out.append(output)
|
||||
return torch.stack(out, dim=0)
|
||||
|
||||
class Mesh_Graphormer_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
mask_bbox_padding=("INT", {"default": 30, "min": 0, "max": 100}),
|
||||
resolution=INPUT.RESOLUTION(),
|
||||
mask_type=INPUT.COMBO(["based_on_depth", "tight_bboxes", "original"]),
|
||||
mask_expand=INPUT.INT(default=5, min=-MAX_RESOLUTION, max=MAX_RESOLUTION),
|
||||
rand_seed=INPUT.INT(default=88, min=0, max=0xffffffffffffffff),
|
||||
detect_thr=INPUT.FLOAT(default=0.6, min=0.1),
|
||||
presence_thr=INPUT.FLOAT(default=0.6, min=0.1)
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
RETURN_NAMES = ("IMAGE", "INPAINTING_MASK")
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, mask_bbox_padding=30, mask_type="based_on_depth", mask_expand=5, resolution=512, rand_seed=88, detect_thr=0.6, presence_thr=0.6, **kwargs):
|
||||
install_deps()
|
||||
from custom_controlnet_aux.mesh_graphormer import MeshGraphormerDetector
|
||||
model = kwargs["model"] if "model" in kwargs \
|
||||
else MeshGraphormerDetector.from_pretrained(detect_thr=detect_thr, presence_thr=presence_thr).to(model_management.get_torch_device())
|
||||
|
||||
depth_map_list = []
|
||||
mask_list = []
|
||||
for single_image in image:
|
||||
np_image = np.asarray(single_image.cpu() * 255., dtype=np.uint8)
|
||||
depth_map, mask, info = model(np_image, output_type="np", detect_resolution=resolution, mask_bbox_padding=mask_bbox_padding, seed=rand_seed)
|
||||
if mask_type == "based_on_depth":
|
||||
H, W = mask.shape[:2]
|
||||
mask = cv2.resize(depth_map.copy(), (W, H))
|
||||
mask[mask > 0] = 255
|
||||
|
||||
elif mask_type == "tight_bboxes":
|
||||
mask = np.zeros_like(mask)
|
||||
hand_bboxes = (info or {}).get("abs_boxes") or []
|
||||
for hand_bbox in hand_bboxes:
|
||||
x_min, x_max, y_min, y_max = hand_bbox
|
||||
mask[y_min:y_max+1, x_min:x_max+1, :] = 255 #HWC
|
||||
|
||||
mask = mask[:, :, :1]
|
||||
depth_map_list.append(torch.from_numpy(depth_map.astype(np.float32) / 255.0))
|
||||
mask_list.append(torch.from_numpy(mask.astype(np.float32) / 255.0))
|
||||
depth_maps, masks = torch.stack(depth_map_list, dim=0), rearrange(torch.stack(mask_list, dim=0), "n h w 1 -> n 1 h w")
|
||||
return depth_maps, expand_mask(masks, mask_expand, tapered_corners=True)
|
||||
|
||||
def normalize_size_base_64(w, h):
|
||||
short_side = min(w, h)
|
||||
remainder = short_side % 64
|
||||
return short_side - remainder + (64 if remainder > 0 else 0)
|
||||
|
||||
class Mesh_Graphormer_With_ImpactDetector_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
types = define_preprocessor_inputs(
|
||||
# Impact pack
|
||||
bbox_threshold=INPUT.FLOAT(default=0.5, min=0.1),
|
||||
bbox_dilation=INPUT.INT(default=10, min=-512, max=512),
|
||||
bbox_crop_factor=INPUT.FLOAT(default=3.0, min=1.0, max=10.0),
|
||||
drop_size=INPUT.INT(default=10, min=1, max=MAX_RESOLUTION),
|
||||
# Mesh Graphormer
|
||||
mask_bbox_padding=INPUT.INT(default=30, min=0, max=100),
|
||||
mask_type=INPUT.COMBO(["based_on_depth", "tight_bboxes", "original"]),
|
||||
mask_expand=INPUT.INT(default=5, min=-MAX_RESOLUTION, max=MAX_RESOLUTION),
|
||||
rand_seed=INPUT.INT(default=88, min=0, max=0xffffffffffffffff),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
types["required"]["bbox_detector"] = ("BBOX_DETECTOR", )
|
||||
return types
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
RETURN_NAMES = ("IMAGE", "INPAINTING_MASK")
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, bbox_detector, bbox_threshold=0.5, bbox_dilation=10, bbox_crop_factor=3.0, drop_size=10, resolution=512, **mesh_graphormer_kwargs):
|
||||
install_deps()
|
||||
from custom_controlnet_aux.mesh_graphormer import MeshGraphormerDetector
|
||||
mesh_graphormer_node = Mesh_Graphormer_Depth_Map_Preprocessor()
|
||||
model = MeshGraphormerDetector.from_pretrained(detect_thr=0.6, presence_thr=0.6).to(model_management.get_torch_device())
|
||||
mesh_graphormer_kwargs["model"] = model
|
||||
|
||||
frames = image
|
||||
depth_maps, masks = [], []
|
||||
for idx in range(len(frames)):
|
||||
frame = frames[idx:idx+1,...] #Impact Pack's BBOX_DETECTOR only supports single batch image
|
||||
bbox_detector.setAux('face') # make default prompt as 'face' if empty prompt for CLIPSeg
|
||||
_, segs = bbox_detector.detect(frame, bbox_threshold, bbox_dilation, bbox_crop_factor, drop_size)
|
||||
bbox_detector.setAux(None)
|
||||
|
||||
n, h, w, _ = frame.shape
|
||||
depth_map, mask = torch.zeros_like(frame), torch.zeros(n, 1, h, w)
|
||||
for i, seg in enumerate(segs):
|
||||
x1, y1, x2, y2 = seg.crop_region
|
||||
cropped_image = frame[:, y1:y2, x1:x2, :] # Never use seg.cropped_image to handle overlapping area
|
||||
mesh_graphormer_kwargs["resolution"] = 0 #Disable resizing
|
||||
sub_depth_map, sub_mask = mesh_graphormer_node.execute(cropped_image, **mesh_graphormer_kwargs)
|
||||
depth_map[:, y1:y2, x1:x2, :] = sub_depth_map
|
||||
mask[:, :, y1:y2, x1:x2] = sub_mask
|
||||
|
||||
depth_maps.append(depth_map)
|
||||
masks.append(mask)
|
||||
|
||||
return (torch.cat(depth_maps), torch.cat(masks))
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MeshGraphormer-DepthMapPreprocessor": Mesh_Graphormer_Depth_Map_Preprocessor,
|
||||
"MeshGraphormer+ImpactDetector-DepthMapPreprocessor": Mesh_Graphormer_With_ImpactDetector_Depth_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MeshGraphormer-DepthMapPreprocessor": "MeshGraphormer Hand Refiner",
|
||||
"MeshGraphormer+ImpactDetector-DepthMapPreprocessor": "MeshGraphormer Hand Refiner With External Detector"
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
import os
|
||||
# Disable NPU device initialization and problematic MMCV ops to prevent RuntimeError
|
||||
os.environ['NPU_DEVICE_COUNT'] = '0'
|
||||
os.environ['MMCV_WITH_OPS'] = '0'
|
||||
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, MAX_RESOLUTION
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Metric3D_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
backbone=INPUT.COMBO(["vit-small", "vit-large", "vit-giant2"]),
|
||||
fx=INPUT.INT(default=1000, min=1, max=MAX_RESOLUTION),
|
||||
fy=INPUT.INT(default=1000, min=1, max=MAX_RESOLUTION),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, backbone="vit-small", fx=1000, fy=1000, resolution=512):
|
||||
from custom_controlnet_aux.metric3d import Metric3DDetector
|
||||
model = Metric3DDetector.from_pretrained(filename=f"metric_depth_{backbone.replace('-', '_')}_800k.pth").to(model_management.get_torch_device())
|
||||
cb = lambda image, **kwargs: model(image, **kwargs)[0]
|
||||
out = common_annotator_call(cb, image, resolution=resolution, fx=fx, fy=fy, depth_and_normal=True)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
class Metric3D_Normal_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
backbone=INPUT.COMBO(["vit-small", "vit-large", "vit-giant2"]),
|
||||
fx=INPUT.INT(default=1000, min=1, max=MAX_RESOLUTION),
|
||||
fy=INPUT.INT(default=1000, min=1, max=MAX_RESOLUTION),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, backbone="vit-small", fx=1000, fy=1000, resolution=512):
|
||||
from custom_controlnet_aux.metric3d import Metric3DDetector
|
||||
model = Metric3DDetector.from_pretrained(filename=f"metric_depth_{backbone.replace('-', '_')}_800k.pth").to(model_management.get_torch_device())
|
||||
cb = lambda image, **kwargs: model(image, **kwargs)[1]
|
||||
out = common_annotator_call(cb, image, resolution=resolution, fx=fx, fy=fy, depth_and_normal=True)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Metric3D-DepthMapPreprocessor": Metric3D_Depth_Map_Preprocessor,
|
||||
"Metric3D-NormalMapPreprocessor": Metric3D_Normal_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Metric3D-DepthMapPreprocessor": "Metric3D Depth Map",
|
||||
"Metric3D-NormalMapPreprocessor": "Metric3D Normal Map"
|
||||
}
|
||||
59
custom_nodes/comfyui_controlnet_aux/node_wrappers/midas.py
Normal file
59
custom_nodes/comfyui_controlnet_aux/node_wrappers/midas.py
Normal file
@@ -0,0 +1,59 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import numpy as np
|
||||
|
||||
class MIDAS_Normal_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
a=INPUT.FLOAT(default=np.pi * 2.0, min=0.0, max=np.pi * 5.0),
|
||||
bg_threshold=INPUT.FLOAT(default=0.1),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, a=np.pi * 2.0, bg_threshold=0.1, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.midas import MidasDetector
|
||||
|
||||
model = MidasDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
#Dirty hack :))
|
||||
cb = lambda image, **kargs: model(image, **kargs)[1]
|
||||
out = common_annotator_call(cb, image, resolution=resolution, a=a, bg_th=bg_threshold, depth_and_normal=True)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
class MIDAS_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
a=INPUT.FLOAT(default=np.pi * 2.0, min=0.0, max=np.pi * 5.0),
|
||||
bg_threshold=INPUT.FLOAT(default=0.1),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, a=np.pi * 2.0, bg_threshold=0.1, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.midas import MidasDetector
|
||||
|
||||
# Ref: https://github.com/lllyasviel/ControlNet/blob/main/gradio_depth2image.py
|
||||
model = MidasDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, a=a, bg_th=bg_threshold)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MiDaS-NormalMapPreprocessor": MIDAS_Normal_Map_Preprocessor,
|
||||
"MiDaS-DepthMapPreprocessor": MIDAS_Depth_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MiDaS-NormalMapPreprocessor": "MiDaS Normal Map",
|
||||
"MiDaS-DepthMapPreprocessor": "MiDaS Depth Map"
|
||||
}
|
||||
31
custom_nodes/comfyui_controlnet_aux/node_wrappers/mlsd.py
Normal file
31
custom_nodes/comfyui_controlnet_aux/node_wrappers/mlsd.py
Normal file
@@ -0,0 +1,31 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import numpy as np
|
||||
|
||||
class MLSD_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
score_threshold=INPUT.FLOAT(default=0.1, min=0.01, max=2.0),
|
||||
dist_threshold=INPUT.FLOAT(default=0.1, min=0.01, max=20.0),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, score_threshold, dist_threshold, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.mlsd import MLSDdetector
|
||||
|
||||
model = MLSDdetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, thr_v=score_threshold, thr_d=dist_threshold)
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"M-LSDPreprocessor": MLSD_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"M-LSDPreprocessor": "M-LSD Lines"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class BAE_Normal_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.normalbae import NormalBaeDetector
|
||||
|
||||
model = NormalBaeDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"BAE-NormalMapPreprocessor": BAE_Normal_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"BAE-NormalMapPreprocessor": "BAE Normal Map"
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class OneFormer_COCO_SemSegPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "semantic_segmentate"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Semantic Segmentation"
|
||||
|
||||
def semantic_segmentate(self, image, resolution=512):
|
||||
from custom_controlnet_aux.oneformer import OneformerSegmentor
|
||||
|
||||
model = OneformerSegmentor.from_pretrained(filename="150_16_swin_l_oneformer_coco_100ep.pth")
|
||||
model = model.to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out,)
|
||||
|
||||
class OneFormer_ADE20K_SemSegPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "semantic_segmentate"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Semantic Segmentation"
|
||||
|
||||
def semantic_segmentate(self, image, resolution=512):
|
||||
from custom_controlnet_aux.oneformer import OneformerSegmentor
|
||||
|
||||
model = OneformerSegmentor.from_pretrained(filename="250_16_swin_l_oneformer_ade20k_160k.pth")
|
||||
model = model.to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"OneFormer-COCO-SemSegPreprocessor": OneFormer_COCO_SemSegPreprocessor,
|
||||
"OneFormer-ADE20K-SemSegPreprocessor": OneFormer_ADE20K_SemSegPreprocessor
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"OneFormer-COCO-SemSegPreprocessor": "OneFormer COCO Segmentor",
|
||||
"OneFormer-ADE20K-SemSegPreprocessor": "OneFormer ADE20K Segmentor"
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import json
|
||||
|
||||
class OpenPose_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
detect_hand=INPUT.COMBO(["enable", "disable"]),
|
||||
detect_body=INPUT.COMBO(["enable", "disable"]),
|
||||
detect_face=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION(),
|
||||
scale_stick_for_xinsr_cn=INPUT.COMBO(["disable", "enable"])
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
|
||||
FUNCTION = "estimate_pose"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def estimate_pose(self, image, detect_hand="enable", detect_body="enable", detect_face="enable", scale_stick_for_xinsr_cn="disable", resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.open_pose import OpenposeDetector
|
||||
|
||||
detect_hand = detect_hand == "enable"
|
||||
detect_body = detect_body == "enable"
|
||||
detect_face = detect_face == "enable"
|
||||
scale_stick_for_xinsr_cn = scale_stick_for_xinsr_cn == "enable"
|
||||
|
||||
model = OpenposeDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
self.openpose_dicts = []
|
||||
def func(image, **kwargs):
|
||||
pose_img, openpose_dict = model(image, **kwargs)
|
||||
self.openpose_dicts.append(openpose_dict)
|
||||
return pose_img
|
||||
|
||||
out = common_annotator_call(func, image, include_hand=detect_hand, include_face=detect_face, include_body=detect_body, image_and_json=True, xinsr_stick_scaling=scale_stick_for_xinsr_cn, resolution=resolution)
|
||||
del model
|
||||
return {
|
||||
'ui': { "openpose_json": [json.dumps(self.openpose_dicts, indent=4)] },
|
||||
"result": (out, self.openpose_dicts)
|
||||
}
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"OpenposePreprocessor": OpenPose_Preprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"OpenposePreprocessor": "OpenPose Pose",
|
||||
}
|
||||
30
custom_nodes/comfyui_controlnet_aux/node_wrappers/pidinet.py
Normal file
30
custom_nodes/comfyui_controlnet_aux/node_wrappers/pidinet.py
Normal file
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class PIDINET_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, safe, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.pidi import PidiNetDetector
|
||||
|
||||
model = PidiNetDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, safe = safe == "enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PiDiNetPreprocessor": PIDINET_Preprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PiDiNetPreprocessor": "PiDiNet Soft-Edge Lines"
|
||||
}
|
||||
@@ -0,0 +1,340 @@
|
||||
import folder_paths
|
||||
import json
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import ImageColor
|
||||
from einops import rearrange
|
||||
import torch
|
||||
import itertools
|
||||
|
||||
from ..src.custom_controlnet_aux.dwpose import draw_poses, draw_animalposes, decode_json_as_poses
|
||||
|
||||
|
||||
"""
|
||||
Format of POSE_KEYPOINT (AP10K keypoints):
|
||||
[{
|
||||
"version": "ap10k",
|
||||
"animals": [
|
||||
[[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]],
|
||||
[[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]],
|
||||
...
|
||||
],
|
||||
"canvas_height": 512,
|
||||
"canvas_width": 768
|
||||
},...]
|
||||
Format of POSE_KEYPOINT (OpenPose keypoints):
|
||||
[{
|
||||
"people": [
|
||||
{
|
||||
'pose_keypoints_2d': [[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]]
|
||||
"face_keypoints_2d": [[x1, y1, 1], [x2, y2, 1],..., [x68, y68, 1]],
|
||||
"hand_left_keypoints_2d": [[x1, y1, 1], [x2, y2, 1],..., [x21, y21, 1]],
|
||||
"hand_right_keypoints_2d":[[x1, y1, 1], [x2, y2, 1],..., [x21, y21, 1]],
|
||||
}
|
||||
],
|
||||
"canvas_height": canvas_height,
|
||||
"canvas_width": canvas_width,
|
||||
},...]
|
||||
"""
|
||||
|
||||
class SavePoseKpsAsJsonFile:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"pose_kps": ("POSE_KEYPOINT",),
|
||||
"filename_prefix": ("STRING", {"default": "PoseKeypoint"})
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save_pose_kps"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
self.type = "output"
|
||||
self.prefix_append = ""
|
||||
def save_pose_kps(self, pose_kps, filename_prefix):
|
||||
filename_prefix += self.prefix_append
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = \
|
||||
folder_paths.get_save_image_path(filename_prefix, self.output_dir, pose_kps[0]["canvas_width"], pose_kps[0]["canvas_height"])
|
||||
file = f"{filename}_{counter:05}.json"
|
||||
with open(os.path.join(full_output_folder, file), 'w') as f:
|
||||
json.dump(pose_kps , f)
|
||||
return {}
|
||||
|
||||
#COCO-Wholebody doesn't have eyebrows as it inherits 68 keypoints format
|
||||
#Perhaps eyebrows can be estimated tho
|
||||
FACIAL_PARTS = ["skin", "left_eye", "right_eye", "nose", "upper_lip", "inner_mouth", "lower_lip"]
|
||||
LAPA_COLORS = dict(
|
||||
skin="rgb(0, 153, 255)",
|
||||
left_eye="rgb(0, 204, 153)",
|
||||
right_eye="rgb(255, 153, 0)",
|
||||
nose="rgb(255, 102, 255)",
|
||||
upper_lip="rgb(102, 0, 51)",
|
||||
inner_mouth="rgb(255, 204, 255)",
|
||||
lower_lip="rgb(255, 0, 102)"
|
||||
)
|
||||
|
||||
#One-based index
|
||||
def kps_idxs(start, end):
|
||||
step = -1 if start > end else 1
|
||||
return list(range(start-1, end+1-1, step))
|
||||
|
||||
#Source: https://www.researchgate.net/profile/Fabrizio-Falchi/publication/338048224/figure/fig1/AS:837860722741255@1576772971540/68-facial-landmarks.jpg
|
||||
FACIAL_PART_RANGES = dict(
|
||||
skin=kps_idxs(1, 17) + kps_idxs(27, 18),
|
||||
nose=kps_idxs(28, 36),
|
||||
left_eye=kps_idxs(37, 42),
|
||||
right_eye=kps_idxs(43, 48),
|
||||
upper_lip=kps_idxs(49, 55) + kps_idxs(65, 61),
|
||||
lower_lip=kps_idxs(61, 68),
|
||||
inner_mouth=kps_idxs(61, 65) + kps_idxs(55, 49)
|
||||
)
|
||||
|
||||
def is_normalized(keypoints) -> bool:
|
||||
point_normalized = [
|
||||
0 <= np.abs(k[0]) <= 1 and 0 <= np.abs(k[1]) <= 1
|
||||
for k in keypoints
|
||||
if k is not None
|
||||
]
|
||||
if not point_normalized:
|
||||
return False
|
||||
return np.all(point_normalized)
|
||||
|
||||
class FacialPartColoringFromPoseKps:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
input_types = {
|
||||
"required": {"pose_kps": ("POSE_KEYPOINT",), "mode": (["point", "polygon"], {"default": "polygon"})}
|
||||
}
|
||||
for facial_part in FACIAL_PARTS:
|
||||
input_types["required"][facial_part] = ("STRING", {"default": LAPA_COLORS[facial_part], "multiline": False})
|
||||
return input_types
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "colorize"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
def colorize(self, pose_kps, mode, **facial_part_colors):
|
||||
pose_frames = pose_kps
|
||||
np_frames = [self.draw_kps(pose_frame, mode, **facial_part_colors) for pose_frame in pose_frames]
|
||||
np_frames = np.stack(np_frames, axis=0)
|
||||
return (torch.from_numpy(np_frames).float() / 255.,)
|
||||
|
||||
def draw_kps(self, pose_frame, mode, **facial_part_colors):
|
||||
width, height = pose_frame["canvas_width"], pose_frame["canvas_height"]
|
||||
canvas = np.zeros((height, width, 3), dtype=np.uint8)
|
||||
for person, part_name in itertools.product(pose_frame["people"], FACIAL_PARTS):
|
||||
n = len(person["face_keypoints_2d"]) // 3
|
||||
facial_kps = rearrange(np.array(person["face_keypoints_2d"]), "(n c) -> n c", n=n, c=3)[:, :2]
|
||||
if is_normalized(facial_kps):
|
||||
facial_kps *= (width, height)
|
||||
facial_kps = facial_kps.astype(np.int32)
|
||||
part_color = ImageColor.getrgb(facial_part_colors[part_name])[:3]
|
||||
part_contours = facial_kps[FACIAL_PART_RANGES[part_name], :]
|
||||
if mode == "point":
|
||||
for pt in part_contours:
|
||||
cv2.circle(canvas, pt, radius=2, color=part_color, thickness=-1)
|
||||
else:
|
||||
cv2.fillPoly(canvas, pts=[part_contours], color=part_color)
|
||||
return canvas
|
||||
|
||||
# https://raw.githubusercontent.com/CMU-Perceptual-Computing-Lab/openpose/master/.github/media/keypoints_pose_18.png
|
||||
BODY_PART_INDEXES = {
|
||||
"Head": (16, 14, 0, 15, 17),
|
||||
"Neck": (0, 1),
|
||||
"Shoulder": (2, 5),
|
||||
"Torso": (2, 5, 8, 11),
|
||||
"RArm": (2, 3),
|
||||
"RForearm": (3, 4),
|
||||
"LArm": (5, 6),
|
||||
"LForearm": (6, 7),
|
||||
"RThigh": (8, 9),
|
||||
"RLeg": (9, 10),
|
||||
"LThigh": (11, 12),
|
||||
"LLeg": (12, 13)
|
||||
}
|
||||
BODY_PART_DEFAULT_W_H = {
|
||||
"Head": "256, 256",
|
||||
"Neck": "100, 100",
|
||||
"Shoulder": '',
|
||||
"Torso": "350, 450",
|
||||
"RArm": "128, 256",
|
||||
"RForearm": "128, 256",
|
||||
"LArm": "128, 256",
|
||||
"LForearm": "128, 256",
|
||||
"RThigh": "128, 256",
|
||||
"RLeg": "128, 256",
|
||||
"LThigh": "128, 256",
|
||||
"LLeg": "128, 256"
|
||||
}
|
||||
|
||||
class SinglePersonProcess:
|
||||
@classmethod
|
||||
def sort_and_get_max_people(s, pose_kps):
|
||||
for idx in range(len(pose_kps)):
|
||||
pose_kps[idx]["people"] = sorted(pose_kps[idx]["people"], key=lambda person:person["pose_keypoints_2d"][0])
|
||||
return pose_kps, max(len(frame["people"]) for frame in pose_kps)
|
||||
|
||||
def __init__(self, pose_kps, person_idx=0) -> None:
|
||||
self.width, self.height = pose_kps[0]["canvas_width"], pose_kps[0]["canvas_height"]
|
||||
self.poses = [
|
||||
self.normalize(pose_frame["people"][person_idx]["pose_keypoints_2d"])
|
||||
if person_idx < len(pose_frame["people"])
|
||||
else None
|
||||
for pose_frame in pose_kps
|
||||
]
|
||||
|
||||
def normalize(self, pose_kps_2d):
|
||||
n = len(pose_kps_2d) // 3
|
||||
pose_kps_2d = rearrange(np.array(pose_kps_2d), "(n c) -> n c", n=n, c=3)
|
||||
pose_kps_2d[np.argwhere(pose_kps_2d[:,2]==0), :] = np.iinfo(np.int32).max // 2 #Safe large value
|
||||
pose_kps_2d = pose_kps_2d[:, :2]
|
||||
if is_normalized(pose_kps_2d):
|
||||
pose_kps_2d *= (self.width, self.height)
|
||||
return pose_kps_2d
|
||||
|
||||
def get_xyxy_bboxes(self, part_name, bbox_size=(128, 256)):
|
||||
width, height = bbox_size
|
||||
xyxy_bboxes = {}
|
||||
for idx, pose in enumerate(self.poses):
|
||||
if pose is None:
|
||||
xyxy_bboxes[idx] = (np.iinfo(np.int32).max // 2,) * 4
|
||||
continue
|
||||
pts = pose[BODY_PART_INDEXES[part_name], :]
|
||||
|
||||
#top_left = np.min(pts[:,0]), np.min(pts[:,1])
|
||||
#bottom_right = np.max(pts[:,0]), np.max(pts[:,1])
|
||||
#pad_width = np.maximum(width - (bottom_right[0]-top_left[0]), 0) / 2
|
||||
#pad_height = np.maximum(height - (bottom_right[1]-top_left[1]), 0) / 2
|
||||
#xyxy_bboxes.append((
|
||||
# top_left[0] - pad_width, top_left[1] - pad_height,
|
||||
# bottom_right[0] + pad_width, bottom_right[1] + pad_height,
|
||||
#))
|
||||
|
||||
x_mid, y_mid = np.mean(pts[:, 0]), np.mean(pts[:, 1])
|
||||
xyxy_bboxes[idx] = (
|
||||
x_mid - width/2, y_mid - height/2,
|
||||
x_mid + width/2, y_mid + height/2
|
||||
)
|
||||
return xyxy_bboxes
|
||||
|
||||
class UpperBodyTrackingFromPoseKps:
|
||||
PART_NAMES = ["Head", "Neck", "Shoulder", "Torso", "RArm", "RForearm", "LArm", "LForearm"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"pose_kps": ("POSE_KEYPOINT",),
|
||||
"id_include": ("STRING", {"default": '', "multiline": False}),
|
||||
**{part_name + "_width_height": ("STRING", {"default": BODY_PART_DEFAULT_W_H[part_name], "multiline": False}) for part_name in s.PART_NAMES}
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("TRACKING", "STRING")
|
||||
RETURN_NAMES = ("tracking", "prompt")
|
||||
FUNCTION = "convert"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
|
||||
def convert(self, pose_kps, id_include, **parts_width_height):
|
||||
parts_width_height = {part_name.replace("_width_height", ''): value for part_name, value in parts_width_height.items()}
|
||||
enabled_part_names = [part_name for part_name in self.PART_NAMES if len(parts_width_height[part_name].strip())]
|
||||
tracked = {part_name: {} for part_name in enabled_part_names}
|
||||
id_include = id_include.strip()
|
||||
id_include = list(map(int, id_include.split(','))) if len(id_include) else []
|
||||
prompt_string = ''
|
||||
pose_kps, max_people = SinglePersonProcess.sort_and_get_max_people(pose_kps)
|
||||
|
||||
for person_idx in range(max_people):
|
||||
if len(id_include) and person_idx not in id_include:
|
||||
continue
|
||||
processor = SinglePersonProcess(pose_kps, person_idx)
|
||||
for part_name in enabled_part_names:
|
||||
bbox_size = tuple(map(int, parts_width_height[part_name].split(',')))
|
||||
part_bboxes = processor.get_xyxy_bboxes(part_name, bbox_size)
|
||||
id_coordinates = {idx: part_bbox+(processor.width, processor.height) for idx, part_bbox in part_bboxes.items()}
|
||||
tracked[part_name][person_idx] = id_coordinates
|
||||
|
||||
for class_name, class_data in tracked.items():
|
||||
for class_id in class_data.keys():
|
||||
class_id_str = str(class_id)
|
||||
# Use the incoming prompt for each class name and ID
|
||||
_class_name = class_name.replace('L', '').replace('R', '').lower()
|
||||
prompt_string += f'"{class_id_str}.{class_name}": "({_class_name})",\n'
|
||||
|
||||
return (tracked, prompt_string)
|
||||
|
||||
|
||||
def numpy2torch(np_image: np.ndarray) -> torch.Tensor:
|
||||
""" [H, W, C] => [B=1, H, W, C]"""
|
||||
return torch.from_numpy(np_image.astype(np.float32) / 255).unsqueeze(0)
|
||||
|
||||
|
||||
class RenderPeopleKps:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"kps": ("POSE_KEYPOINT",),
|
||||
"render_body": ("BOOLEAN", {"default": True}),
|
||||
"render_hand": ("BOOLEAN", {"default": True}),
|
||||
"render_face": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "render"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
|
||||
def render(self, kps, render_body, render_hand, render_face) -> tuple[np.ndarray]:
|
||||
if isinstance(kps, list):
|
||||
kps = kps[0]
|
||||
|
||||
poses, _, height, width = decode_json_as_poses(kps)
|
||||
np_image = draw_poses(
|
||||
poses,
|
||||
height,
|
||||
width,
|
||||
render_body,
|
||||
render_hand,
|
||||
render_face,
|
||||
)
|
||||
return (numpy2torch(np_image),)
|
||||
|
||||
class RenderAnimalKps:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"kps": ("POSE_KEYPOINT",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "render"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
|
||||
def render(self, kps) -> tuple[np.ndarray]:
|
||||
if isinstance(kps, list):
|
||||
kps = kps[0]
|
||||
|
||||
_, poses, height, width = decode_json_as_poses(kps)
|
||||
np_image = draw_animalposes(poses, height, width)
|
||||
return (numpy2torch(np_image),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SavePoseKpsAsJsonFile": SavePoseKpsAsJsonFile,
|
||||
"FacialPartColoringFromPoseKps": FacialPartColoringFromPoseKps,
|
||||
"UpperBodyTrackingFromPoseKps": UpperBodyTrackingFromPoseKps,
|
||||
"RenderPeopleKps": RenderPeopleKps,
|
||||
"RenderAnimalKps": RenderAnimalKps,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SavePoseKpsAsJsonFile": "Save Pose Keypoints",
|
||||
"FacialPartColoringFromPoseKps": "Colorize Facial Parts from PoseKPS",
|
||||
"UpperBodyTrackingFromPoseKps": "Upper Body Tracking From PoseKps (InstanceDiffusion)",
|
||||
"RenderPeopleKps": "Render Pose JSON (Human)",
|
||||
"RenderAnimalKps": "Render Pose JSON (Animal)",
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class PyraCanny_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
low_threshold=INPUT.INT(default=64, max=255),
|
||||
high_threshold=INPUT.INT(default=128, max=255),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, low_threshold=64, high_threshold=128, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.pyracanny import PyraCannyDetector
|
||||
|
||||
return (common_annotator_call(PyraCannyDetector(), image, low_threshold=low_threshold, high_threshold=high_threshold, resolution=resolution), )
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PyraCannyPreprocessor": PyraCanny_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PyraCannyPreprocessor": "PyraCanny"
|
||||
}
|
||||
46
custom_nodes/comfyui_controlnet_aux/node_wrappers/recolor.py
Normal file
46
custom_nodes/comfyui_controlnet_aux/node_wrappers/recolor.py
Normal file
@@ -0,0 +1,46 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
|
||||
class ImageLuminanceDetector:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
#https://github.com/Mikubill/sd-webui-controlnet/blob/416c345072c9c2066101e225964e3986abe6945e/scripts/processor.py#L1229
|
||||
return define_preprocessor_inputs(
|
||||
gamma_correction=INPUT.FLOAT(default=1.0, min=0.1, max=2.0),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Recolor"
|
||||
|
||||
def execute(self, image, gamma_correction=1.0, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.recolor import Recolorizer
|
||||
return (common_annotator_call(Recolorizer(), image, mode="luminance", gamma_correction=gamma_correction , resolution=resolution), )
|
||||
|
||||
class ImageIntensityDetector:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
#https://github.com/Mikubill/sd-webui-controlnet/blob/416c345072c9c2066101e225964e3986abe6945e/scripts/processor.py#L1229
|
||||
return define_preprocessor_inputs(
|
||||
gamma_correction=INPUT.FLOAT(default=1.0, min=0.1, max=2.0),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Recolor"
|
||||
|
||||
def execute(self, image, gamma_correction=1.0, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.recolor import Recolorizer
|
||||
return (common_annotator_call(Recolorizer(), image, mode="intensity", gamma_correction=gamma_correction , resolution=resolution), )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageLuminanceDetector": ImageLuminanceDetector,
|
||||
"ImageIntensityDetector": ImageIntensityDetector
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageLuminanceDetector": "Image Luminance",
|
||||
"ImageIntensityDetector": "Image Intensity"
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, nms
|
||||
import comfy.model_management as model_management
|
||||
import cv2
|
||||
|
||||
class Scribble_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.scribble import ScribbleDetector
|
||||
|
||||
model = ScribbleDetector()
|
||||
return (common_annotator_call(model, image, resolution=resolution), )
|
||||
|
||||
class Scribble_XDoG_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
threshold=INPUT.INT(default=32, min=1, max=64),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, threshold=32, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.scribble import ScribbleXDog_Detector
|
||||
|
||||
model = ScribbleXDog_Detector()
|
||||
return (common_annotator_call(model, image, resolution=resolution, thr_a=threshold), )
|
||||
|
||||
class Scribble_PiDiNet_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe=(["enable", "disable"],),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, safe="enable", resolution=512):
|
||||
def model(img, **kwargs):
|
||||
from custom_controlnet_aux.pidi import PidiNetDetector
|
||||
pidinet = PidiNetDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
result = pidinet(img, scribble=True, **kwargs)
|
||||
result = nms(result, 127, 3.0)
|
||||
result = cv2.GaussianBlur(result, (0, 0), 3.0)
|
||||
result[result > 4] = 255
|
||||
result[result < 255] = 0
|
||||
return result
|
||||
return (common_annotator_call(model, image, resolution=resolution, safe=safe=="enable"),)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ScribblePreprocessor": Scribble_Preprocessor,
|
||||
"Scribble_XDoG_Preprocessor": Scribble_XDoG_Preprocessor,
|
||||
"Scribble_PiDiNet_Preprocessor": Scribble_PiDiNet_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ScribblePreprocessor": "Scribble Lines",
|
||||
"Scribble_XDoG_Preprocessor": "Scribble XDoG Lines",
|
||||
"Scribble_PiDiNet_Preprocessor": "Scribble PiDiNet Lines"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class SAM_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/others"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.sam import SamDetector
|
||||
|
||||
mobile_sam = SamDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(mobile_sam, image, resolution=resolution)
|
||||
del mobile_sam
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SAMPreprocessor": SAM_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SAMPreprocessor": "SAM Segmentor"
|
||||
}
|
||||
27
custom_nodes/comfyui_controlnet_aux/node_wrappers/shuffle.py
Normal file
27
custom_nodes/comfyui_controlnet_aux/node_wrappers/shuffle.py
Normal file
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, MAX_RESOLUTION
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Shuffle_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
resolution=INPUT.RESOLUTION(),
|
||||
seed=INPUT.SEED()
|
||||
)
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "preprocess"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/T2IAdapter-only"
|
||||
|
||||
def preprocess(self, image, resolution=512, seed=0):
|
||||
from custom_controlnet_aux.shuffle import ContentShuffleDetector
|
||||
|
||||
return (common_annotator_call(ContentShuffleDetector(), image, resolution=resolution, seed=seed), )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ShufflePreprocessor": Shuffle_Preprocessor
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ShufflePreprocessor": "Content Shuffle"
|
||||
}
|
||||
30
custom_nodes/comfyui_controlnet_aux/node_wrappers/teed.py
Normal file
30
custom_nodes/comfyui_controlnet_aux/node_wrappers/teed.py
Normal file
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class TEED_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe_steps=INPUT.INT(default=2, max=10),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, safe_steps=2, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.teed import TEDDetector
|
||||
|
||||
model = TEDDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, safe_steps=safe_steps)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"TEEDPreprocessor": TEED_Preprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TEED_Preprocessor": "TEED Soft-Edge Lines",
|
||||
}
|
||||
73
custom_nodes/comfyui_controlnet_aux/node_wrappers/tile.py
Normal file
73
custom_nodes/comfyui_controlnet_aux/node_wrappers/tile.py
Normal file
@@ -0,0 +1,73 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
|
||||
|
||||
class Tile_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
pyrUp_iters=INPUT.INT(default=3, min=1, max=10),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/tile"
|
||||
|
||||
def execute(self, image, pyrUp_iters, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.tile import TileDetector
|
||||
|
||||
return (common_annotator_call(TileDetector(), image, pyrUp_iters=pyrUp_iters, resolution=resolution),)
|
||||
|
||||
class TTPlanet_TileGF_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
scale_factor=INPUT.FLOAT(default=1.00, min=1.000, max=8.00),
|
||||
blur_strength=INPUT.FLOAT(default=2.0, min=1.0, max=10.0),
|
||||
radius=INPUT.INT(default=7, min=1, max=20),
|
||||
eps=INPUT.FLOAT(default=0.01, min=0.001, max=0.1, step=0.001),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/tile"
|
||||
|
||||
def execute(self, image, scale_factor, blur_strength, radius, eps, **kwargs):
|
||||
from custom_controlnet_aux.tile import TTPlanet_Tile_Detector_GF
|
||||
|
||||
return (common_annotator_call(TTPlanet_Tile_Detector_GF(), image, scale_factor=scale_factor, blur_strength=blur_strength, radius=radius, eps=eps),)
|
||||
|
||||
class TTPlanet_TileSimple_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
scale_factor=INPUT.FLOAT(default=1.00, min=1.000, max=8.00),
|
||||
blur_strength=INPUT.FLOAT(default=2.0, min=1.0, max=10.0),
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/tile"
|
||||
|
||||
def execute(self, image, scale_factor, blur_strength):
|
||||
from custom_controlnet_aux.tile import TTPLanet_Tile_Detector_Simple
|
||||
|
||||
return (common_annotator_call(TTPLanet_Tile_Detector_Simple(), image, scale_factor=scale_factor, blur_strength=blur_strength),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"TilePreprocessor": Tile_Preprocessor,
|
||||
"TTPlanet_TileGF_Preprocessor": TTPlanet_TileGF_Preprocessor,
|
||||
"TTPlanet_TileSimple_Preprocessor": TTPlanet_TileSimple_Preprocessor
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TilePreprocessor": "Tile",
|
||||
"TTPlanet_TileGF_Preprocessor": "TTPlanet Tile GuidedFilter",
|
||||
"TTPlanet_TileSimple_Preprocessor": "TTPlanet Tile Simple"
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
import os
|
||||
# Disable NPU device initialization and problematic MMCV ops to prevent RuntimeError
|
||||
os.environ['NPU_DEVICE_COUNT'] = '0'
|
||||
os.environ['MMCV_WITH_OPS'] = '0'
|
||||
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Uniformer_SemSegPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "semantic_segmentate"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Semantic Segmentation"
|
||||
|
||||
def semantic_segmentate(self, image, resolution=512):
|
||||
from custom_controlnet_aux.uniformer import UniformerSegmentor
|
||||
|
||||
model = UniformerSegmentor.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"UniFormer-SemSegPreprocessor": Uniformer_SemSegPreprocessor,
|
||||
"SemSegPreprocessor": Uniformer_SemSegPreprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UniFormer-SemSegPreprocessor": "UniFormer Segmentor",
|
||||
"SemSegPreprocessor": "Semantic Segmentor (legacy, alias for UniFormer)",
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
from ..utils import common_annotator_call
|
||||
import comfy.model_management as model_management
|
||||
import torch
|
||||
import numpy as np
|
||||
from einops import rearrange
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Unimatch_OptFlowPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": dict(
|
||||
image=("IMAGE",),
|
||||
ckpt_name=(
|
||||
["gmflow-scale1-mixdata.pth", "gmflow-scale2-mixdata.pth", "gmflow-scale2-regrefine6-mixdata.pth"],
|
||||
{"default": "gmflow-scale2-regrefine6-mixdata.pth"}
|
||||
),
|
||||
backward_flow=("BOOLEAN", {"default": False}),
|
||||
bidirectional_flow=("BOOLEAN", {"default": False})
|
||||
)
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OPTICAL_FLOW", "IMAGE")
|
||||
RETURN_NAMES = ("OPTICAL_FLOW", "PREVIEW_IMAGE")
|
||||
FUNCTION = "estimate"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Optical Flow"
|
||||
|
||||
def estimate(self, image, ckpt_name, backward_flow=False, bidirectional_flow=False):
|
||||
assert len(image) > 1, "[Unimatch] Requiring as least two frames as an optical flow estimator. Only use this node on video input."
|
||||
from custom_controlnet_aux.unimatch import UnimatchDetector
|
||||
tensor_images = image
|
||||
model = UnimatchDetector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())
|
||||
flows, vis_flows = [], []
|
||||
for i in range(len(tensor_images) - 1):
|
||||
image0, image1 = np.asarray(image[i:i+2].cpu() * 255., dtype=np.uint8)
|
||||
flow, vis_flow = model(image0, image1, output_type="np", pred_bwd_flow=backward_flow, pred_bidir_flow=bidirectional_flow)
|
||||
flows.append(torch.from_numpy(flow).float())
|
||||
vis_flows.append(torch.from_numpy(vis_flow).float() / 255.)
|
||||
del model
|
||||
return (torch.stack(flows, dim=0), torch.stack(vis_flows, dim=0))
|
||||
|
||||
class MaskOptFlow:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": dict(optical_flow=("OPTICAL_FLOW",), mask=("MASK",))
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OPTICAL_FLOW", "IMAGE")
|
||||
RETURN_NAMES = ("OPTICAL_FLOW", "PREVIEW_IMAGE")
|
||||
FUNCTION = "mask_opt_flow"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Optical Flow"
|
||||
|
||||
def mask_opt_flow(self, optical_flow, mask):
|
||||
from custom_controlnet_aux.unimatch import flow_to_image
|
||||
assert len(mask) >= len(optical_flow), f"Not enough masks to mask optical flow: {len(mask)} vs {len(optical_flow)}"
|
||||
mask = mask[:optical_flow.shape[0]]
|
||||
mask = F.interpolate(mask, optical_flow.shape[1:3])
|
||||
mask = rearrange(mask, "n 1 h w -> n h w 1")
|
||||
vis_flows = torch.stack([torch.from_numpy(flow_to_image(flow)).float() / 255. for flow in optical_flow.numpy()], dim=0)
|
||||
vis_flows *= mask
|
||||
optical_flow *= mask
|
||||
return (optical_flow, vis_flows)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Unimatch_OptFlowPreprocessor": Unimatch_OptFlowPreprocessor,
|
||||
"MaskOptFlow": MaskOptFlow
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Unimatch_OptFlowPreprocessor": "Unimatch Optical Flow",
|
||||
"MaskOptFlow": "Mask Optical Flow (DragNUWA)"
|
||||
}
|
||||
27
custom_nodes/comfyui_controlnet_aux/node_wrappers/zoe.py
Normal file
27
custom_nodes/comfyui_controlnet_aux/node_wrappers/zoe.py
Normal file
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Zoe_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.zoe import ZoeDetector
|
||||
|
||||
model = ZoeDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Zoe-DepthMapPreprocessor": Zoe_Depth_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Zoe-DepthMapPreprocessor": "Zoe Depth Map"
|
||||
}
|
||||
Reference in New Issue
Block a user