107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
import os
|
|
from typing import Dict
|
|
|
|
import folder_paths
|
|
|
|
from ..utils import load_module
|
|
|
|
custom_nodes = folder_paths.get_folder_paths("custom_nodes")
|
|
preprocessors_dir_names = ["ControlNetPreprocessors", "comfyui_controlnet_aux"]
|
|
|
|
preprocessors: list[str] = []
|
|
_preprocessors_map = {
|
|
"canny": "CannyEdgePreprocessor",
|
|
"canny_pyra": "PyraCannyPreprocessor",
|
|
"lineart": "LineArtPreprocessor",
|
|
"lineart_anime": "AnimeLineArtPreprocessor",
|
|
"lineart_manga": "Manga2Anime_LineArt_Preprocessor",
|
|
"lineart_any": "AnyLineArtPreprocessor_aux",
|
|
"scribble": "ScribblePreprocessor",
|
|
"scribble_xdog": "Scribble_XDoG_Preprocessor",
|
|
"scribble_pidi": "Scribble_PiDiNet_Preprocessor",
|
|
"scribble_hed": "FakeScribblePreprocessor",
|
|
"hed": "HEDPreprocessor",
|
|
"pidi": "PiDiNetPreprocessor",
|
|
"mlsd": "M-LSDPreprocessor",
|
|
"pose": "DWPreprocessor",
|
|
"openpose": "OpenposePreprocessor",
|
|
"dwpose": "DWPreprocessor",
|
|
"pose_dense": "DensePosePreprocessor",
|
|
"pose_animal": "AnimalPosePreprocessor",
|
|
"normalmap_bae": "BAE-NormalMapPreprocessor",
|
|
"normalmap_dsine": "DSINE-NormalMapPreprocessor",
|
|
"normalmap_midas": "MiDaS-NormalMapPreprocessor",
|
|
"depth": "DepthAnythingV2Preprocessor",
|
|
"depth_anything": "DepthAnythingPreprocessor",
|
|
"depth_anything_v2": "DepthAnythingV2Preprocessor",
|
|
"depth_anything_zoe": "Zoe_DepthAnythingPreprocessor",
|
|
"depth_zoe": "Zoe-DepthMapPreprocessor",
|
|
"depth_midas": "MiDaS-DepthMapPreprocessor",
|
|
"depth_leres": "LeReS-DepthMapPreprocessor",
|
|
"depth_metric3d": "Metric3D-DepthMapPreprocessor",
|
|
"depth_meshgraphormer": "MeshGraphormer-DepthMapPreprocessor",
|
|
"seg_ofcoco": "OneFormer-COCO-SemSegPreprocessor",
|
|
"seg_ofade20k": "OneFormer-ADE20K-SemSegPreprocessor",
|
|
"seg_ufade20k": "UniFormer-SemSegPreprocessor",
|
|
"seg_animeface": "AnimeFace_SemSegPreprocessor",
|
|
"shuffle": "ShufflePreprocessor",
|
|
"teed": "TEEDPreprocessor",
|
|
"color": "ColorPreprocessor",
|
|
"sam": "SAMPreprocessor",
|
|
"tile": "TilePreprocessor"
|
|
}
|
|
|
|
|
|
def apply_preprocessor(image, preprocessor, resolution=512):
|
|
raise NotImplementedError("apply_preprocessor is not implemented")
|
|
|
|
|
|
try:
|
|
module_path = None
|
|
|
|
for custom_node in custom_nodes:
|
|
custom_node = custom_node if not os.path.islink(custom_node) else os.readlink(custom_node)
|
|
for module_dir in preprocessors_dir_names:
|
|
if module_dir in os.listdir(custom_node):
|
|
module_path = os.path.abspath(os.path.join(custom_node, module_dir))
|
|
break
|
|
|
|
if module_path is None:
|
|
raise Exception("Could not find ControlNetPreprocessors nodes")
|
|
|
|
module = load_module(module_path)
|
|
print("Loaded ControlNetPreprocessors nodes from", module_path)
|
|
|
|
nodes: Dict = getattr(module, "NODE_CLASS_MAPPINGS")
|
|
available_preprocessors: list[str] = getattr(module, "PREPROCESSOR_OPTIONS")
|
|
|
|
AIO_Preprocessor = nodes.get("AIO_Preprocessor", None)
|
|
if AIO_Preprocessor is None:
|
|
raise Exception("Could not find AIO_Preprocessor node")
|
|
|
|
for name, preprocessor in _preprocessors_map.items():
|
|
if preprocessor in available_preprocessors:
|
|
preprocessors.append(name)
|
|
|
|
aio_preprocessor = AIO_Preprocessor()
|
|
|
|
def apply_preprocessor(image, preprocessor, resolution=512):
|
|
if preprocessor == "None":
|
|
return image
|
|
|
|
if preprocessor not in preprocessors:
|
|
raise Exception(f"Preprocessor {preprocessor} is not implemented")
|
|
|
|
preprocessor_cls = _preprocessors_map[preprocessor]
|
|
args = {"preprocessor": preprocessor_cls, "image": image, "resolution": resolution}
|
|
|
|
function_name = AIO_Preprocessor.FUNCTION
|
|
res = getattr(aio_preprocessor, function_name)(**args)
|
|
if isinstance(res, dict):
|
|
res = res["result"]
|
|
|
|
return res[0]
|
|
|
|
except Exception as e:
|
|
print(e)
|