feat: SEGSUpscaler

fix: Make Tile SEGS - Ensuring the overlap of an irregular mask.

https://github.com/ltdrdata/ComfyUI-Impact-Pack/issues/497
https://github.com/ltdrdata/ComfyUI-Inspire-Pack/issues/79

README.md

@@ -108,6 +108,8 @@ This custom node helps to conveniently enhance images through Detector, Detailer
     * From SEG_ELT - Extract detailed information from SEG_ELT.
     * Edit SEG_ELT - Modify some of the information in SEG_ELT.
     * Dilate SEG_ELT - Dilate the mask of SEG_ELT.
+    * From SEG_ELT bbox - Extract coordinate from bbox in SEG_ELT
+    * From SEG_ELT crop_region - Extract coordinate from crop_region in SEG_ELT
 
 * Mask Manipulation
   * Dilate Mask - Dilate Mask. 

__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
 
@@ -229,6 +230,8 @@ NODE_CLASS_MAPPINGS = {
     "ImpactDilateMaskInSEGS": DilateMaskInSEGS,
     "ImpactGaussianBlurMaskInSEGS": GaussianBlurMaskInSEGS,
     "ImpactScaleBy_BBOX_SEG_ELT": SEG_ELT_BBOX_ScaleBy,
+    "ImpactFrom_SEG_ELT_bbox": From_SEG_ELT_bbox,
+    "ImpactFrom_SEG_ELT_crop_region": From_SEG_ELT_crop_region,
 
     "BboxDetectorCombined_v2": BboxDetectorCombined,
     "SegmDetectorCombined_v2": SegmDetectorCombined,
@@ -256,6 +259,8 @@ NODE_CLASS_MAPPINGS = {
     "ImpactWildcardProcessor": ImpactWildcardProcessor,
     "ImpactWildcardEncode": ImpactWildcardEncode,
 
+    "SEGSUpscaler": SEGSUpscaler,
+    "SEGSUpscalerPipe": SEGSUpscalerPipe,
     "SEGSDetailer": SEGSDetailer,
     "SEGSPaste": SEGSPaste,
     "SEGSPreview": SEGSPreview,
@@ -361,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)",
@@ -402,6 +409,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ImpactAssembleSEGS": "Assemble (SEGS)",
     "ImpactFrom_SEG_ELT": "From SEG_ELT",
     "ImpactEdit_SEG_ELT": "Edit SEG_ELT",
+    "ImpactFrom_SEG_ELT_bbox": "From SEG_ELT bbox",
+    "ImpactFrom_SEG_ELT_crop_region": "From SEG_ELT crop_region",
     "ImpactDilate_Mask_SEG_ELT": "Dilate Mask (SEG_ELT)",
     "ImpactScaleBy_BBOX_SEG_ELT": "ScaleBy BBOX (SEG_ELT)",
     "ImpactDilateMask": "Dilate Mask",

modules/impact/config.py

@@ -2,7 +2,7 @@ import configparser
 import os
 
 
-version_code = [4, 80]
+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

modules/impact/core.py

@@ -1609,6 +1609,12 @@ class ControlNetAdvancedWrapper:
         else:
             self.control_image = None
 
+    def doit_ipadapter(self, model):
+        if self.prev_control_net is not None:
+            return self.prev_control_net.doit_ipadapter(model)
+        else:
+            return model, []
+
     def apply(self, positive, negative, image, mask=None, use_acn=False):
         cnet_image_list = []
         prev_cnet_images = []
@@ -1811,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

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)

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
@@ -755,6 +756,44 @@ class From_SEG_ELT:
         return (seg_elt, cropped_image, to_tensor(seg_elt.cropped_mask), seg_elt.crop_region, seg_elt.bbox, seg_elt.control_net_wrapper, seg_elt.confidence, seg_elt.label,)
 
 
+class From_SEG_ELT_bbox:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                     "bbox": ("SEG_ELT_bbox", ),
+                     },
+                }
+
+    RETURN_TYPES = ("INT", "INT", "INT", "INT")
+    RETURN_NAMES = ("left", "top", "right", "bottom")
+
+    FUNCTION = "doit"
+
+    CATEGORY = "ImpactPack/Util"
+
+    def doit(self, bbox):
+        return bbox
+
+
+class From_SEG_ELT_crop_region:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                     "crop_region": ("SEG_ELT_crop_region", ),
+                     },
+                }
+
+    RETURN_TYPES = ("INT", "INT", "INT", "INT")
+    RETURN_NAMES = ("left", "top", "right", "bottom")
+
+    FUNCTION = "doit"
+
+    CATEGORY = "ImpactPack/Util"
+
+    def doit(self, crop_region):
+        return crop_region
+
+
 class Edit_SEG_ELT:
     @classmethod
     def INPUT_TYPES(s):
@@ -1441,7 +1480,7 @@ class MakeTileSEGS:
 
     def doit(self, images, bbox_size, crop_factor, min_overlap, filter_segs_dilation, mask_irregularity=0, irregular_mask_mode="Reuse fast", filter_in_segs_opt=None, filter_out_segs_opt=None):
         if bbox_size <= 2*min_overlap:
-            new_min_overlap = 2 / bbox_size
+            new_min_overlap = bbox_size / 2
             print(f"[MakeTileSEGS] min_overlap should be greater than bbox_size. (value changed: {min_overlap} => {new_min_overlap})")
             min_overlap = new_min_overlap
 
@@ -1461,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)
@@ -1495,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:
@@ -1590,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)

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