├── requirements.txt ├── README.md ├── __init__.py └── preprocessor.py /requirements.txt: -------------------------------------------------------------------------------- 1 | opencv-python 2 | controlnet-aux==0.0.6 3 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ControlNet Preprocessors for ComfyUI 2 | ## All available preprocessors 3 | 4 | ![screenshot](https://github.com/kenneth2001/comfyui_controlnet_preprocessors/assets/24566737/b9bd6973-d7ea-4df6-aae2-4cfb94cff052) 5 | 6 | ## Installation 7 | 8 | ``` 9 | cd ComfyUI/custom_nodes 10 | git clone https://github.com/kenneth2001/comfyui_controlnet_preprocessors 11 | ``` 12 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | from custom_nodes.comfyui_controlnet_preprocessors import preprocessor 2 | 3 | # NODE_CLASS_MAPPINGS = { 4 | # "Canny_Detector_Preprocessor": Canny_Detector_Preprocessor, 5 | # "Depth_detetor_Preprocessor": Depth_detetor_Preprocessor, 6 | # "Normal_Bae_Detector_Preprocessor": Normal_Bae_Detector_Preprocessor, 7 | # "Openpose_Detector_Preprocessor": Openpose_Detector_Preprocessor, 8 | # "MLSD_Detector_Preprocessor": MLSD_Detector_Preprocessor, 9 | # "Lineart_Detector_Preprocessor": Lineart_Detector_Preprocessor, 10 | # "Softedge_Detector_Preprocessor": Softedge_Detector_Preprocessor, 11 | # "Scribble_Detector_Preprocessor": Scribble_Detector_Preprocessor, 12 | # "Shuffle_Detector_Preprocessor": Shuffle_Detector_Preprocessor, 13 | # } 14 | 15 | NODE_CLASS_MAPPINGS = { 16 | **preprocessor.NODE_CLASS_MAPPINGS 17 | } 18 | -------------------------------------------------------------------------------- /preprocessor.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import cv2 3 | from controlnet_aux import CannyDetector, MidasDetector, ZoeDetector, LeresDetector, NormalBaeDetector, OpenposeDetector, MLSDdetector, LineartDetector, LineartAnimeDetector, HEDdetector, PidiNetDetector, ContentShuffleDetector 4 | import numpy as np 5 | from torch.hub import download_url_to_file, get_dir 6 | import os 7 | from urllib.parse import urlparse 8 | 9 | 10 | def img_tensor_to_np(img_tensor): 11 | img_tensor = img_tensor.clone() * 255.0 12 | return img_tensor.squeeze().numpy().astype(np.uint8) 13 | 14 | 15 | def img_np_to_tensor(img_np_list): 16 | return torch.from_numpy(img_np_list.astype(np.float32) / 255.0).unsqueeze(0) 17 | 18 | 19 | def load_file_from_url(url, model_dir=None, progress=True, file_name=None): 20 | if model_dir is None: # use the pytorch hub_dir 21 | hub_dir = get_dir() 22 | model_dir = os.path.join(hub_dir, 'checkpoints') 23 | 24 | os.makedirs(model_dir, exist_ok=True) 25 | 26 | parts = urlparse(url) 27 | filename = os.path.basename(parts.path) 28 | if file_name is not None: 29 | filename = file_name 30 | cached_file = os.path.abspath(os.path.join(model_dir, filename)) 31 | if not os.path.exists(cached_file): 32 | print(f'Downloading: "{url}" to {cached_file}\n') 33 | download_url_to_file( 34 | url, cached_file, hash_prefix=None, progress=progress) 35 | return cached_file 36 | 37 | 38 | class Canny_Detector_Preprocessor: 39 | @classmethod 40 | def INPUT_TYPES(s): 41 | return {"required": {"image": ("IMAGE",), }} 42 | 43 | RETURN_TYPES = ("IMAGE",) 44 | FUNCTION = "controlnet_preprocessor" 45 | 46 | CATEGORY = "controlnet_preprocessor" 47 | 48 | def controlnet_preprocessor(self, image): 49 | image_np = img_tensor_to_np(image) 50 | 51 | canny = CannyDetector() 52 | processed_image = canny(image_np) 53 | 54 | processed_image_np = np.array(processed_image) 55 | processed_image_tensor = img_np_to_tensor(processed_image_np) 56 | return (processed_image_tensor,) 57 | 58 | 59 | class Depth_Detector_Preprocessor: 60 | @classmethod 61 | def INPUT_TYPES(s): 62 | return {"required": {"image": ("IMAGE",), 63 | "mode": (["depth_midas", "depth_zoe", "depth_leres", "depth_leres++"], {"default": "depth_midas"}), }} 64 | 65 | RETURN_TYPES = ("IMAGE",) 66 | FUNCTION = "controlnet_preprocessor" 67 | 68 | CATEGORY = "controlnet_preprocessor" 69 | 70 | def controlnet_preprocessor(self, image, mode): 71 | image_np = img_tensor_to_np(image) 72 | 73 | try: 74 | if mode == 'depth_midas': 75 | depth = MidasDetector.from_pretrained('ckpts') 76 | elif mode == 'depth_zoe': 77 | depth = ZoeDetector.from_pretrained('ckpts') 78 | elif mode in ['depth_leres', 'depth_leres++']: 79 | depth = LeresDetector.from_pretrained('ckpts') 80 | except: 81 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/ZoeD_M12_N.pt', 82 | model_dir='ckpts', file_name='ZoeD_M12_N.pt') 83 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/dpt_hybrid-midas-501f0c75.pt', 84 | model_dir='ckpts', file_name='dpt_hybrid-midas-501f0c75.pt') 85 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/res101.pth', 86 | model_dir='ckpts', file_name='res101.pth') 87 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/latest_net_G.pth', 88 | model_dir='ckpts', file_name='latest_net_G.pth') 89 | if mode == 'depth_midas': 90 | depth = MidasDetector.from_pretrained('ckpts') 91 | elif mode == 'depth_zoe': 92 | depth = ZoeDetector.from_pretrained('ckpts') 93 | elif mode in ['depth_leres', 'depth_leres++']: 94 | depth = LeresDetector.from_pretrained('ckpts') 95 | 96 | if mode == 'depth_leres++': 97 | processed_image = depth(image_np, boost=True) 98 | elif mode == 'depth_leres': 99 | processed_image = depth(image_np, boost=False) 100 | else: 101 | processed_image = depth(image_np) 102 | 103 | processed_image_np = np.array(processed_image) 104 | processed_image_tensor = img_np_to_tensor(processed_image_np) 105 | return (processed_image_tensor,) 106 | 107 | 108 | class Normal_Bae_Detector_Preprocessor: 109 | @classmethod 110 | def INPUT_TYPES(s): 111 | return {"required": {"image": ("IMAGE",), }} 112 | 113 | RETURN_TYPES = ("IMAGE",) 114 | FUNCTION = "controlnet_preprocessor" 115 | 116 | CATEGORY = "controlnet_preprocessor" 117 | 118 | def controlnet_preprocessor(self, image): 119 | image_np = img_tensor_to_np(image) 120 | 121 | try: 122 | normal_bae = NormalBaeDetector.from_pretrained('ckpts') 123 | except: 124 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/scannet.pt', 125 | model_dir='ckpts', file_name='scannet.pt') 126 | normal_bae = NormalBaeDetector.from_pretrained('ckpts') 127 | 128 | processed_image = normal_bae(image_np) 129 | 130 | processed_image_np = np.array(processed_image) 131 | processed_image_tensor = img_np_to_tensor(processed_image_np) 132 | return (processed_image_tensor,) 133 | 134 | 135 | class Openpose_Detector_Preprocessor: 136 | @classmethod 137 | def INPUT_TYPES(s): 138 | return {"required": {"image": ("IMAGE",), 139 | "include_face": ([True, False], {"default": False}), 140 | "include_hand": ([True, False], {"default": False}), 141 | "include_body": ([True, False], {"default": True})}} 142 | 143 | RETURN_TYPES = ("IMAGE",) 144 | FUNCTION = "controlnet_preprocessor" 145 | 146 | CATEGORY = "controlnet_preprocessor" 147 | 148 | def controlnet_preprocessor(self, image, include_face, include_hand, include_body): 149 | image_np = img_tensor_to_np(image) 150 | 151 | try: 152 | openpose = OpenposeDetector.from_pretrained('ckpts') 153 | except: 154 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/body_pose_model.pth', 155 | model_dir='ckpts', file_name='body_pose_model.pth') 156 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/hand_pose_model.pth', 157 | model_dir='ckpts', file_name='hand_pose_model.pth') 158 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/facenet.pth', 159 | model_dir='ckpts', file_name='facenet.pth') 160 | openpose = OpenposeDetector.from_pretrained('ckpts') 161 | 162 | processed_image = openpose( 163 | image_np, include_face=include_face, include_hand=include_hand, include_body=include_body) 164 | 165 | processed_image_np = np.array(processed_image) 166 | processed_image_tensor = img_np_to_tensor(processed_image_np) 167 | return (processed_image_tensor,) 168 | 169 | 170 | class MLSD_Detector_Preprocessor: 171 | @classmethod 172 | def INPUT_TYPES(s): 173 | return {"required": {"image": ("IMAGE",), }} 174 | 175 | RETURN_TYPES = ("IMAGE",) 176 | FUNCTION = "controlnet_preprocessor" 177 | 178 | CATEGORY = "controlnet_preprocessor" 179 | 180 | def controlnet_preprocessor(self, image): 181 | image_np = img_tensor_to_np(image) 182 | 183 | try: 184 | mlsd = MLSDdetector.from_pretrained('ckpts') 185 | except: 186 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/mlsd_large_512_fp32.pth', 187 | model_dir='ckpts', file_name='mlsd_large_512_fp32.pth') 188 | mlsd = MLSDdetector.from_pretrained('ckpts') 189 | 190 | processed_image = mlsd(image_np) 191 | 192 | processed_image_np = np.array(processed_image) 193 | processed_image_tensor = img_np_to_tensor(processed_image_np) 194 | return (processed_image_tensor,) 195 | 196 | 197 | class Lineart_Detector_Preprocessor: 198 | @classmethod 199 | def INPUT_TYPES(s): 200 | return {"required": {"image": ("IMAGE",), 201 | "mode": (["coarse", "anime", "realistic"], {"default": "anime"}), }} 202 | 203 | RETURN_TYPES = ("IMAGE",) 204 | FUNCTION = "controlnet_preprocessor" 205 | 206 | CATEGORY = "controlnet_preprocessor" 207 | 208 | def controlnet_preprocessor(self, image, mode): 209 | image_np = img_tensor_to_np(image) 210 | 211 | try: 212 | if mode in ('coarse', 'realistic'): 213 | linear_art = LineartDetector.from_pretrained('ckpts') 214 | elif mode == 'anime': 215 | linear_art = LineartAnimeDetector.from_pretrained('ckpts') 216 | except: 217 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/netG.pth', 218 | model_dir='ckpts', file_name='netG.pth') 219 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/sk_model.pth', 220 | model_dir='ckpts', file_name='sk_model.pth') 221 | linear_art = LineartDetector.from_pretrained('ckpts') if mode in ( 222 | 'coarse', 'realistic') else LineartAnimeDetector.from_pretrained('ckpts') 223 | 224 | if mode in ('coarse', 'realistic'): 225 | processed_image = linear_art( 226 | image_np, coarse=True if mode == 'coarse' else False) 227 | elif mode == 'anime': 228 | processed_image = linear_art(image_np) 229 | 230 | processed_image_np = np.array(processed_image) 231 | processed_image_tensor = img_np_to_tensor(processed_image_np) 232 | return (processed_image_tensor,) 233 | 234 | 235 | class Softedge_Detector_Preprocessor: 236 | @classmethod 237 | def INPUT_TYPES(s): 238 | return {"required": {"image": ("IMAGE",), 239 | "mode": (["softedge_hed", "softedge_hedsafe", "softedge_pidinet", "softedge_pidsafe"], {"default": "softedge_hed"}), }} 240 | 241 | RETURN_TYPES = ("IMAGE",) 242 | FUNCTION = "controlnet_preprocessor" 243 | 244 | CATEGORY = "controlnet_preprocessor" 245 | 246 | def controlnet_preprocessor(self, image, mode): 247 | image_np = img_tensor_to_np(image) 248 | 249 | try: 250 | if mode in ('softedge_hed', 'softedge_hedsafe'): 251 | softedge = HEDdetector.from_pretrained('ckpts') 252 | elif mode in ('softedge_pidinet', 'softedge_pidsafe'): 253 | softedge = PidiNetDetector.from_pretrained('ckpts') 254 | except: 255 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetHED.pth', 256 | model_dir='ckpts', file_name='ControlNetHED.pth') 257 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/table5_pidinet.pth', 258 | model_dir='ckpts', file_name='table5_pidinet.pth') 259 | softedge = HEDdetector.from_pretrained('ckpts') if mode in ( 260 | 'softedge_hed', 'softedge_hedsafe') else PidiNetDetector.from_pretrained('ckpts') 261 | 262 | processed_image = softedge(image_np, safe=True if mode in ( 263 | 'softedge_hedsafe', 'softedge_pidsafe') else False, scribble=False) 264 | 265 | processed_image_np = np.array(processed_image) 266 | processed_image_tensor = img_np_to_tensor(processed_image_np) 267 | return (processed_image_tensor,) 268 | 269 | 270 | class Scribble_Detector_Preprocessor: 271 | @classmethod 272 | def INPUT_TYPES(s): 273 | return {"required": {"image": ("IMAGE",), 274 | "mode": (["scribble_hed", "scribble_hedsafe", "scribble_pidinet", "scribble_pidsafe"], {"default": "scribble_hed"})}} 275 | 276 | RETURN_TYPES = ("IMAGE",) 277 | FUNCTION = "controlnet_preprocessor" 278 | 279 | CATEGORY = "controlnet_preprocessor" 280 | 281 | def controlnet_preprocessor(self, image, mode): 282 | image_np = img_tensor_to_np(image) 283 | 284 | try: 285 | if mode in ('scribble_hed', 'scribble_hedsafe'): 286 | scribble = HEDdetector.from_pretrained('ckpts') 287 | elif mode in ('scribble_pidinet', 'scribble_pidsafe'): 288 | scribble = PidiNetDetector.from_pretrained('ckpts') 289 | except: 290 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetHED.pth', 291 | model_dir='ckpts', file_name='ControlNetHED.pth') 292 | load_file_from_url('https://huggingface.co/lllyasviel/Annotators/resolve/main/table5_pidinet.pth', 293 | model_dir='ckpts', file_name='table5_pidinet.pth') 294 | scribble = HEDdetector.from_pretrained('ckpts') if mode in ( 295 | 'scribble_hed', 'scribble_hedsafe') else PidiNetDetector.from_pretrained('ckpts') 296 | 297 | processed_image = scribble(image_np, safe=True if mode in ( 298 | 'scribble_hedsafe', 'scribble_pidsafe') else False, scribble=True) 299 | 300 | processed_image_np = np.array(processed_image) 301 | processed_image_tensor = img_np_to_tensor(processed_image_np) 302 | return (processed_image_tensor,) 303 | 304 | 305 | class Shuffle_Detector_Preprocessor: 306 | @classmethod 307 | def INPUT_TYPES(s): 308 | return {"required": {"image": ("IMAGE",), }} 309 | 310 | RETURN_TYPES = ("IMAGE",) 311 | FUNCTION = "controlnet_preprocessor" 312 | 313 | CATEGORY = "controlnet_preprocessor" 314 | 315 | def controlnet_preprocessor(self, image): 316 | image_np = img_tensor_to_np(image) 317 | 318 | shuffle = ContentShuffleDetector() 319 | processed_image = shuffle(image_np) 320 | 321 | processed_image_np = np.array(processed_image) 322 | processed_image_tensor = img_np_to_tensor(processed_image_np) 323 | return (processed_image_tensor,) 324 | 325 | 326 | NODE_CLASS_MAPPINGS = { 327 | "Canny_Detector_Preprocessor": Canny_Detector_Preprocessor, 328 | "Depth_Detector_Preprocessor": Depth_Detector_Preprocessor, 329 | "Normal_Bae_Detector_Preprocessor": Normal_Bae_Detector_Preprocessor, 330 | "Openpose_Detector_Preprocessor": Openpose_Detector_Preprocessor, 331 | "MLSD_Detector_Preprocessor": MLSD_Detector_Preprocessor, 332 | "Lineart_Detector_Preprocessor": Lineart_Detector_Preprocessor, 333 | "Softedge_Detector_Preprocessor": Softedge_Detector_Preprocessor, 334 | "Scribble_Detector_Preprocessor": Scribble_Detector_Preprocessor, 335 | "Shuffle_Detector_Preprocessor": Shuffle_Detector_Preprocessor, 336 | } 337 | --------------------------------------------------------------------------------