Add Anyline Preprocessor (#2904)
Chenlei Hu
GitHub
parent
681de2c4f0
commit
69d80e5a57
|
|
@ -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 @@ class TEEDDector:
|
|||
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 @@ class TEEDDector:
|
|||
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
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from .teed import *
|
||||
from .pulid import *
|
||||
from .inpaint import *
|
||||
from .lama_inpaint import *
|
||||
|
|
|
|||
|
|
@ -1141,33 +1141,33 @@ legacy_preprocessors = {
|
|||
"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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
Loading…
Reference in New Issue