From decff0d4c6ff98a6a26414a8d60cfca333233f80 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Sat, 27 Jan 2024 12:27:48 +0900 Subject: [PATCH] improve: Regional Sampler - use impact_sampling instead of KSamplerAdvanced. - Handle the _sde sampler in a more appropriate manner. --- __init__.py | 7 +- modules/impact/animatediff_nodes.py | 145 ++++++++++++++++++ modules/impact/config.py | 2 +- modules/impact/core.py | 158 +++----------------- modules/impact/impact_pack.py | 12 +- modules/impact/impact_sampling.py | 222 ++++++++++++++++++++++++++++ modules/impact/segs_nodes.py | 92 +----------- modules/impact/special_samplers.py | 70 +++++---- 8 files changed, 439 insertions(+), 269 deletions(-) create mode 100644 modules/impact/animatediff_nodes.py create mode 100644 modules/impact/impact_sampling.py diff --git a/__init__.py b/__init__.py index 692fb65..6569b79 100644 --- a/__init__.py +++ b/__init__.py @@ -106,6 +106,7 @@ 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 * import threading @@ -139,6 +140,7 @@ NODE_CLASS_MAPPINGS = { "DetailerForEachDebug": DetailerForEachTest, "DetailerForEachPipe": DetailerForEachPipe, "DetailerForEachDebugPipe": DetailerForEachTestPipe, + "DetailerForEachPipeForAnimateDiff": DetailerForEachPipeForAnimateDiff, "SAMDetectorCombined": SAMDetectorCombined, "SAMDetectorSegmented": SAMDetectorSegmented, @@ -349,12 +351,13 @@ NODE_DISPLAY_NAME_MAPPINGS = { "DetailerForEachPipe": "Detailer (SEGS/pipe)", "DetailerForEachDebug": "DetailerDebug (SEGS)", "DetailerForEachDebugPipe": "DetailerDebug (SEGS/pipe)", - "SEGSDetailerForAnimateDiff": "Detailer For AnimateDiff (SEGS/pipe)", + "SEGSDetailerForAnimateDiff": "SEGSDetailer For AnimateDiff (SEGS/pipe)", + "DetailerForEachPipeForAnimateDiff": "Detailer For AnimateDiff (SEGS/pipe)", "SAMDetectorCombined": "SAMDetector (combined)", "SAMDetectorSegmented": "SAMDetector (segmented)", "FaceDetailerPipe": "FaceDetailer (pipe)", - "MaskDetailerPipe": "MaskDetailer (Pipe)", + "MaskDetailerPipe": "MaskDetailer (pipe)", "FromDetailerPipeSDXL": "FromDetailer (SDXL/pipe)", "BasicPipeToDetailerPipeSDXL": "BasicPipe -> DetailerPipe (SDXL)", diff --git a/modules/impact/animatediff_nodes.py b/modules/impact/animatediff_nodes.py new file mode 100644 index 0000000..c26228b --- /dev/null +++ b/modules/impact/animatediff_nodes.py @@ -0,0 +1,145 @@ +from nodes import MAX_RESOLUTION +from impact.utils import * +import impact.core as core +from impact.core import SEG +from impact.segs_nodes import SEGSPaste + + +class SEGSDetailerForAnimateDiff: + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "image_frames": ("IMAGE", ), + "segs": ("SEGS", ), + "guide_size": ("FLOAT", {"default": 256, "min": 64, "max": MAX_RESOLUTION, "step": 8}), + "guide_size_for": ("BOOLEAN", {"default": True, "label_on": "bbox", "label_off": "crop_region"}), + "max_size": ("FLOAT", {"default": 768, "min": 64, "max": MAX_RESOLUTION, "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}), + "basic_pipe": ("BASIC_PIPE",), + "refiner_ratio": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0}) + }, + "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": 0, "min": 0, "max": 100, "step": 1}), + } + } + + RETURN_TYPES = ("SEGS",) + RETURN_NAMES = ("segs",) + OUTPUT_IS_LIST = (False,) + + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Detailer" + + @staticmethod + def do_detail(image_frames, segs, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, + denoise, basic_pipe, refiner_ratio=None, refiner_basic_pipe_opt=None, inpaint_model=False, noise_mask_feather=0): + + model, clip, vae, positive, negative = basic_pipe + if refiner_basic_pipe_opt is None: + refiner_model, refiner_clip, refiner_positive, refiner_negative = None, None, None, None + else: + refiner_model, refiner_clip, _, refiner_positive, refiner_negative = refiner_basic_pipe_opt + + segs = core.segs_scale_match(segs, image_frames.shape) + + new_segs = [] + + for seg in segs[1]: + cropped_image_frames = None + + for image in image_frames: + image = image.unsqueeze(0) + cropped_image = seg.cropped_image if seg.cropped_image is not None else crop_tensor4(image, seg.crop_region) + cropped_image = to_tensor(cropped_image) + if cropped_image_frames is None: + cropped_image_frames = cropped_image + else: + cropped_image_frames = torch.concat((cropped_image_frames, cropped_image), dim=0) + + cropped_image_frames = cropped_image_frames.numpy() + enhanced_image_tensor = core.enhance_detail_for_animatediff(cropped_image_frames, model, clip, vae, guide_size, guide_size_for, max_size, + seg.bbox, seed, steps, cfg, sampler_name, scheduler, + positive, negative, denoise, seg.cropped_mask, + refiner_ratio=refiner_ratio, refiner_model=refiner_model, + refiner_clip=refiner_clip, refiner_positive=refiner_positive, + refiner_negative=refiner_negative, + inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather) + + if enhanced_image_tensor is None: + new_cropped_image = cropped_image_frames + else: + new_cropped_image = enhanced_image_tensor.numpy() + + new_seg = SEG(new_cropped_image, seg.cropped_mask, seg.confidence, seg.crop_region, seg.bbox, seg.label, None) + new_segs.append(new_seg) + + return (segs[0], new_segs) + + def doit(self, image_frames, segs, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, + denoise, basic_pipe, refiner_ratio=None, refiner_basic_pipe_opt=None, inpaint_model=False, noise_mask_feather=0): + + segs = SEGSDetailerForAnimateDiff.do_detail(image_frames, segs, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, + scheduler, denoise, basic_pipe, refiner_ratio, refiner_basic_pipe_opt, + inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather) + + return (segs,) + + +class DetailerForEachPipeForAnimateDiff: + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "image_frames": ("IMAGE", ), + "segs": ("SEGS", ), + "guide_size": ("FLOAT", {"default": 384, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 8}), + "guide_size_for": ("BOOLEAN", {"default": True, "label_on": "bbox", "label_off": "crop_region"}), + "max_size": ("FLOAT", {"default": 1024, "min": 64, "max": nodes.MAX_RESOLUTION, "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}), + "basic_pipe": ("BASIC_PIPE", ), + "refiner_ratio": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0}), + }, + "optional": { + "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": 0, "min": 0, "max": 100, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE", "SEGS", "BASIC_PIPE") + RETURN_NAMES = ("image", "segs", "basic_pipe") + OUTPUT_IS_LIST = (False, False, False, True) + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Detailer" + + @staticmethod + def doit(image_frames, segs, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, + denoise, feather, basic_pipe, refiner_ratio=None, detailer_hook=None, refiner_basic_pipe_opt=None, + inpaint_model=False, noise_mask_feather=0): + + enhanced_segs = [] + for sub_seg in segs[1]: + single_seg = segs[0], [sub_seg] + enhanced_seg = SEGSDetailerForAnimateDiff().do_detail(image_frames, single_seg, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, + denoise, basic_pipe, refiner_ratio, refiner_basic_pipe_opt, inpaint_model, noise_mask_feather) + + image_frames = SEGSPaste.doit(image_frames, enhanced_seg, feather, alpha=255)[0] + enhanced_segs += enhanced_seg[1] + + new_segs = segs[0], enhanced_segs + return image_frames, new_segs, basic_pipe diff --git a/modules/impact/config.py b/modules/impact/config.py index 1168634..fc40699 100644 --- a/modules/impact/config.py +++ b/modules/impact/config.py @@ -2,7 +2,7 @@ import configparser import os -version_code = [4, 68] +version_code = [4, 69] 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 be15ecf..c37c976 100644 --- a/modules/impact/segs_nodes.py +++ b/modules/impact/segs_nodes.py @@ -1,8 +1,6 @@ import os import sys -import torch - import impact.impact_server from nodes import MAX_RESOLUTION @@ -39,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}), } } @@ -125,94 +123,6 @@ class SEGSDetailer: return (segs, cnet_pil_list) -class SEGSDetailerForAnimateDiff: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "image_frames": ("IMAGE", ), - "segs": ("SEGS", ), - "guide_size": ("FLOAT", {"default": 256, "min": 64, "max": MAX_RESOLUTION, "step": 8}), - "guide_size_for": ("BOOLEAN", {"default": True, "label_on": "bbox", "label_off": "crop_region"}), - "max_size": ("FLOAT", {"default": 768, "min": 64, "max": MAX_RESOLUTION, "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}), - "basic_pipe": ("BASIC_PIPE",), - "refiner_ratio": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0}) - }, - "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}), - } - } - - RETURN_TYPES = ("SEGS",) - RETURN_NAMES = ("segs",) - OUTPUT_IS_LIST = (False,) - - FUNCTION = "doit" - - CATEGORY = "ImpactPack/Detailer" - - @staticmethod - def do_detail(image_frames, segs, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, - denoise, basic_pipe, refiner_ratio=None, refiner_basic_pipe_opt=None, inpaint_model=False, noise_mask_feather=0): - - model, clip, vae, positive, negative = basic_pipe - if refiner_basic_pipe_opt is None: - refiner_model, refiner_clip, refiner_positive, refiner_negative = None, None, None, None - else: - refiner_model, refiner_clip, _, refiner_positive, refiner_negative = refiner_basic_pipe_opt - - segs = core.segs_scale_match(segs, image_frames.shape) - - new_segs = [] - - for seg in segs[1]: - cropped_image_frames = None - - for image in image_frames: - image = image.unsqueeze(0) - cropped_image = seg.cropped_image if seg.cropped_image is not None else crop_tensor4(image, seg.crop_region) - cropped_image = to_tensor(cropped_image) - if cropped_image_frames is None: - cropped_image_frames = cropped_image - else: - cropped_image_frames = torch.concat((cropped_image_frames, cropped_image), dim=0) - - cropped_image_frames = cropped_image_frames.numpy() - enhanced_image_tensor = core.enhance_detail_for_animatediff(cropped_image_frames, model, clip, vae, guide_size, guide_size_for, max_size, - seg.bbox, seed, steps, cfg, sampler_name, scheduler, - positive, negative, denoise, seg.cropped_mask, - refiner_ratio=refiner_ratio, refiner_model=refiner_model, - refiner_clip=refiner_clip, refiner_positive=refiner_positive, - refiner_negative=refiner_negative, - inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather) - - if enhanced_image_tensor is None: - new_cropped_image = cropped_image_frames - else: - new_cropped_image = enhanced_image_tensor.numpy() - - new_seg = SEG(new_cropped_image, seg.cropped_mask, seg.confidence, seg.crop_region, seg.bbox, seg.label, None) - new_segs.append(new_seg) - - return (segs[0], new_segs) - - def doit(self, image_frames, segs, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, - denoise, basic_pipe, refiner_ratio=None, refiner_basic_pipe_opt=None, inpaint_model=False, noise_mask_feather=0): - - segs = SEGSDetailerForAnimateDiff.do_detail(image_frames, segs, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, - scheduler, denoise, basic_pipe, refiner_ratio, refiner_basic_pipe_opt, - inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather) - - return (segs,) - - class SEGSPaste: @classmethod def INPUT_TYPES(s): 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 +