Add Anyline Preprocessor (#2904)

annotator/teed/__init__.py

@@ -1,36 +1,52 @@
-"""
-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",
         )
         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
+            return edge

scripts/preprocessor/__init__.py

@@ -1,3 +1,4 @@
+from .teed import *
 from .pulid import *
 from .inpaint import *
 from .lama_inpaint import *

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,

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

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())