Spaces:
Runtime error
Runtime error
| import torch | |
| import numpy as np | |
| import comfy.model_management as model_management | |
| import comfy.utils | |
| # Requires comfyui_controlnet_aux funcsions and classes | |
| from ..utils import common_annotator_call, MAX_RESOLUTION | |
| 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 AnyLinePreprocessor: | |
| def INPUT_TYPES(s): | |
| return { | |
| "required": { | |
| "image": ("IMAGE",), | |
| "merge_with_lineart": (["lineart_standard", "lineart_realisitic", "lineart_anime", "manga_line"], {"default": "lineart_standard"}), | |
| "resolution": ("INT", {"default": 1280, "min": 512, "max": MAX_RESOLUTION, "step": 8}) | |
| }, | |
| "optional": { | |
| "lineart_lower_bound": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.01}), | |
| "lineart_upper_bound": ("FLOAT", {"default": 1, "min": 0, "max": 1, "step": 0.01}), | |
| "object_min_size": ("INT", {"default": 36, "min": 1, "max": MAX_RESOLUTION}), | |
| "object_connectivity": ("INT", {"default": 1, "min": 1, "max": MAX_RESOLUTION}), | |
| } | |
| } | |
| RETURN_TYPES = ("IMAGE",) | |
| RETURN_NAMES = ("image",) | |
| FUNCTION = "get_anyline" | |
| CATEGORY = "ControlNet Preprocessors/Line Extractors" | |
| def __init__(self): | |
| self.device = model_management.get_torch_device() | |
| def get_anyline(self, image, merge_with_lineart, resolution, lineart_lower_bound=0, lineart_upper_bound=1, object_min_size=36, object_connectivity=1): | |
| from controlnet_aux.teed import TEDDetector | |
| from skimage import morphology | |
| pbar = comfy.utils.ProgressBar(3) | |
| # Process the image with MTEED model | |
| mteed_model = TEDDetector.from_pretrained("TheMistoAI/MistoLine", "MTEED.pth", subfolder="Anyline").to(self.device) | |
| mteed_result = common_annotator_call(mteed_model, image, resolution=resolution, show_pbar=False) | |
| mteed_result = mteed_result.numpy() | |
| del mteed_model | |
| pbar.update(1) | |
| # Process the image with the lineart standard preprocessor | |
| if merge_with_lineart == "lineart_standard": | |
| from controlnet_aux.lineart_standard import LineartStandardDetector | |
| lineart_standard_detector = LineartStandardDetector() | |
| lineart_result = common_annotator_call(lineart_standard_detector, image, guassian_sigma=2, intensity_threshold=3, resolution=resolution, show_pbar=False).numpy() | |
| del lineart_standard_detector | |
| else: | |
| from controlnet_aux.lineart import LineartDetector | |
| from controlnet_aux.lineart_anime import LineartAnimeDetector | |
| from controlnet_aux.manga_line import LineartMangaDetector | |
| lineart_detector = dict(lineart_realisitic=LineartDetector, lineart_anime=LineartAnimeDetector, manga_line=LineartMangaDetector)[merge_with_lineart] | |
| lineart_detector = lineart_detector.from_pretrained().to(self.device) | |
| lineart_result = common_annotator_call(lineart_detector, image, resolution=resolution, show_pbar=False).numpy() | |
| del lineart_detector | |
| pbar.update(1) | |
| final_result = [] | |
| for i in range(len(image)): | |
| _lineart_result = get_intensity_mask(lineart_result[i], lower_bound=lineart_lower_bound, upper_bound=lineart_upper_bound) | |
| _cleaned = morphology.remove_small_objects(_lineart_result.astype(bool), min_size=object_min_size, connectivity=object_connectivity) | |
| _lineart_result = _lineart_result * _cleaned | |
| _mteed_result = mteed_result[i] | |
| # Combine the results | |
| final_result.append(torch.from_numpy(combine_layers(_mteed_result, _lineart_result))) | |
| pbar.update(1) | |
| return (torch.stack(final_result),) | |
| NODE_CLASS_MAPPINGS = { | |
| "AnyLineArtPreprocessor_aux": AnyLinePreprocessor | |
| } | |
| NODE_DISPLAY_NAME_MAPPINGS = { | |
| "AnyLineArtPreprocessor_aux": "AnyLine Lineart" | |
| } | |