From 923188f1a2bab1b74e3982a2b1bab7b904a2e60c Mon Sep 17 00:00:00 2001 From: huchenlei Date: Sun, 19 May 2024 11:07:54 -0400 Subject: [PATCH] Add Anyline Preprocessor --- annotator/teed/__init__.py | 44 +++++--- scripts/preprocessor/__init__.py | 1 + .../legacy/preprocessor_compiled.py | 54 ++++----- scripts/preprocessor/legacy/processor.py | 15 --- scripts/preprocessor/teed.py | 106 ++++++++++++++++++ 5 files changed, 165 insertions(+), 55 deletions(-) create mode 100644 scripts/preprocessor/teed.py diff --git a/annotator/teed/__init__.py b/annotator/teed/__init__.py index 47517043f..2b5c0d368 100644 --- a/annotator/teed/__init__.py +++ b/annotator/teed/__init__.py @@ -1,28 +1,34 @@ -""" -Hello, welcome on board, -""" from __future__ import print_function import os import cv2 import numpy as np - import torch - -from annotator.teed.ted import TED # TEED architecture from einops import rearrange + from modules import devices -from annotator.util import load_model,safe_step +from modules.safe import unsafe_torch_load +from annotator.teed.ted import TED # TEED architecture +from annotator.util import load_model, safe_step from annotator.annotator_path import models_path -class TEEDDector: + +class TEEDDetector: """https://github.com/xavysp/TEED""" model_dir = os.path.join(models_path, "TEED") - def __init__(self): + def __init__(self, mteed: bool = False): self.device = devices.get_device_for("controlnet") self.model = TED().to(self.device).eval() + + if mteed: + self.load_mteed_model() + else: + self.load_teed_model() + + def load_teed_model(self): + """Load vanilla TEED model""" remote_url = os.environ.get( "CONTROLNET_TEED_MODEL_URL", "https://huggingface.co/bdsqlsz/qinglong_controlnet-lllite/resolve/main/Annotators/7_model.pth", @@ -30,7 +36,17 @@ def __init__(self): model_path = load_model( "7_model.pth", remote_url=remote_url, model_dir=self.model_dir ) - self.model.load_state_dict(torch.load(model_path)) + self.model.load_state_dict(unsafe_torch_load(model_path)) + + def load_mteed_model(self): + """Load MTEED model for Anyline""" + remote_url = ( + "https://huggingface.co/TheMistoAI/MistoLine/resolve/main/Anyline/MTEED.pth" + ) + model_path = load_model( + "MTEED.pth", remote_url=remote_url, model_dir=self.model_dir + ) + self.model.load_state_dict(unsafe_torch_load(model_path)) def unload_model(self): if self.model is not None: @@ -43,13 +59,15 @@ def __call__(self, image: np.ndarray, safe_steps: int = 2) -> np.ndarray: H, W, _ = image.shape with torch.no_grad(): image_teed = torch.from_numpy(image.copy()).float().to(self.device) - image_teed = rearrange(image_teed, 'h w c -> 1 c h w') + image_teed = rearrange(image_teed, "h w c -> 1 c h w") edges = self.model(image_teed) edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges] - edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges] + edges = [ + cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges + ] edges = np.stack(edges, axis=2) edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64))) if safe_steps != 0: edge = safe_step(edge, safe_steps) edge = (edge * 255.0).clip(0, 255).astype(np.uint8) - return edge \ No newline at end of file + return edge diff --git a/scripts/preprocessor/__init__.py b/scripts/preprocessor/__init__.py index b330e73ce..081af9779 100644 --- a/scripts/preprocessor/__init__.py +++ b/scripts/preprocessor/__init__.py @@ -1,3 +1,4 @@ +from .teed import * from .pulid import * from .inpaint import * from .lama_inpaint import * diff --git a/scripts/preprocessor/legacy/preprocessor_compiled.py b/scripts/preprocessor/legacy/preprocessor_compiled.py index 31406b9f9..b5a69164a 100644 --- a/scripts/preprocessor/legacy/preprocessor_compiled.py +++ b/scripts/preprocessor/legacy/preprocessor_compiled.py @@ -1141,33 +1141,33 @@ "SoftEdge" ] }, - "te_hed": { - "label": "softedge_teed", - "call_function": te_hed, - "unload_function": unload_te_hed, - "managed_model": "model_te_hed", - "model_free": False, - "no_control_mode": False, - "resolution": { - "label": "Resolution", - "value": 512, - "minimum": 64, - "maximum": 2048 - }, - "slider_1": { - "label": "Safe Steps", - "minimum": 0, - "maximum": 10, - "value": 2, - "step": 1 - }, - "slider_2": None, - "slider_3": None, - "priority": 0, - "tags": [ - "SoftEdge" - ] - }, + # "te_hed": { + # "label": "softedge_teed", + # "call_function": te_hed, + # "unload_function": unload_te_hed, + # "managed_model": "model_te_hed", + # "model_free": False, + # "no_control_mode": False, + # "resolution": { + # "label": "Resolution", + # "value": 512, + # "minimum": 64, + # "maximum": 2048 + # }, + # "slider_1": { + # "label": "Safe Steps", + # "minimum": 0, + # "maximum": 10, + # "value": 2, + # "step": 1 + # }, + # "slider_2": None, + # "slider_3": None, + # "priority": 0, + # "tags": [ + # "SoftEdge" + # ] + # }, "color": { "label": "t2ia_color_grid", "call_function": color, diff --git a/scripts/preprocessor/legacy/processor.py b/scripts/preprocessor/legacy/processor.py index dbf55e235..e9861fb3b 100644 --- a/scripts/preprocessor/legacy/processor.py +++ b/scripts/preprocessor/legacy/processor.py @@ -625,21 +625,6 @@ def unload_densepose(): from annotator.densepose import unload_model unload_model() -model_te_hed = None - -def te_hed(img, res=512, thr_a=2, **kwargs): - img, remove_pad = resize_image_with_pad(img, res) - global model_te_hed - if model_te_hed is None: - from annotator.teed import TEEDDector - model_te_hed = TEEDDector() - result = model_te_hed(img, safe_steps=int(thr_a)) - return remove_pad(result), True - -def unload_te_hed(): - if model_te_hed is not None: - model_te_hed.unload_model() - class InsightFaceModel: def __init__(self, face_analysis_model_name: str = "buffalo_l"): self.model = None diff --git a/scripts/preprocessor/teed.py b/scripts/preprocessor/teed.py new file mode 100644 index 000000000..3fa8a9938 --- /dev/null +++ b/scripts/preprocessor/teed.py @@ -0,0 +1,106 @@ +import numpy as np +from skimage import morphology + +from annotator.teed import TEEDDetector +from annotator.util import HWC3 +from scripts.supported_preprocessor import Preprocessor, PreprocessorParameter +from scripts.utils import resize_image_with_pad + + +class PreprocessorTEED(Preprocessor): + def __init__(self): + super().__init__(name="softedge_teed") + self.tags = ["SoftEdge"] + self.slider_1 = PreprocessorParameter( + label="Safe Steps", + minimum=0, + maximum=10, + value=2, + step=1, + ) + self.model = None + + def __call__( + self, + input_image, + resolution, + slider_1=None, + slider_2=None, + slider_3=None, + **kwargs + ): + img, remove_pad = resize_image_with_pad(input_image, resolution) + if self.model is None: + self.model = TEEDDetector() + + result = self.model(img, safe_steps=int(slider_1)) + return remove_pad(result) + + +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 PreprocessorAnyline(Preprocessor): + def __init__(self): + super().__init__(name="softedge_anyline") + self.tags = ["SoftEdge"] + self.slider_resolution = PreprocessorParameter( + label="Resolution", + minimum=64, + maximum=2048, + value=1280, + step=8, + visible=True, + ) + self.slider_1 = PreprocessorParameter( + label="Safe Steps", + minimum=0, + maximum=10, + value=2, + step=1, + ) + self.preprocessor_deps = ["lineart_standard"] + self.model = None + + def __call__( + self, + input_image, + resolution, + slider_1=None, + slider_2=None, + slider_3=None, + **kwargs + ): + img, remove_pad = resize_image_with_pad(input_image, resolution) + if self.model is None: + self.model = TEEDDetector(mteed=True) + + mteed_result = self.model(img, safe_steps=int(slider_1)) + mteed_result = HWC3(mteed_result) + lineart_preprocessor = Preprocessor.get_preprocessor("lineart_standard") + assert lineart_preprocessor is not None + lineart_result = lineart_preprocessor(img, resolution) + lineart_result = get_intensity_mask( + lineart_result, lower_bound=0, upper_bound=1 + ) + cleaned = morphology.remove_small_objects( + lineart_result.astype(bool), min_size=36, connectivity=1 + ) + lineart_result = lineart_result * cleaned + final_result = combine_layers(mteed_result, lineart_result) + return remove_pad(final_result) + + +Preprocessor.add_supported_preprocessor(PreprocessorTEED()) +Preprocessor.add_supported_preprocessor(PreprocessorAnyline())