# pylint: skip-file """ Copyright 2024 Lvmin Zhang, fannovel16, Mikubill, Benjamin Berman Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ import logging as log from ..controlnet_aux.utils import ResizeMode, safe_numpy import numpy as np import torch import cv2 from ..controlnet_aux.utils import get_unique_axis0 from ..controlnet_aux.lvminthin import nake_nms, lvmin_thin MAX_IMAGEGEN_RESOLUTION = 8192 #https://github.com/comfyanonymous/ComfyUI/blob/c910b4a01ca58b04e5d4ab4c747680b996ada02b/nodes.py#L42 RESIZE_MODES = [ResizeMode.RESIZE.value, ResizeMode.INNER_FIT.value, ResizeMode.OUTER_FIT.value] #Port from https://github.com/Mikubill/sd-webui-controlnet/blob/e67e017731aad05796b9615dc6eadce911298ea1/internal_controlnet/external_code.py#L89 class PixelPerfectResolution: @classmethod def INPUT_TYPES(s): return { "required": { "original_image": ("IMAGE", ), "image_gen_width": ("INT", {"default": 512, "min": 64, "max": MAX_IMAGEGEN_RESOLUTION, "step": 8}), "image_gen_height": ("INT", {"default": 512, "min": 64, "max": MAX_IMAGEGEN_RESOLUTION, "step": 8}), #https://github.com/comfyanonymous/ComfyUI/blob/c910b4a01ca58b04e5d4ab4c747680b996ada02b/nodes.py#L854 "resize_mode": (RESIZE_MODES, {"default": ResizeMode.RESIZE.value}) } } RETURN_TYPES = ("INT",) RETURN_NAMES = ("RESOLUTION (INT)", ) FUNCTION = "execute" CATEGORY = "ControlNet Preprocessors" def execute(self, original_image, image_gen_width, image_gen_height, resize_mode): _, raw_H, raw_W, _ = original_image.shape k0 = float(image_gen_height) / float(raw_H) k1 = float(image_gen_width) / float(raw_W) if resize_mode == ResizeMode.OUTER_FIT.value: estimation = min(k0, k1) * float(min(raw_H, raw_W)) else: estimation = max(k0, k1) * float(min(raw_H, raw_W)) log.debug(f"Pixel Perfect Computation:") log.debug(f"resize_mode = {resize_mode}") log.debug(f"raw_H = {raw_H}") log.debug(f"raw_W = {raw_W}") log.debug(f"target_H = {image_gen_height}") log.debug(f"target_W = {image_gen_width}") log.debug(f"estimation = {estimation}") return (int(np.round(estimation)), ) class HintImageEnchance: @classmethod def INPUT_TYPES(s): return { "required": { "hint_image": ("IMAGE", ), "image_gen_width": ("INT", {"default": 512, "min": 64, "max": MAX_IMAGEGEN_RESOLUTION, "step": 8}), "image_gen_height": ("INT", {"default": 512, "min": 64, "max": MAX_IMAGEGEN_RESOLUTION, "step": 8}), #https://github.com/comfyanonymous/ComfyUI/blob/c910b4a01ca58b04e5d4ab4c747680b996ada02b/nodes.py#L854 "resize_mode": (RESIZE_MODES, {"default": ResizeMode.RESIZE.value}) } } RETURN_TYPES = ("IMAGE",) FUNCTION = "execute" CATEGORY = "ControlNet Preprocessors" def execute(self, hint_image, image_gen_width, image_gen_height, resize_mode): outs = [] for single_hint_image in hint_image: np_hint_image = np.asarray(single_hint_image * 255., dtype=np.uint8) if resize_mode == ResizeMode.RESIZE.value: np_hint_image = self.execute_resize(np_hint_image, image_gen_width, image_gen_height) elif resize_mode == ResizeMode.OUTER_FIT.value: np_hint_image = self.execute_outer_fit(np_hint_image, image_gen_width, image_gen_height) else: np_hint_image = self.execute_inner_fit(np_hint_image, image_gen_width, image_gen_height) outs.append(torch.from_numpy(np_hint_image.astype(np.float32) / 255.0)) return (torch.stack(outs, dim=0),) def execute_resize(self, detected_map, w, h): detected_map = self.high_quality_resize(detected_map, (w, h)) detected_map = safe_numpy(detected_map) return detected_map def execute_outer_fit(self, detected_map, w, h): old_h, old_w, _ = detected_map.shape old_w = float(old_w) old_h = float(old_h) k0 = float(h) / old_h k1 = float(w) / old_w safeint = lambda x: int(np.round(x)) k = min(k0, k1) borders = np.concatenate([detected_map[0, :, :], detected_map[-1, :, :], detected_map[:, 0, :], detected_map[:, -1, :]], axis=0) high_quality_border_color = np.median(borders, axis=0).astype(detected_map.dtype) if len(high_quality_border_color) == 4: # Inpaint hijack high_quality_border_color[3] = 255 high_quality_background = np.tile(high_quality_border_color[None, None], [h, w, 1]) detected_map = self.high_quality_resize(detected_map, (safeint(old_w * k), safeint(old_h * k))) new_h, new_w, _ = detected_map.shape pad_h = max(0, (h - new_h) // 2) pad_w = max(0, (w - new_w) // 2) high_quality_background[pad_h:pad_h + new_h, pad_w:pad_w + new_w] = detected_map detected_map = high_quality_background detected_map = safe_numpy(detected_map) return detected_map def execute_inner_fit(self, detected_map, w, h): old_h, old_w, _ = detected_map.shape old_w = float(old_w) old_h = float(old_h) k0 = float(h) / old_h k1 = float(w) / old_w safeint = lambda x: int(np.round(x)) k = max(k0, k1) detected_map = self.high_quality_resize(detected_map, (safeint(old_w * k), safeint(old_h * k))) new_h, new_w, _ = detected_map.shape pad_h = max(0, (new_h - h) // 2) pad_w = max(0, (new_w - w) // 2) detected_map = detected_map[pad_h:pad_h+h, pad_w:pad_w+w] detected_map = safe_numpy(detected_map) return detected_map def high_quality_resize(self, x, size): # Written by lvmin # Super high-quality control map up-scaling, considering binary, seg, and one-pixel edges inpaint_mask = None if x.ndim == 3 and x.shape[2] == 4: inpaint_mask = x[:, :, 3] x = x[:, :, 0:3] if x.shape[0] != size[1] or x.shape[1] != size[0]: new_size_is_smaller = (size[0] * size[1]) < (x.shape[0] * x.shape[1]) new_size_is_bigger = (size[0] * size[1]) > (x.shape[0] * x.shape[1]) unique_color_count = len(get_unique_axis0(x.reshape(-1, x.shape[2]))) is_one_pixel_edge = False is_binary = False if unique_color_count == 2: is_binary = np.min(x) < 16 and np.max(x) > 240 if is_binary: xc = x xc = cv2.erode(xc, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1) xc = cv2.dilate(xc, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1) one_pixel_edge_count = np.where(xc < x)[0].shape[0] all_edge_count = np.where(x > 127)[0].shape[0] is_one_pixel_edge = one_pixel_edge_count * 2 > all_edge_count if 2 < unique_color_count < 200: interpolation = cv2.INTER_NEAREST elif new_size_is_smaller: interpolation = cv2.INTER_AREA else: interpolation = cv2.INTER_CUBIC # Must be CUBIC because we now use nms. NEVER CHANGE THIS y = cv2.resize(x, size, interpolation=interpolation) if inpaint_mask is not None: inpaint_mask = cv2.resize(inpaint_mask, size, interpolation=interpolation) if is_binary: y = np.mean(y.astype(np.float32), axis=2).clip(0, 255).astype(np.uint8) if is_one_pixel_edge: y = nake_nms(y) _, y = cv2.threshold(y, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) y = lvmin_thin(y, prunings=new_size_is_bigger) else: _, y = cv2.threshold(y, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) y = np.stack([y] * 3, axis=2) else: y = x if inpaint_mask is not None: inpaint_mask = (inpaint_mask > 127).astype(np.float32) * 255.0 inpaint_mask = inpaint_mask[:, :, None].clip(0, 255).astype(np.uint8) y = np.concatenate([y, inpaint_mask], axis=2) return y class ImageGenResolutionFromLatent: @classmethod def INPUT_TYPES(s): return { "required": { "latent": ("LATENT", ) } } RETURN_TYPES = ("INT", "INT") RETURN_NAMES = ("IMAGE_GEN_WIDTH (INT)", "IMAGE_GEN_HEIGHT (INT)") FUNCTION = "execute" CATEGORY = "ControlNet Preprocessors" def execute(self, latent): _, _, H, W = latent["samples"].shape return (W * 8, H * 8) class ImageGenResolutionFromImage: @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE", ) } } RETURN_TYPES = ("INT", "INT") RETURN_NAMES = ("IMAGE_GEN_WIDTH (INT)", "IMAGE_GEN_HEIGHT (INT)") FUNCTION = "execute" CATEGORY = "ControlNet Preprocessors" def execute(self, image): _, H, W, _ = image.shape return (W, H) NODE_CLASS_MAPPINGS = { "PixelPerfectResolution": PixelPerfectResolution, "ImageGenResolutionFromImage": ImageGenResolutionFromImage, "ImageGenResolutionFromLatent": ImageGenResolutionFromLatent, "HintImageEnchance": HintImageEnchance } NODE_DISPLAY_NAME_MAPPINGS = { "PixelPerfectResolution": "Pixel Perfect Resolution", "ImageGenResolutionFromImage": "Generation Resolution From Image", "ImageGenResolutionFromLatent": "Generation Resolution From Latent", "HintImageEnchance": "Enhance And Resize Hint Images" }