From bd493d42d41fbb508e8e83e9915d13fecbde79aa Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Thu, 19 Oct 2023 00:21:59 +0900 Subject: [PATCH] feat: support Inpaint Preprocessors for SEGS --- README.md | 2 +- __init__.py | 2 +- segs_support.py | 54 +++++++++++++++++++++++++++++++++++++------------ 3 files changed, 43 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index fbefbe1..213a1b7 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ This repository offers various extension nodes for ComfyUI. Nodes here have diff * `DW Preprocessor Provider (SEGS)`, `MiDaS Depth Map Preprocessor Provider (SEGS)`, `LeReS Depth Map Preprocessor Provider (SEGS)`, `MediaPipe FaceMesh Preprocessor Provider (SEGS)`, `HED Preprocessor Provider (SEGS)`, `Fake Scribble Preprocessor (SEGS)`, `AnimeLineArt Preprocessor Provider (SEGS)`, `Manga2Anime LineArt Preprocessor Provider (SEGS)`, `LineArt Preprocessor Provider (SEGS)`, - `Color Preprocessor Provider (SEGS)` + `Color Preprocessor Provider (SEGS)`, `Inpaint Preprocessor Provider (SEGS)`, * `MediaPipeFaceMeshDetectorProvider`: This node provides `BBOX_DETECTOR` and `SEGM_DETECTOR` that can be used in Impact Pack's Detector using the `MediaPipe-FaceMesh Preprocessor` of ControlNet Auxiliary Preprocessors. * A1111 Compatibility support - These nodes assists in replicating the creation of A1111 in ComfyUI exactly. diff --git a/__init__.py b/__init__.py index 74e7284..e5787e5 100644 --- a/__init__.py +++ b/__init__.py @@ -7,7 +7,7 @@ import importlib -print(f"### Loading: ComfyUI-Inspire-Pack (V0.31)") +print(f"### Loading: ComfyUI-Inspire-Pack (V0.32)") node_list = [ "lora_block_weight", diff --git a/segs_support.py b/segs_support.py index f05ed1b..8f3568f 100644 --- a/segs_support.py +++ b/segs_support.py @@ -1,6 +1,6 @@ import nodes import numpy as np - +import torch def normalize_size_base_64(w, h): short_side = min(w, h) @@ -50,7 +50,7 @@ class MediaPipe_FaceMesh_Preprocessor_wrapper: self.min_confidence = min_confidence self.upscale_factor = upscale_factor - def apply(self, image): + def apply(self, image, mask=None): if 'MediaPipe-FaceMeshPreprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use MediaPipeFaceMeshDetector, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -63,7 +63,7 @@ class MediaPipe_FaceMesh_Preprocessor_wrapper: class AnimeLineArt_Preprocessor_wrapper: - def apply(self, image): + def apply(self, image, mask=None): if 'AnimeLineArtPreprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use AnimeLineArt_Preprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -73,7 +73,7 @@ class AnimeLineArt_Preprocessor_wrapper: class Manga2Anime_LineArt_Preprocessor_wrapper: - def apply(self, image): + def apply(self, image, mask=None): if 'Manga2Anime_LineArt_Preprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use Manga2Anime_LineArt_Preprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -83,7 +83,7 @@ class Manga2Anime_LineArt_Preprocessor_wrapper: class Color_Preprocessor_wrapper: - def apply(self, image): + def apply(self, image, mask=None): if 'ColorPreprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use Color_Preprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -92,11 +92,23 @@ class Color_Preprocessor_wrapper: return obj.execute(image, resolution=resolution)[0] +class InpaintPreprocessor_wrapper: + def apply(self, image, mask=None): + if 'InpaintPreprocessor' not in nodes.NODE_CLASS_MAPPINGS: + raise Exception(f"[ERROR] To use InpaintPreprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") + + obj = nodes.NODE_CLASS_MAPPINGS['InpaintPreprocessor']() + if mask is None: + mask = torch.ones((image.shape[1], image.shape[2]), dtype=torch.float32, device="cpu").unsqueeze(0) + + return obj.preprocess(image, mask)[0] + + class LineArt_Preprocessor_wrapper: def __init__(self, coarse): self.coarse = coarse - def apply(self, image): + def apply(self, image, mask=None): if 'LineArtPreprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use LineArt_Preprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -114,7 +126,7 @@ class OpenPose_Preprocessor_wrapper: self.detect_face = detect_face self.upscale_factor = upscale_factor - def apply(self, image): + def apply(self, image, mask=None): if 'OpenposePreprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use OpenPose_Preprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -137,7 +149,7 @@ class DWPreprocessor_wrapper: self.detect_face = detect_face self.upscale_factor = upscale_factor - def apply(self, image): + def apply(self, image, mask=None): if 'DWPreprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use DWPreprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -159,7 +171,7 @@ class LeReS_DepthMap_Preprocessor_wrapper: self.rm_background = rm_background self.boost = boost - def apply(self, image): + def apply(self, image, mask=None): if 'LeReS-DepthMapPreprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use LeReS_DepthMap_Preprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -175,7 +187,7 @@ class MiDaS_DepthMap_Preprocessor_wrapper: self.a = a self.bg_threshold = bg_threshold - def apply(self, image): + def apply(self, image, mask=None): if 'MiDaS-DepthMapPreprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use MiDaS_DepthMap_Preprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -185,7 +197,7 @@ class MiDaS_DepthMap_Preprocessor_wrapper: class Zoe_DepthMap_Preprocessor_wrapper: - def apply(self, image): + def apply(self, image, mask=None): if 'Zoe-DepthMapPreprocessor' not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use Zoe_DepthMap_Preprocessor_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -199,7 +211,7 @@ class HED_Preprocessor_wrapper: self.safe = safe self.nodename = nodename - def apply(self, image): + def apply(self, image, mask=None): if self.nodename not in nodes.NODE_CLASS_MAPPINGS: raise Exception(f"[ERROR] To use {self.nodename}_Provider, you need to install 'ComfyUI's ControlNet Auxiliary Preprocessors.'") @@ -213,7 +225,7 @@ class Canny_Preprocessor_wrapper: self.low_threshold = low_threshold self.high_threshold = high_threshold - def apply(self, image): + def apply(self, image, mask=None): obj = nodes.NODE_CLASS_MAPPINGS['Canny']() return obj.detect_edge(image, self.low_threshold, self.high_threshold)[0] @@ -465,6 +477,20 @@ class Color_Preprocessor_Provider_for_SEGS: return (obj, ) +class InpaintPreprocessor_Provider_for_SEGS: + @classmethod + def INPUT_TYPES(s): + return {"required": {}} + RETURN_TYPES = ("SEGS_PREPROCESSOR",) + FUNCTION = "doit" + + CATEGORY = "InspirePack/SEGS/ControlNet" + + def doit(self): + obj = InpaintPreprocessor_wrapper() + return (obj, ) + + NODE_CLASS_MAPPINGS = { "OpenPose_Preprocessor_Provider_for_SEGS //Inspire": OpenPose_Preprocessor_Provider_for_SEGS, "DWPreprocessor_Provider_for_SEGS //Inspire": DWPreprocessor_Provider_for_SEGS, @@ -480,6 +506,7 @@ NODE_CLASS_MAPPINGS = { "Manga2Anime_LineArt_Preprocessor_Provider_for_SEGS //Inspire": Manga2Anime_LineArt_Preprocessor_Provider_for_SEGS, "LineArt_Preprocessor_Provider_for_SEGS //Inspire": LineArt_Preprocessor_Provider_for_SEGS, "Color_Preprocessor_Provider_for_SEGS //Inspire": Color_Preprocessor_Provider_for_SEGS, + "InpaintPreprocessor_Provider_for_SEGS //Inspire": InpaintPreprocessor_Provider_for_SEGS, } NODE_DISPLAY_NAME_MAPPINGS = { "OpenPose_Preprocessor_Provider_for_SEGS //Inspire": "OpenPose Preprocessor Provider (SEGS)", @@ -495,5 +522,6 @@ NODE_DISPLAY_NAME_MAPPINGS = { "Manga2Anime_LineArt_Preprocessor_Provider_for_SEGS //Inspire": "Manga2Anime LineArt Preprocessor Provider (SEGS)", "LineArt_Preprocessor_Provider_for_SEGS //Inspire": "LineArt Preprocessor Provider (SEGS)", "Color_Preprocessor_Provider_for_SEGS //Inspire": "Color Preprocessor Provider (SEGS)", + "InpaintPreprocessor_Provider_for_SEGS //Inspire": "Inpaint Preprocessor Provider (SEGS)", "MediaPipeFaceMeshDetectorProvider //Inspire": "MediaPipeFaceMesh Detector Provider", }