diff --git a/modules/impact/animatediff_nodes.py b/modules/impact/animatediff_nodes.py index 57d8d0c..c26228b 100644 --- a/modules/impact/animatediff_nodes.py +++ b/modules/impact/animatediff_nodes.py @@ -26,7 +26,7 @@ class SEGSDetailerForAnimateDiff: "optional": { "refiner_basic_pipe_opt": ("BASIC_PIPE",), # TODO: "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), - # TODO: "noise_mask_feather": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + # TODO: "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), } } @@ -116,7 +116,7 @@ class DetailerForEachPipeForAnimateDiff: "detailer_hook": ("DETAILER_HOOK",), "refiner_basic_pipe_opt": ("BASIC_PIPE",), # "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), - # "noise_mask_feather": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + # "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), } } diff --git a/modules/impact/config.py b/modules/impact/config.py index fc40699..e54110e 100644 --- a/modules/impact/config.py +++ b/modules/impact/config.py @@ -2,7 +2,7 @@ import configparser import os -version_code = [4, 69] +version_code = [4, 70] 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 fa0b2a0..023a3d2 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -1,15 +1,9 @@ -import copy -import os - -import numpy -import torch from segment_anything import SamPredictor -import torch.nn.functional as F from impact.utils import * from collections import namedtuple import numpy as np -from skimage.measure import label, regionprops +from skimage.measure import label import nodes import comfy_extras.nodes_upscale_model as model_upscale @@ -21,7 +15,9 @@ import cv2 import time from comfy import model_management from impact import utils -from scipy.ndimage import distance_transform_edt +from impact import impact_sampling +from concurrent.futures import ThreadPoolExecutor + SEG = namedtuple("SEG", ['cropped_image', 'cropped_mask', 'confidence', 'crop_region', 'bbox', 'label', 'control_net_wrapper'], @@ -69,44 +65,6 @@ def erosion_mask(mask, grow_mask_by): return mask_erosion[:, :, :w, :h].round().cpu() -def ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, - refiner_ratio=None, refiner_model=None, refiner_clip=None, refiner_positive=None, - refiner_negative=None): - if refiner_ratio is None or refiner_model is None or refiner_clip is None or refiner_positive is None or refiner_negative is None: - refined_latent = \ - nodes.KSampler().sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, - denoise)[0] - else: - advanced_steps = math.floor(steps / denoise) - start_at_step = advanced_steps - steps - end_at_step = start_at_step + math.floor(steps * (1.0 - refiner_ratio)) - - print(f"pre: {start_at_step} .. {end_at_step} / {advanced_steps}") - temp_latent = \ - nodes.KSamplerAdvanced().sample(model, "enable", seed, advanced_steps, cfg, sampler_name, scheduler, - positive, negative, latent_image, start_at_step, end_at_step, - "enable")[0] - - if 'noise_mask' in latent_image: - # noise_latent = \ - # nodes.KSamplerAdvanced().sample(refiner_model, "enable", seed, advanced_steps, cfg, sampler_name, - # scheduler, refiner_positive, refiner_negative, latent_image, end_at_step, - # end_at_step, "enable")[0] - - latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() - temp_latent = \ - latent_compositor.composite(latent_image, temp_latent, 0, 0, False, latent_image['noise_mask'])[0] - - print(f"post: {end_at_step} .. {advanced_steps + 1} / {advanced_steps}") - refined_latent = \ - nodes.KSamplerAdvanced().sample(refiner_model, "disable", seed, advanced_steps, cfg, sampler_name, scheduler, - refiner_positive, refiner_negative, temp_latent, end_at_step, - advanced_steps + 1, - "disable")[0] - - return refined_latent - - class REGIONAL_PROMPT: def __init__(self, mask, sampler): mask = make_2d_mask(mask) @@ -288,9 +246,8 @@ def enhance_detail(image, model, clip, vae, guide_size, guide_size_for_bbox, max model2, seed2, steps2, cfg2, sampler_name2, scheduler2, positive2, negative2, upscaled_latent2, denoise2 = \ model, seed + i, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise - refined_latent = ksampler_wrapper(model2, seed2, steps2, cfg2, sampler_name2, scheduler2, positive2, negative2, - refined_latent, denoise2, - refiner_ratio, refiner_model, refiner_clip, refiner_positive, refiner_negative) + refined_latent = impact_sampling.ksampler_wrapper(model2, seed2, steps2, cfg2, sampler_name2, scheduler2, positive2, negative2, + refined_latent, denoise2, refiner_ratio, refiner_model, refiner_clip, refiner_positive, refiner_negative) if detailer_hook is not None: refined_latent = detailer_hook.pre_decode(refined_latent) @@ -369,7 +326,7 @@ def enhance_detail_for_animatediff(image_frames, model, clip, vae, guide_size, g print(f"Detailer: segment upscale for ({bbox_w, bbox_h}) | crop region {w, h} x {upscale} -> {new_w, new_h}") # upscale the mask tensor by a factor of 2 using bilinear interpolation - if isinstance(noise_mask, numpy.ndarray): + if isinstance(noise_mask, np.ndarray): noise_mask = torch.from_numpy(noise_mask) if len(noise_mask.shape) == 2: @@ -424,9 +381,8 @@ def enhance_detail_for_animatediff(image_frames, model, clip, vae, guide_size, g if detailer_hook is not None: latent = detailer_hook.post_encode(latent) - refined_latent = ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, - latent, denoise, - refiner_ratio, refiner_model, refiner_clip, refiner_positive, refiner_negative) + refined_latent = impact_sampling.ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, + latent, denoise, refiner_ratio, refiner_model, refiner_clip, refiner_positive, refiner_negative) if detailer_hook is not None: refined_latent = detailer_hook.pre_decode(refined_latent) @@ -474,6 +430,7 @@ def sam_predict(predictor, points, plabs, bbox, threshold): selected = False max_score = 0 + max_mask = None for idx in range(len(scores)): if scores[idx] > max_score: max_score = scores[idx] @@ -485,7 +442,7 @@ def sam_predict(predictor, points, plabs, bbox, threshold): else: pass - if not selected: + if not selected and max_mask is not None: total_masks.append(max_mask) return total_masks @@ -745,7 +702,7 @@ def segs_scale_match(segs, target_shape): new_seg = SEG(cropped_image, cropped_mask, seg.confidence, crop_region, bbox, seg.label, seg.control_net_wrapper) new_segs.append(new_seg) - return ((th, tw), new_segs) + return (th, tw), new_segs # Used Python's slicing feature. stacked_masks[2::3] means starting from index 2, selecting every third tensor with a step size of 3. @@ -1168,7 +1125,7 @@ def segs_to_masklist(segs): masks = [] for seg in segs[1]: - if isinstance(seg.cropped_mask, numpy.ndarray): + if isinstance(seg.cropped_mask, np.ndarray): cropped_mask = torch.from_numpy(seg.cropped_mask) else: cropped_mask = seg.cropped_mask @@ -1217,78 +1174,6 @@ def vae_encode(vae, pixels, use_tile, hook, tile_size=512): return samples -class KSamplerWrapper: - params = None - - def __init__(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise): - self.params = model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise - - def sample(self, latent_image, hook=None): - model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise = self.params - - if hook is not None: - model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise = \ - hook.pre_ksample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, - denoise) - - return nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, - denoise=denoise)[0] - - -class KSamplerAdvancedWrapper: - params = None - - def __init__(self, model, cfg, sampler_name, scheduler, positive, negative): - self.params = model, cfg, sampler_name, scheduler, positive, negative - - def sample_advanced(self, add_noise, seed, steps, latent_image, start_at_step, end_at_step, - return_with_leftover_noise, hook=None, recover_special_sampler=False): - model, cfg, sampler_name, scheduler, positive, negative = self.params - - if hook is not None: - model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent = \ - hook.pre_ksample_advanced(model, add_noise, seed, steps, cfg, sampler_name, scheduler, - positive, negative, latent_image, start_at_step, end_at_step, - return_with_leftover_noise) - - if recover_special_sampler and sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu', 'dpmpp_2m_sde', 'dpmpp_2m_sde_gpu', 'dpmpp_3m_sde', 'dpmpp_3m_sde_gpu']: - base_image = latent_image.copy() - else: - base_image = None - - try: - latent_image = nodes.KSamplerAdvanced().sample(model, add_noise, seed, steps, cfg, sampler_name, scheduler, - positive, negative, latent_image, start_at_step, end_at_step, - return_with_leftover_noise)[0] - except ValueError as e: - if str(e) == 'sigma_min and sigma_max must not be 0': - print(f"\nWARN: sampling skipped - sigma_min and sigma_max are 0") - return latent_image - - if recover_special_sampler and sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu', 'dpmpp_2m_sde', 'dpmpp_2m_sde_gpu', 'dpmpp_3m_sde', 'dpmpp_3m_sde_gpu']: - compensate = 0 if sampler_name in ['uni_pc', 'uni_pc_bh2'] else 2 - sampler_name = 'dpmpp_fast' if sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu'] else 'dpmpp_2m' - latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() - - noise_mask = latent_image['noise_mask'] - - if len(noise_mask.shape) == 4: - noise_mask = noise_mask.squeeze(0).squeeze(0) - - latent_image = \ - latent_compositor.composite(base_image, latent_image, 0, 0, False, noise_mask)[0] - - try: - latent_image = nodes.KSamplerAdvanced().sample(model, add_noise, seed, steps, cfg, sampler_name, scheduler, - positive, negative, latent_image, start_at_step-compensate, end_at_step, - return_with_leftover_noise)[0] - except ValueError as e: - if str(e) == 'sigma_min and sigma_max must not be 0': - print(f"\nWARN: sampling skipped - sigma_min and sigma_max are 0") - - return latent_image - - def latent_upscale_on_pixel_space_shape(samples, scale_method, w, h, vae, use_tile=False, tile_size=512, save_temp_prefix=None, hook=None): pixels = vae_decode(vae, samples, use_tile, hook, tile_size=tile_size) @@ -1323,7 +1208,7 @@ def latent_upscale_on_pixel_space2(samples, scale_method, scale_factor, vae, use def latent_upscale_on_pixel_space(samples, scale_method, scale_factor, vae, use_tile=False, tile_size=512, save_temp_prefix=None, hook=None): - return latent_upscale_on_pixel_space2(samples, scale_method, scale_factor, vae, use_tile, tile_size, save_temp_prefix, hook)[0] + return latent_upscale_on_pixel_space2(samples, scale_method, scale_factor, vae, use_tile, tile_size, save_temp_prefix, hook)[0] def latent_upscale_on_pixel_space_with_model_shape(samples, scale_method, upscale_model, new_w, new_h, vae, @@ -1692,10 +1577,8 @@ class TiledKSamplerWrapper: hook.pre_ksample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise) - return \ - TiledKSampler().sample(model, seed, tile_width, tile_height, tiling_strategy, steps, cfg, sampler_name, - scheduler, - positive, negative, latent_image, denoise)[0] + return TiledKSampler().sample(model, seed, tile_width, tile_height, tiling_strategy, steps, cfg, sampler_name, + scheduler, positive, negative, latent_image, denoise)[0] class PixelTiledKSampleUpscaler: @@ -1722,10 +1605,8 @@ class PixelTiledKSampleUpscaler: scale_method, model, vae, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise = self.params tile_width, tile_height, tiling_strategy = self.tile_params - return \ - TiledKSampler().sample(model, seed, tile_width, tile_height, tiling_strategy, steps, cfg, sampler_name, - scheduler, - positive, negative, latent, denoise)[0] + return TiledKSampler().sample(model, seed, tile_width, tile_height, tiling_strategy, steps, cfg, sampler_name, + scheduler, positive, negative, latent, denoise)[0] def upscale(self, step_info, samples, upscale_factor, save_temp_prefix=None): scale_method, model, vae, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise = self.params @@ -1838,9 +1719,6 @@ def update_node_status(node, text, progress=None): }, PromptServer.instance.client_id) -from concurrent.futures import ThreadPoolExecutor - - def random_mask_raw(mask, bbox, factor): x1, y1, x2, y2 = bbox w = x2 - x1 diff --git a/modules/impact/impact_pack.py b/modules/impact/impact_pack.py index e354e39..7d3c613 100644 --- a/modules/impact/impact_pack.py +++ b/modules/impact/impact_pack.py @@ -176,7 +176,7 @@ class DetailerForEach: "optional": { "detailer_hook": ("DETAILER_HOOK",), "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), - "noise_mask_feather": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), } } @@ -331,7 +331,7 @@ class DetailerForEachPipe: "detailer_hook": ("DETAILER_HOOK",), "refiner_basic_pipe_opt": ("BASIC_PIPE",), "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), - "noise_mask_feather": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), } } @@ -418,7 +418,7 @@ class FaceDetailer: "segm_detector_opt": ("SEGM_DETECTOR", ), "detailer_hook": ("DETAILER_HOOK",), "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), - "noise_mask_feather": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), }} RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "MASK", "DETAILER_PIPE", "IMAGE") @@ -1189,7 +1189,7 @@ class FaceDetailerPipe: }, "optional": { "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), - "noise_mask_feather": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), } } @@ -1279,7 +1279,7 @@ class MaskDetailerPipe: "refiner_basic_pipe_opt": ("BASIC_PIPE", ), "detailer_hook": ("DETAILER_HOOK",), "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), - "noise_mask_feather": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), } } @@ -1288,7 +1288,7 @@ class MaskDetailerPipe: OUTPUT_IS_LIST = (False, True, True, False, False) FUNCTION = "doit" - CATEGORY = "ImpactPack/__for_test" + CATEGORY = "ImpactPack/Detailer" def doit(self, image, mask, basic_pipe, guide_size, guide_size_for, max_size, mask_mode, seed, steps, cfg, sampler_name, scheduler, denoise, diff --git a/modules/impact/impact_sampling.py b/modules/impact/impact_sampling.py new file mode 100644 index 0000000..21e933e --- /dev/null +++ b/modules/impact/impact_sampling.py @@ -0,0 +1,222 @@ +import nodes +from comfy.k_diffusion import sampling as k_diffusion_sampling +from comfy import samplers +from comfy_extras import nodes_custom_sampler + +import torch +import math + + +def calculate_sigmas(model, sampler, scheduler, steps): + discard_penultimate_sigma = False + if sampler in ['dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2']: + steps += 1 + discard_penultimate_sigma = True + + sigmas = samplers.calculate_sigmas_scheduler(model.model, scheduler, steps) + + if discard_penultimate_sigma: + sigmas = torch.cat([sigmas[:-2], sigmas[-1:]]) + return sigmas + + +def get_noise_sampler(x, cpu, total_sigmas, **kwargs): + if 'extra_args' in kwargs and 'seed' in kwargs['extra_args']: + sigma_min, sigma_max = total_sigmas[total_sigmas > 0].min(), total_sigmas.max() + seed = kwargs['extra_args'].get("seed", None) + return k_diffusion_sampling.BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=cpu) + return None + + +def ksampler(sampler_name, total_sigmas, extra_options={}, inpaint_options={}): + if sampler_name == "dpmpp_sde": + def sample_dpmpp_sde(model, x, sigmas, **kwargs): + noise_sampler = get_noise_sampler(x, True, total_sigmas, **kwargs) + if noise_sampler is not None: + kwargs['noise_sampler'] = noise_sampler + + return k_diffusion_sampling.sample_dpmpp_sde(model, x, sigmas, **kwargs) + + sampler_function = sample_dpmpp_sde + + elif sampler_name == "dpmpp_sde_gpu": + def sample_dpmpp_sde(model, x, sigmas, **kwargs): + noise_sampler = get_noise_sampler(x, False, total_sigmas, **kwargs) + if noise_sampler is not None: + kwargs['noise_sampler'] = noise_sampler + + return k_diffusion_sampling.sample_dpmpp_sde_gpu(model, x, sigmas, **kwargs) + + sampler_function = sample_dpmpp_sde + + elif sampler_name == "dpmpp_2m_sde": + def sample_dpmpp_sde(model, x, sigmas, **kwargs): + noise_sampler = get_noise_sampler(x, True, total_sigmas, **kwargs) + if noise_sampler is not None: + kwargs['noise_sampler'] = noise_sampler + + return k_diffusion_sampling.sample_dpmpp_2m_sde(model, x, sigmas, **kwargs) + + sampler_function = sample_dpmpp_sde + + elif sampler_name == "dpmpp_2m_sde_gpu": + def sample_dpmpp_sde(model, x, sigmas, **kwargs): + noise_sampler = get_noise_sampler(x, False, total_sigmas, **kwargs) + if noise_sampler is not None: + kwargs['noise_sampler'] = noise_sampler + + return k_diffusion_sampling.sample_dpmpp_2m_sde_gpu(model, x, sigmas, **kwargs) + + sampler_function = sample_dpmpp_sde + + elif sampler_name == "dpmpp_3m_sde": + def sample_dpmpp_sde(model, x, sigmas, **kwargs): + noise_sampler = get_noise_sampler(x, True, total_sigmas, **kwargs) + if noise_sampler is not None: + kwargs['noise_sampler'] = noise_sampler + + return k_diffusion_sampling.sample_dpmpp_2m_sde(model, x, sigmas, **kwargs) + + sampler_function = sample_dpmpp_sde + + elif sampler_name == "dpmpp_3m_sde_gpu": + def sample_dpmpp_sde(model, x, sigmas, **kwargs): + noise_sampler = get_noise_sampler(x, False, total_sigmas, **kwargs) + if noise_sampler is not None: + kwargs['noise_sampler'] = noise_sampler + + return k_diffusion_sampling.sample_dpmpp_2m_sde_gpu(model, x, sigmas, **kwargs) + + sampler_function = sample_dpmpp_sde + else: + return samplers.ksampler(sampler_name, extra_options, inpaint_options) + + return samplers.KSAMPLER(sampler_function, extra_options, inpaint_options) + + +def separated_sample(model, add_noise, seed, steps, cfg, sampler_name, scheduler, positive, negative, + latent_image, start_at_step, end_at_step, return_with_leftover_noise, sigma_ratio=1.0): + total_sigmas = calculate_sigmas(model, sampler_name, scheduler, steps) + + sigmas = total_sigmas[start_at_step:end_at_step+1] * sigma_ratio + impact_sampler = ksampler(sampler_name, total_sigmas) + + if len(sigmas) == 0 or (len(sigmas) == 1 and sigmas[0] == 0): + return latent_image + + res = nodes_custom_sampler.SamplerCustom().sample(model, add_noise, seed, cfg, positive, negative, impact_sampler, sigmas, latent_image) + + if return_with_leftover_noise: + return res[0] + else: + return res[1] + + +def ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, + refiner_ratio=None, refiner_model=None, refiner_clip=None, refiner_positive=None, refiner_negative=None): + + if refiner_ratio is None or refiner_model is None or refiner_clip is None or refiner_positive is None or refiner_negative is None: + refined_latent = nodes.KSampler().sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise)[0] + else: + advanced_steps = math.floor(steps / denoise) + start_at_step = advanced_steps - steps + end_at_step = start_at_step + math.floor(steps * (1.0 - refiner_ratio)) + + print(f"pre: {start_at_step} .. {end_at_step} / {advanced_steps}") + temp_latent = separated_sample(model, True, seed, advanced_steps, cfg, sampler_name, scheduler, + positive, negative, latent_image, start_at_step, end_at_step, True) + + if 'noise_mask' in latent_image: + # noise_latent = \ + # impact_sampling.separated_sample(refiner_model, "enable", seed, advanced_steps, cfg, sampler_name, + # scheduler, refiner_positive, refiner_negative, latent_image, end_at_step, + # end_at_step, "enable") + + latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() + temp_latent = latent_compositor.composite(latent_image, temp_latent, 0, 0, False, latent_image['noise_mask'])[0] + + print(f"post: {end_at_step} .. {advanced_steps + 1} / {advanced_steps}") + refined_latent = separated_sample(refiner_model, False, seed, advanced_steps, cfg, sampler_name, scheduler, + refiner_positive, refiner_negative, temp_latent, end_at_step, advanced_steps + 1, False) + + return refined_latent + + +class KSamplerAdvancedWrapper: + params = None + + def __init__(self, model, cfg, sampler_name, scheduler, positive, negative): + self.params = model, cfg, sampler_name, scheduler, positive, negative + + def sample_advanced(self, add_noise, seed, steps, latent_image, start_at_step, end_at_step, return_with_leftover_noise, hook=None, + recovery_mode="ratio additional", recovery_sampler="AUTO", recovery_sigma_ratio=1.0): + + model, cfg, sampler_name, scheduler, positive, negative = self.params + + if hook is not None: + model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent = hook.pre_ksample_advanced(model, add_noise, seed, steps, cfg, sampler_name, scheduler, + positive, negative, latent_image, start_at_step, end_at_step, + return_with_leftover_noise) + + if recovery_mode != 'DISABLE' and sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu', 'dpmpp_2m_sde', 'dpmpp_2m_sde_gpu', 'dpmpp_3m_sde', 'dpmpp_3m_sde_gpu']: + base_image = latent_image.copy() + if recovery_mode == "ratio between": + sigma_ratio = 1.0 - recovery_sigma_ratio + else: + sigma_ratio = 1.0 + else: + base_image = None + sigma_ratio = 1.0 + + try: + if sigma_ratio > 0: + latent_image = separated_sample(model, add_noise, seed, steps, cfg, sampler_name, scheduler, + positive, negative, latent_image, start_at_step, end_at_step, + return_with_leftover_noise, sigma_ratio=sigma_ratio) + except ValueError as e: + if str(e) == 'sigma_min and sigma_max must not be 0': + print(f"\nWARN: sampling skipped - sigma_min and sigma_max are 0") + return latent_image + + if (recovery_sigma_ratio > 0 and recovery_mode != 'DISABLE' and + sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu', 'dpmpp_2m_sde', 'dpmpp_2m_sde_gpu', 'dpmpp_3m_sde', 'dpmpp_3m_sde_gpu']): + compensate = 0 if sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu', 'dpmpp_2m_sde', 'dpmpp_2m_sde_gpu', 'dpmpp_3m_sde', 'dpmpp_3m_sde_gpu'] else 2 + if recovery_sampler == "AUTO": + recovery_sampler = 'dpm_fast' if sampler_name in ['uni_pc', 'uni_pc_bh2', 'dpmpp_sde', 'dpmpp_sde_gpu'] else 'dpmpp_2m' + + latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() + + noise_mask = latent_image['noise_mask'] + + if len(noise_mask.shape) == 4: + noise_mask = noise_mask.squeeze(0).squeeze(0) + + latent_image = latent_compositor.composite(base_image, latent_image, 0, 0, False, noise_mask)[0] + + try: + latent_image = separated_sample(model, add_noise, seed, steps, cfg, recovery_sampler, scheduler, + positive, negative, latent_image, start_at_step-compensate, end_at_step, + return_with_leftover_noise, sigma_ratio=recovery_sigma_ratio) + except ValueError as e: + if str(e) == 'sigma_min and sigma_max must not be 0': + print(f"\nWARN: sampling skipped - sigma_min and sigma_max are 0") + + return latent_image + + +class KSamplerWrapper: + params = None + + def __init__(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise): + self.params = model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise + + def sample(self, latent_image, hook=None): + model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise = self.params + + if hook is not None: + model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise = \ + hook.pre_ksample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, + denoise) + + return nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, + denoise=denoise)[0] diff --git a/modules/impact/segs_nodes.py b/modules/impact/segs_nodes.py index 856a7bc..c37c976 100644 --- a/modules/impact/segs_nodes.py +++ b/modules/impact/segs_nodes.py @@ -37,7 +37,7 @@ class SEGSDetailer: "optional": { "refiner_basic_pipe_opt": ("BASIC_PIPE",), "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), - "noise_mask_feather": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}), + "noise_mask_feather": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), } } diff --git a/modules/impact/special_samplers.py b/modules/impact/special_samplers.py index 29c9d21..e93c34e 100644 --- a/modules/impact/special_samplers.py +++ b/modules/impact/special_samplers.py @@ -1,11 +1,10 @@ -import time - -import comfy import math import impact.core as core from impact.utils import * from nodes import MAX_RESOLUTION import nodes +from impact.impact_sampling import KSamplerWrapper, KSamplerAdvancedWrapper + class TiledKSamplerProvider: @classmethod @@ -57,7 +56,7 @@ class KSamplerProvider: def doit(self, seed, steps, cfg, sampler_name, scheduler, denoise, basic_pipe): model, _, _, positive, negative = basic_pipe - sampler = core.KSamplerWrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise) + sampler = KSamplerWrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise) return (sampler, ) @@ -79,7 +78,7 @@ class KSamplerAdvancedProvider: def doit(self, cfg, sampler_name, scheduler, basic_pipe): model, _, _, positive, negative = basic_pipe - sampler = core.KSamplerAdvancedWrapper(model, cfg, sampler_name, scheduler, positive, negative) + sampler = KSamplerAdvancedWrapper(model, cfg, sampler_name, scheduler, positive, negative) return (sampler, ) @@ -167,10 +166,10 @@ class TwoAdvancedSamplersForMask: return_with_leftover_noise = "enable" if i+1 != adv_steps else "disable" new_latent_image['noise_mask'] = inv_mask - new_latent_image = base_sampler.sample_advanced(add_noise, seed, adv_steps, new_latent_image, i, i + 1, "enable", recover_special_sampler=True) + new_latent_image = base_sampler.sample_advanced(add_noise, seed, adv_steps, new_latent_image, i, i + 1, "enable", recovery_mode="ratio additional") new_latent_image['noise_mask'] = mask_erosion - new_latent_image = mask_sampler.sample_advanced("disable", seed, adv_steps, new_latent_image, i, i + 1, return_with_leftover_noise, recover_special_sampler=True) + new_latent_image = mask_sampler.sample_advanced("disable", seed, adv_steps, new_latent_image, i, i + 1, return_with_leftover_noise, recovery_mode="ratio additional") del new_latent_image['noise_mask'] @@ -287,6 +286,9 @@ class RegionalSampler: "regional_prompts": ("REGIONAL_PROMPTS", ), "overlap_factor": ("INT", {"default": 10, "min": 0, "max": 10000}), "restore_latent": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "additional_mode": (["DISABLE", "ratio additional", "ratio between"], {"default": "ratio between"}), + "additional_sampler": (["AUTO", "euler", "heun", "heunpp2", "dpm_2", "dpm_fast", "dpmpp_2m", "ddpm"],), + "additional_sigma_ratio": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), }, "hidden": {"unique_id": "UNIQUE_ID"}, } @@ -314,7 +316,8 @@ class RegionalSampler: return mask_erosion[:, :, :w, :h].round() - def doit(self, seed, seed_2nd, seed_2nd_mode, steps, base_only_steps, denoise, samples, base_sampler, regional_prompts, overlap_factor, restore_latent, unique_id=None): + def doit(self, seed, seed_2nd, seed_2nd_mode, steps, base_only_steps, denoise, samples, base_sampler, regional_prompts, overlap_factor, restore_latent, + additional_mode, additional_sampler, additional_sigma_ratio, unique_id=None): if restore_latent: latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() else: @@ -332,12 +335,12 @@ class RegionalSampler: region_len = len(regional_prompts) total = steps*region_len - leftover_noise = 'disable' + leftover_noise = False if base_only_steps > 0: if seed_2nd_mode == 'ignore': - leftover_noise = 'enable' + leftover_noise = True - samples = base_sampler.sample_advanced("enable", seed, adv_steps, samples, start_at_step, start_at_step + base_only_steps, leftover_noise, recover_special_sampler=False) + samples = base_sampler.sample_advanced(True, seed, adv_steps, samples, start_at_step, start_at_step + base_only_steps, leftover_noise, recovery_mode="DISABLE") if seed_2nd_mode == "seed+seed_2nd": seed += seed_2nd @@ -353,16 +356,17 @@ class RegionalSampler: new_latent_image = samples.copy() base_latent_image = None - if leftover_noise != 'enable': - add_noise = "enable" + if not leftover_noise: + add_noise = True else: - add_noise = "disable" + add_noise = False for i in range(start_at_step+base_only_steps, adv_steps): core.update_node_status(unique_id, f"{i}/{steps} steps | ", ((i-start_at_step)*region_len)/total) new_latent_image['noise_mask'] = inv_mask - new_latent_image = base_sampler.sample_advanced(add_noise, seed, adv_steps, new_latent_image, i, i + 1, "enable", recover_special_sampler=True) + new_latent_image = base_sampler.sample_advanced(add_noise, seed, adv_steps, new_latent_image, i, i + 1, True, + recovery_mode=additional_mode, recovery_sampler=additional_sampler, recovery_sigma_ratio=additional_sigma_ratio) if restore_latent: if 'noise_mask' in new_latent_image: @@ -379,8 +383,8 @@ class RegionalSampler: region_mask = regional_prompt.get_mask_erosion(overlap_factor).squeeze(0).squeeze(0) new_latent_image['noise_mask'] = region_mask - new_latent_image = regional_prompt.sampler.sample_advanced("disable", seed, adv_steps, new_latent_image, - i, i + 1, "enable", recover_special_sampler=True) + new_latent_image = regional_prompt.sampler.sample_advanced(False, seed, adv_steps, new_latent_image, i, i + 1, True, + recovery_mode=additional_mode, recovery_sampler=additional_sampler, recovery_sigma_ratio=additional_sigma_ratio) if restore_latent: del new_latent_image['noise_mask'] @@ -389,7 +393,7 @@ class RegionalSampler: j += 1 - add_noise = 'disable' + add_noise = False # finalize core.update_node_status(unique_id, f"finalize") @@ -399,7 +403,8 @@ class RegionalSampler: base_latent_image = new_latent_image new_latent_image['noise_mask'] = inv_mask - new_latent_image = base_sampler.sample_advanced("disable", seed, adv_steps, new_latent_image, adv_steps, adv_steps+1, "disable", recover_special_sampler=False) + new_latent_image = base_sampler.sample_advanced(False, seed, adv_steps, new_latent_image, adv_steps, adv_steps+1, False, + recovery_mode=additional_mode, recovery_sampler=additional_sampler, recovery_sigma_ratio=additional_sigma_ratio) core.update_node_status(unique_id, f"{steps}/{steps} steps", total) core.update_node_status(unique_id, "", None) @@ -428,6 +433,9 @@ class RegionalSamplerAdvanced: "latent_image": ("LATENT", ), "base_sampler": ("KSAMPLER_ADVANCED", ), "regional_prompts": ("REGIONAL_PROMPTS", ), + "additional_mode": (["DISABLE", "ratio additional", "ratio between"], {"default": "ratio between"}), + "additional_sampler": (["AUTO", "euler", "heun", "heunpp2", "dpm_2", "dpm_fast", "dpmpp_2m", "ddpm"],), + "additional_sigma_ratio": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), }, "hidden": {"unique_id": "UNIQUE_ID"}, } @@ -437,8 +445,9 @@ class RegionalSamplerAdvanced: CATEGORY = "ImpactPack/Regional" - def doit(self, add_noise, noise_seed, steps, start_at_step, end_at_step, overlap_factor, restore_latent, - return_with_leftover_noise, latent_image, base_sampler, regional_prompts, unique_id): + def doit(self, add_noise, noise_seed, steps, start_at_step, end_at_step, overlap_factor, restore_latent, return_with_leftover_noise, latent_image, base_sampler, regional_prompts, + additional_mode, additional_sampler, additional_sigma_ratio, unique_id): + if restore_latent: latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() else: @@ -458,13 +467,14 @@ class RegionalSamplerAdvanced: base_latent_image = None region_masks = {} - for i in range(start_at_step, end_at_step): + for i in range(start_at_step, end_at_step-1): core.update_node_status(unique_id, f"{start_at_step+i}/{end_at_step} steps | ", ((i-start_at_step)*region_len)/total) - cur_add_noise = "enable" if i == start_at_step and add_noise else "disable" + cur_add_noise = True if i == start_at_step and add_noise else False new_latent_image['noise_mask'] = inv_mask - new_latent_image = base_sampler.sample_advanced(cur_add_noise, noise_seed, steps, new_latent_image, i, i + 1, "enable", recover_special_sampler=True) + new_latent_image = base_sampler.sample_advanced(cur_add_noise, noise_seed, steps, new_latent_image, i, i + 1, True, + recovery_mode=additional_mode, recovery_sampler=additional_sampler, recovery_sigma_ratio=additional_sigma_ratio) if restore_latent: del new_latent_image['noise_mask'] @@ -484,8 +494,8 @@ class RegionalSamplerAdvanced: region_mask = region_masks[j] new_latent_image['noise_mask'] = region_mask - new_latent_image = regional_prompt.sampler.sample_advanced("disable", noise_seed, steps, new_latent_image, - i, i + 1, "enable", recover_special_sampler=True) + new_latent_image = regional_prompt.sampler.sample_advanced(False, noise_seed, steps, new_latent_image, i, i + 1, True, + recovery_mode=additional_mode, recovery_sampler=additional_sampler, recovery_sigma_ratio=additional_sigma_ratio) if restore_latent: del new_latent_image['noise_mask'] @@ -502,7 +512,8 @@ class RegionalSamplerAdvanced: base_latent_image = new_latent_image new_latent_image['noise_mask'] = inv_mask - new_latent_image = base_sampler.sample_advanced("disable", noise_seed, steps, new_latent_image, end_at_step, end_at_step+1, return_with_leftover_noise, recover_special_sampler=False) + new_latent_image = base_sampler.sample_advanced(False, noise_seed, steps, new_latent_image, end_at_step-1, end_at_step, return_with_leftover_noise, + recovery_mode=additional_mode, recovery_sampler=additional_sampler, recovery_sigma_ratio=additional_sigma_ratio) core.update_node_status(unique_id, f"{end_at_step}/{end_at_step} steps", total) core.update_node_status(unique_id, "", None) @@ -539,7 +550,7 @@ class KSamplerBasicPipe: def sample(self, basic_pipe, seed, steps, cfg, sampler_name, scheduler, latent_image, denoise=1.0): model, clip, vae, positive, negative = basic_pipe latent = nodes.KSampler().sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise)[0] - return (basic_pipe, latent, vae) + return basic_pipe, latent, vae class KSamplerAdvancedBasicPipe: @@ -579,4 +590,5 @@ class KSamplerAdvancedBasicPipe: return_with_leftover_noise = "disable" latent = nodes.KSamplerAdvanced().sample(model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise)[0] - return (basic_pipe, latent, vae) + return basic_pipe, latent, vae +