From 7f21da704424df2c7cb22865568903ece348134c Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Fri, 22 Sep 2023 13:00:48 +0900 Subject: [PATCH] feat: RegionalSamplerAdvanced improve: RegionalSampler --- README.md | 10 +- __init__.py | 1 + js/impact-pack.js | 3 +- modules/impact/config.py | 2 +- modules/impact/core.py | 51 +++++--- modules/impact/special_samplers.py | 193 +++++++++++++++++++++++++++-- 6 files changed, 232 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 0cd4a40..385a576 100644 --- a/README.md +++ b/README.md @@ -146,8 +146,14 @@ This takes latent as input and outputs latent as the result. * If the `Inspire Pack` is installed, you can use **Lora Block Weight** in the form of `LBW=lbw spec;` * ``, ``, `` -* RegionalSampler, CombineRegionalPrompts, RegionalPrompt - experimental feature -- multiple region version of TwoAdvancedSamplersForMask +* Regional Sampling - These nodes offer the capability to divide regions and perform partial sampling using a mask. Unlike TwoSamplersForMask, sampling for each region is applied during each step. + * RegionalPrompt - This node combines a **mask** for specifying regions and the **sampler** to apply to each region to create `REGIONAL_PROMPTS`. + * CombineRegionalPrompts - Combine multiple `REGIONAL_PROMPTS` to create a single `REGIONAL_PROMPTS`. + * RegionalSampler - This node performs sampling using a base sampler and regional prompts. Sampling by the base sampler is executed at each step, while sampling for each region is performed through the sampler bound to each region. + * overlap_factor - Specifies the amount of overlap for each region to blend well with the area outside the mask. + * latent_restore - When sampling each region, restore the areas outside the mask to the base latent, preventing additional noise from being introduced outside the mask during region sampling. + * RegionalSamplerAdvanced - This is the Advanced version of the RegionalSampler. You can control it using `step` instead of `denoise`. + * NOTE: The `sde` sampler and `uni_pc` sampler introduce additional noise during each step of the sampling process. To mitigate this, when sampling each region, the `uni_pc` sampler applies additional `dpmpp_fast`, and the sde sampler applies the `dpmpp_2m` sampler as an additional measure. * KSampler (pipe), KSampler (advanced/pipe) diff --git a/__init__.py b/__init__.py index 5fe719e..045b6ad 100644 --- a/__init__.py +++ b/__init__.py @@ -225,6 +225,7 @@ NODE_CLASS_MAPPINGS = { "ImpactMakeImageList": MakeImageList, "RegionalSampler": RegionalSampler, + "RegionalSamplerAdvanced": RegionalSamplerAdvanced, "CombineRegionalPrompts": CombineRegionalPrompts, "RegionalPrompt": RegionalPrompt, diff --git a/js/impact-pack.js b/js/impact-pack.js index a8a035a..54adeee 100644 --- a/js/impact-pack.js +++ b/js/impact-pack.js @@ -187,7 +187,8 @@ app.registerExtension({ }, async beforeRegisterNodeDef(nodeType, nodeData, app) { - if (nodeData.name == "IterativeLatentUpscale" || nodeData.name == "IterativeImageUpscale" || nodeData.name == "RegionalSampler") { + if (nodeData.name == "IterativeLatentUpscale" || nodeData.name == "IterativeImageUpscale" + || nodeData.name == "RegionalSampler"|| nodeData.name == "RegionalSamplerAdvanced") { impactProgressBadge.addStatusHandler(nodeType); } diff --git a/modules/impact/config.py b/modules/impact/config.py index 54d8901..36b4053 100644 --- a/modules/impact/config.py +++ b/modules/impact/config.py @@ -2,7 +2,7 @@ import configparser import os -version = "V4.8.5" +version = "V4.9" dependency_version = 11 diff --git a/modules/impact/core.py b/modules/impact/core.py index 3007741..8adad59 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -53,9 +53,9 @@ def ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, 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] + 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 = \ @@ -65,14 +65,14 @@ def ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() temp_latent = \ - latent_compositor.composite(latent_image, temp_latent, 0, 0, False, latent_image['noise_mask'])[0] + 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] + 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 @@ -251,9 +251,6 @@ def composite_to(dest_latent, crop_region, src_latent): # composite to original latent lc = nodes.LatentComposite() - - # 현재 mask 를 고려한 composite 가 없음... 이거 처리 필요. - orig_image = lc.composite(dest_latent, src_latent, x1, y1) return orig_image[0] @@ -925,7 +922,7 @@ class KSamplerAdvancedWrapper: 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): + return_with_leftover_noise, hook=None, recover_special_sampler=False): model, cfg, sampler_name, scheduler, positive, negative = self.params if hook is not None: @@ -934,9 +931,33 @@ class KSamplerAdvancedWrapper: positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise) - return 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] + 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 + + 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] + + 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' + print(f"recover latent!!: {sampler_name} ->") + 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] + 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] + + return latent_image class PixelKSampleHook: diff --git a/modules/impact/special_samplers.py b/modules/impact/special_samplers.py index 40351ad..6b62476 100644 --- a/modules/impact/special_samplers.py +++ b/modules/impact/special_samplers.py @@ -1,3 +1,5 @@ +import time + import comfy import math import impact.core as core @@ -165,10 +167,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") + 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['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) + 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) del new_latent_image['noise_mask'] @@ -225,7 +227,8 @@ class RegionalSampler: "samples": ("LATENT", ), "base_sampler": ("KSAMPLER_ADVANCED", ), "regional_prompts": ("REGIONAL_PROMPTS", ), - "overlap_factor": ("INT", {"default": 10, "min": 0, "max": 10000}) + "overlap_factor": ("INT", {"default": 10, "min": 0, "max": 10000}), + "latent_restore": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}) }, "hidden": {"unique_id": "UNIQUE_ID"}, } @@ -253,7 +256,11 @@ class RegionalSampler: return mask_erosion[:, :, :w, :h].round() - def doit(self, seed, steps, denoise, samples, base_sampler, regional_prompts, overlap_factor, unique_id): + def doit(self, seed, steps, denoise, samples, base_sampler, regional_prompts, overlap_factor, latent_restore, unique_id=None): + if latent_restore: + latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() + else: + latent_compositor = None masks = [regional_prompt.mask.numpy() for regional_prompt in regional_prompts] masks = [np.ceil(mask).astype(np.int32) for mask in masks] @@ -268,31 +275,159 @@ class RegionalSampler: total = steps*region_len new_latent_image = samples.copy() + base_latent_image = None for i in range(start_at_step, adv_steps): core.update_node_status(unique_id, f"{i}/{steps} steps | ", (i*region_len)/total) add_noise = "enable" if i == start_at_step else "disable" - 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") + new_latent_image = base_sampler.sample_advanced(add_noise, seed, adv_steps, new_latent_image, i, i + 1, "enable", recover_special_sampler=True) + + if latent_restore: + del new_latent_image['noise_mask'] + base_latent_image = new_latent_image.copy() j = 1 for regional_prompt in regional_prompts: + if latent_restore: + new_latent_image = base_latent_image.copy() + core.update_node_status(unique_id, f"{i}/{steps} steps | {j}/{region_len}", (i*region_len + j)/total) - new_latent_image['noise_mask'] = regional_prompt.get_mask_erosion(overlap_factor) + + 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, return_with_leftover_noise) + i, i + 1, "enable", recover_special_sampler=True) + + if latent_restore: + del new_latent_image['noise_mask'] + base_latent_image = latent_compositor.composite(base_latent_image, new_latent_image, 0, 0, False, region_mask)[0] + new_latent_image = base_latent_image + j += 1 + # finalize + core.update_node_status(unique_id, f"finalize") + if base_latent_image is not None: + new_latent_image = base_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) + core.update_node_status(unique_id, f"{steps}/{steps} steps", total) core.update_node_status(unique_id, "", None) - del new_latent_image['noise_mask'] + if latent_restore: + new_latent_image = base_latent_image + + if 'noise_mask' in new_latent_image: + del new_latent_image['noise_mask'] return (new_latent_image, ) + +class RegionalSamplerAdvanced: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "add_noise": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}), + "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}), + "overlap_factor": ("INT", {"default": 10, "min": 0, "max": 10000}), + "latent_restore": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "return_with_leftover_noise": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + "latent_image": ("LATENT", ), + "base_sampler": ("KSAMPLER_ADVANCED", ), + "regional_prompts": ("REGIONAL_PROMPTS", ), + }, + "hidden": {"unique_id": "UNIQUE_ID"}, + } + + RETURN_TYPES = ("LATENT", ) + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Regional" + + def doit(self, add_noise, noise_seed, steps, start_at_step, end_at_step, overlap_factor, latent_restore, + return_with_leftover_noise, latent_image, base_sampler, regional_prompts, unique_id): + if latent_restore: + latent_compositor = nodes.NODE_CLASS_MAPPINGS['LatentCompositeMasked']() + else: + latent_compositor = None + + masks = [regional_prompt.mask.numpy() for regional_prompt in regional_prompts] + masks = [np.ceil(mask).astype(np.int32) for mask in masks] + combined_mask = torch.from_numpy(np.bitwise_or.reduce(masks)) + + inv_mask = torch.where(combined_mask == 0, torch.tensor(1.0), torch.tensor(0.0)) + + region_len = len(regional_prompts) + end_at_step = min(steps, end_at_step) + total = (end_at_step - start_at_step) * region_len + + new_latent_image = latent_image.copy() + base_latent_image = None + region_masks = {} + + for i in range(start_at_step, end_at_step): + 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" + + 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) + + if latent_restore: + del new_latent_image['noise_mask'] + base_latent_image = new_latent_image.copy() + + j = 1 + for regional_prompt in regional_prompts: + if latent_restore: + new_latent_image = base_latent_image.copy() + + core.update_node_status(unique_id, f"{start_at_step+i}/{end_at_step} steps | {j}/{region_len}", ((i-start_at_step)*region_len + j)/total) + + if j not in region_masks: + region_mask = regional_prompt.get_mask_erosion(overlap_factor).squeeze(0).squeeze(0) + region_masks[j] = region_mask + else: + 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) + + if latent_restore: + del new_latent_image['noise_mask'] + base_latent_image = latent_compositor.composite(base_latent_image, new_latent_image, 0, 0, False, region_mask)[0] + new_latent_image = base_latent_image + + j += 1 + + # finalize + core.update_node_status(unique_id, f"finalize") + if base_latent_image is not None: + new_latent_image = base_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, "disable", recover_special_sampler=False) + + core.update_node_status(unique_id, f"{end_at_step}/{end_at_step} steps", total) + core.update_node_status(unique_id, "", None) + + if latent_restore: + new_latent_image = base_latent_image + + if 'noise_mask' in new_latent_image: + del new_latent_image['noise_mask'] + + return (new_latent_image, ) + + class KSamplerBasicPipe: @classmethod def INPUT_TYPES(s): @@ -357,3 +492,43 @@ class KSamplerAdvancedBasicPipe: 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) + + +class KSamplerAdvancedBasicPipe: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"basic_pipe": ("BASIC_PIPE",), + "add_noise": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "noise_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, ), + "latent_image": ("LATENT", ), + "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}), + "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}), + "return_with_leftover_noise": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + } + } + + RETURN_TYPES = ("BASIC_PIPE", "LATENT", "VAE") + FUNCTION = "sample" + + CATEGORY = "sampling" + + def sample(self, basic_pipe, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0): + model, clip, vae, positive, negative = basic_pipe + + if add_noise: + add_noise = "enabled" + else: + add_noise = "disabled" + + if return_with_leftover_noise: + return_with_leftover_noise = "enabled" + else: + return_with_leftover_noise = "disabled" + + 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)