diff --git a/__init__.py b/__init__.py index b11a58a..83717de 100644 --- a/__init__.py +++ b/__init__.py @@ -96,17 +96,18 @@ def setup_js(): setup_js() -from impact.impact_pack import * -from impact.detectors import * -from impact.pipe import * -from impact.logics import * -from impact.util_nodes import * -from impact.segs_nodes import * -from impact.special_samplers import * -from impact.hf_nodes import * -from impact.bridge_nodes import * -from impact.hook_nodes import * -from impact.animatediff_nodes import * +from .modules.impact.impact_pack import * +from .modules.impact.detectors import * +from .modules.impact.pipe import * +from .modules.impact.logics import * +from .modules.impact.util_nodes import * +from .modules.impact.segs_nodes import * +from .modules.impact.special_samplers import * +from .modules.impact.hf_nodes import * +from .modules.impact.bridge_nodes import * +from .modules.impact.hook_nodes import * +from .modules.impact.animatediff_nodes import * +from .modules.impact.segs_upscaler import * import threading @@ -258,6 +259,8 @@ NODE_CLASS_MAPPINGS = { "ImpactWildcardProcessor": ImpactWildcardProcessor, "ImpactWildcardEncode": ImpactWildcardEncode, + "SEGSUpscaler": SEGSUpscaler, + "SEGSUpscalerPipe": SEGSUpscalerPipe, "SEGSDetailer": SEGSDetailer, "SEGSPaste": SEGSPaste, "SEGSPreview": SEGSPreview, @@ -363,6 +366,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "DetailerForEachDebugPipe": "DetailerDebug (SEGS/pipe)", "SEGSDetailerForAnimateDiff": "SEGSDetailer For AnimateDiff (SEGS/pipe)", "DetailerForEachPipeForAnimateDiff": "Detailer For AnimateDiff (SEGS/pipe)", + "SEGSUpscaler": "Upscaler (SEGS)", + "SEGSUpscalerPipe": "Upscaler (SEGS/pipe)", "SAMDetectorCombined": "SAMDetector (combined)", "SAMDetectorSegmented": "SAMDetector (segmented)", diff --git a/modules/impact/config.py b/modules/impact/config.py index c25269c..3de4a5e 100644 --- a/modules/impact/config.py +++ b/modules/impact/config.py @@ -2,7 +2,7 @@ import configparser import os -version_code = [4, 81] +version_code = [4, 82] version = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '') dependency_version = 20 diff --git a/modules/impact/core.py b/modules/impact/core.py index 447061e..c59323d 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -1817,7 +1817,7 @@ def random_mask_raw(mask, bbox, factor): w = x2 - x1 h = y2 - y1 - factor = int(min(w, h) * factor / 4) + factor = max(6, int(min(w, h) * factor / 4)) def draw_random_circle(center, radius): i, j = center diff --git a/modules/impact/hook_nodes.py b/modules/impact/hook_nodes.py index c218da4..4e4707f 100644 --- a/modules/impact/hook_nodes.py +++ b/modules/impact/hook_nodes.py @@ -73,11 +73,11 @@ class PreviewDetailerHookProvider: "hidden": {"unique_id": "UNIQUE_ID"}, } - RETURN_TYPES = ("DETAILER_HOOK", ) + RETURN_TYPES = ("DETAILER_HOOK", "UPSCALER_HOOK") FUNCTION = "doit" CATEGORY = "ImpactPack/Util" def doit(self, quality, unique_id): hook = hooks.PreviewDetailerHook(unique_id, quality) - return (hook, ) + return (hook, hook) diff --git a/modules/impact/segs_nodes.py b/modules/impact/segs_nodes.py index 20e3da9..d3fd95d 100644 --- a/modules/impact/segs_nodes.py +++ b/modules/impact/segs_nodes.py @@ -5,11 +5,12 @@ import impact.impact_server from nodes import MAX_RESOLUTION from impact.utils import * -import impact.core as core -from impact.core import SEG +from . import core +from .core import SEG import impact.utils as utils from . import defs - +from . import segs_upscaler +import math class SEGSDetailer: @classmethod @@ -1499,6 +1500,12 @@ class MakeTileSEGS: elif irregular_mask_mode == "All random fast": mask_quality = 512 + # compensate overlap/bbox_size for irregular mask + if mask_irregularity > 0: + compensate = max(6, int(mask_quality * mask_irregularity / 4)) + min_overlap += compensate + bbox_size += compensate*2 + # create exclusion mask if filter_out_segs_opt is not None: exclusion_mask = core.segs_to_combined_mask(filter_out_segs_opt) @@ -1533,8 +1540,8 @@ class MakeTileSEGS: print(f"[MaskTileSEGS] bbox_size is greater than resolution (value changed: {bbox_size} => {new_bbox_size}") bbox_size = new_bbox_size - n_horizontal = int(w / (bbox_size - min_overlap)) - n_vertical = int(h / (bbox_size - min_overlap)) + n_horizontal = math.ceil(w / (bbox_size - min_overlap)) + n_vertical = math.ceil(h / (bbox_size - min_overlap)) w_overlap_sum = (bbox_size * n_horizontal) - w if w_overlap_sum < 0: @@ -1628,3 +1635,137 @@ class MakeTileSEGS: res = (ih, iw), new_segs # segs return (res,) + + +class SEGSUpscaler: + @classmethod + def INPUT_TYPES(s): + resampling_methods = ["lanczos", "nearest", "bilinear", "bicubic"] + + return {"required": { + "image": ("IMAGE",), + "segs": ("SEGS",), + "model": ("MODEL",), + "clip": ("CLIP",), + "vae": ("VAE",), + "rescale_factor": ("FLOAT", {"default": 2, "min": 0.01, "max": 100.0, "step": 0.01}), + "resampling_method": (resampling_methods,), + "supersample": (["true", "false"],), + "rounding_modulus": ("INT", {"default": 8, "min": 8, "max": 1024, "step": 8}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "denoise": ("FLOAT", {"default": 0.5, "min": 0.0001, "max": 1.0, "step": 0.01}), + "feather": ("INT", {"default": 5, "min": 0, "max": 100, "step": 1}), + "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + "noise_mask": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), + }, + "optional": { + "upscale_model_opt": ("UPSCALE_MODEL",), + "upscaler_hook_opt": ("UPSCALER_HOOK",), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Upscale" + + @staticmethod + def doit(image, segs, model, clip, vae, rescale_factor, resampling_method, supersample, rounding_modulus, + seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, feather, inpaint_model, noise_mask, noise_mask_feather, + upscale_model_opt=None, upscaler_hook_opt=None): + + new_image = segs_upscaler.upscaler(image, upscale_model_opt, rescale_factor, resampling_method, supersample, rounding_modulus) + + segs = core.segs_scale_match(segs, new_image.shape) + + ordered_segs = segs[1] + + for i, seg in enumerate(ordered_segs): + cropped_image = crop_ndarray4(new_image.numpy(), seg.crop_region) + cropped_image = to_tensor(cropped_image) + mask = to_tensor(seg.cropped_mask) + mask = tensor_gaussian_blur_mask(mask, feather) + + is_mask_all_zeros = (seg.cropped_mask == 0).all().item() + if is_mask_all_zeros: + print(f"SEGSUpscaler: segment skip [empty mask]") + continue + + if noise_mask: + cropped_mask = seg.cropped_mask + else: + cropped_mask = None + + seg_seed = seed + i + + enhanced_image = segs_upscaler.img2img_segs(cropped_image, model, clip, vae, seg_seed, steps, cfg, sampler_name, scheduler, + positive, negative, denoise, + noise_mask=cropped_mask, control_net_wrapper=seg.control_net_wrapper, + inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather) + if not (enhanced_image is None): + new_image = new_image.cpu() + enhanced_image = enhanced_image.cpu() + left = seg.crop_region[0] + top = seg.crop_region[1] + tensor_paste(new_image, enhanced_image, (left, top), mask) + + if upscaler_hook_opt is not None: + upscaler_hook_opt.post_paste(new_image) + + enhanced_img = tensor_convert_rgb(new_image) + + return (enhanced_img,) + + +class SEGSUpscalerPipe: + @classmethod + def INPUT_TYPES(s): + resampling_methods = ["lanczos", "nearest", "bilinear", "bicubic"] + + return {"required": { + "image": ("IMAGE",), + "segs": ("SEGS",), + "basic_pipe": ("BASIC_PIPE",), + "rescale_factor": ("FLOAT", {"default": 2, "min": 0.01, "max": 100.0, "step": 0.01}), + "resampling_method": (resampling_methods,), + "supersample": (["true", "false"],), + "rounding_modulus": ("INT", {"default": 8, "min": 8, "max": 1024, "step": 8}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), + "denoise": ("FLOAT", {"default": 0.5, "min": 0.0001, "max": 1.0, "step": 0.01}), + "feather": ("INT", {"default": 5, "min": 0, "max": 100, "step": 1}), + "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + "noise_mask": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), + }, + "optional": { + "upscale_model_opt": ("UPSCALE_MODEL",), + "upscaler_hook_opt": ("UPSCALER_HOOK",), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Upscale" + + @staticmethod + def doit(image, segs, basic_pipe, rescale_factor, resampling_method, supersample, rounding_modulus, + seed, steps, cfg, sampler_name, scheduler, denoise, feather, inpaint_model, noise_mask, noise_mask_feather, + upscale_model_opt=None, upscaler_hook_opt=None): + + model, clip, vae, positive, negative = basic_pipe + + return SEGSUpscaler.doit(image, segs, model, clip, vae, rescale_factor, resampling_method, supersample, rounding_modulus, + seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, feather, inpaint_model, noise_mask, noise_mask_feather, + upscale_model_opt=upscale_model_opt, upscaler_hook_opt=upscaler_hook_opt) diff --git a/modules/impact/segs_upscaler.py b/modules/impact/segs_upscaler.py new file mode 100644 index 0000000..165ce77 --- /dev/null +++ b/modules/impact/segs_upscaler.py @@ -0,0 +1,111 @@ +from impact.utils import * +from impact import impact_sampling +from comfy_extras.chainner_models import model_loading +from comfy import model_management +import nodes + + +# Implementation based on `https://github.com/lingondricka2/Upscaler-Detailer` + +# code from comfyroll ---> +# https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/blob/main/nodes/functions_upscale.py + +def upscale_with_model(upscale_model, image): + device = model_management.get_torch_device() + upscale_model.to(device) + in_img = image.movedim(-1,-3).to(device) + free_memory = model_management.get_free_memory(device) + + tile = 512 + overlap = 32 + + oom = True + while oom: + try: + steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap) + pbar = comfy.utils.ProgressBar(steps) + s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar) + oom = False + except model_management.OOM_EXCEPTION as e: + tile //= 2 + if tile < 128: + raise e + + upscale_model.cpu() + s = torch.clamp(s.movedim(-3, -1), min=0, max=1.0) + return s + + +def apply_resize_image(image: Image.Image, original_width, original_height, rounding_modulus, mode='scale', supersample='true', factor: int = 2, width: int = 1024, height: int = 1024, + resample='bicubic'): + # Calculate the new width and height based on the given mode and parameters + if mode == 'rescale': + new_width, new_height = int(original_width * factor), int(original_height * factor) + else: + m = rounding_modulus + original_ratio = original_height / original_width + height = int(width * original_ratio) + + new_width = width if width % m == 0 else width + (m - width % m) + new_height = height if height % m == 0 else height + (m - height % m) + + # Define a dictionary of resampling filters + resample_filters = {'nearest': 0, 'bilinear': 2, 'bicubic': 3, 'lanczos': 1} + + # Apply supersample + if supersample == 'true': + image = image.resize((new_width * 8, new_height * 8), resample=Image.Resampling(resample_filters[resample])) + + # Resize the image using the given resampling filter + resized_image = image.resize((new_width, new_height), resample=Image.Resampling(resample_filters[resample])) + + return resized_image + + +def upscaler(image, upscale_model, rescale_factor, resampling_method, supersample, rounding_modulus): + if upscale_model is not None: + up_image = upscale_with_model(upscale_model, image) + else: + up_image = image + + pil_img = tensor2pil(image) + original_width, original_height = pil_img.size + scaled_image = pil2tensor(apply_resize_image(tensor2pil(up_image), original_width, original_height, rounding_modulus, 'rescale', + supersample, rescale_factor, 1024, resampling_method)) + return scaled_image + +# <--- + + +def img2img_segs(image, model, clip, vae, seed, steps, cfg, sampler_name, scheduler, + positive, negative, denoise, noise_mask, control_net_wrapper=None, + inpaint_model=False, noise_mask_feather=0): + if noise_mask is not None: + noise_mask = tensor_gaussian_blur_mask(noise_mask, noise_mask_feather) + noise_mask = noise_mask.squeeze(3) + + if control_net_wrapper is not None: + positive, negative, _ = control_net_wrapper.apply(positive, negative, image, noise_mask) + + # prepare mask + if noise_mask is not None and inpaint_model: + positive, negative, latent_image = nodes.InpaintModelConditioning().encode(positive, negative, image, vae, noise_mask) + else: + latent_image = to_latent_image(image, vae) + if noise_mask is not None: + latent_image['noise_mask'] = noise_mask + + refined_latent = latent_image + + # ksampler + refined_latent = impact_sampling.ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, refined_latent, denoise) + + # non-latent downscale - latent downscale cause bad quality + refined_image = vae.decode(refined_latent['samples']) + + # prevent mixing of device + refined_image = refined_image.cpu() + + # don't convert to latent - latent break image + # preserving pil is much better + return refined_image