From 47ded8aab544e300830af65c65115011d2d7edd0 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Fri, 12 May 2023 22:27:05 +0900 Subject: [PATCH] add nodes: TwoSamplersForMaskUpscalerProvider, TwoSamplersForMaskUpscalerProviderPipe, TiledKSamplerProvider --- README.md | 7 +- __init__.py | 12 +++- impact_config.py | 2 +- impact_core.py | 170 ++++++++++++++++++++++++++++++++++++++++--- impact_pack.py | 184 ++++++++++++++++++++++++++++++++++------------- install.py | 2 +- legacy_nodes.py | 56 ++++++++++++++- 7 files changed, 369 insertions(+), 64 deletions(-) diff --git a/README.md b/README.md index fc90555..c1b41a9 100644 --- a/README.md +++ b/README.md @@ -53,7 +53,12 @@ This takes latent as input and outputs latent as the result. * TwoSamplersForMask - This node can apply two samplers depending on the mask area. The base_sampler is applied to the area where the mask is 0, while the mask_sampler is applied to the area where the mask is 1. * Note: The latent encoded through VAEEncodeForInpaint cannot be used. -* KSamplerProvider - This is a wrapper that enables KSampler to be used in TwoSamplersForMask. +* KSamplerProvider - This is a wrapper that enables KSampler to be used in TwoSamplersForMask TwoSamplersForMaskUpscalerProvider. +* TiledKSamplerProvider - ComfyUI_TiledKSampler is a wrapper that provides KSAMPLER. + * You need to install the [ComfyUI_TiledKSampler](https://github.com/BlenderNeko/ComfyUI_TiledKSampler) node extension. + +* TwoSamplersForMaskUpscalerProvider - This is an Upscaler that extends TwoSamplersForMask to be used in Iterative Upscale. + * TwoSamplersForMaskUpscalerProviderPipe - pipe version of TwoSamplersForMaskUpscalerProvider. # Depercated * The following nodes have been kept only for compatibility with existing workflows, and are no longer supported. Please replace them with new nodes. diff --git a/__init__.py b/__init__.py index ca53bb6..3521d66 100644 --- a/__init__.py +++ b/__init__.py @@ -91,6 +91,8 @@ NODE_CLASS_MAPPINGS = { "IterativeImageUpscale": IterativeImageUpscale, "PixelTiledKSampleUpscalerProvider": PixelTiledKSampleUpscalerProvider, "PixelTiledKSampleUpscalerProviderPipe": PixelTiledKSampleUpscalerProviderPipe, + "TwoSamplersForMaskUpscalerProvider": TwoSamplersForMaskUpscalerProvider, + "TwoSamplersForMaskUpscalerProviderPipe": TwoSamplersForMaskUpscalerProviderPipe, "PixelKSampleHookCombine": PixelKSampleHookCombine, "DenoiseScheduleHookProvider": DenoiseScheduleHookProvider, @@ -104,8 +106,6 @@ NODE_CLASS_MAPPINGS = { "MaskToSEGS": MaskToSEGS, "ToBinaryMask": ToBinaryMask, - "MaskPainter": MaskPainter, - "BboxDetectorSEGS": BboxDetectorForEach, "SegmDetectorSEGS": SegmDetectorForEach, "ONNXDetectorSEGS": ONNXDetectorForEach, @@ -116,7 +116,11 @@ NODE_CLASS_MAPPINGS = { "KSamplerProvider": KSamplerProvider, "TwoSamplersForMask": TwoSamplersForMask, + "TiledKSamplerProvider": TiledKSamplerProvider, + #"PreviewBridge": PreviewBridge, + + "MaskPainter": legacy_nodes.MaskPainter, "MMDetLoader": legacy_nodes.MMDetLoader, "SegsMaskCombine": legacy_nodes.SegsMaskCombine, "BboxDetectorForEach": legacy_nodes.BboxDetectorForEach, @@ -154,6 +158,10 @@ NODE_DISPLAY_NAME_MAPPINGS = { "IterativeLatentUpscale": "Iterative Upscale (Latent)", "IterativeImageUpscale": "Iterative Upscale (Image)", + "TwoSamplersForMaskUpscalerProvider": "TwoSamplersForMask Upscaler Provider", + "TwoSamplersForMaskUpscalerProviderPipe": "TwoSamplersForMask Upscaler Provider (pipe)", + + "MaskPainter": "MaskPainter (Legacy)", "MMDetLoader": "MMDetLoader (Legacy)", "SegsMaskCombine": "SegsMaskCombine (Legacy)", "BboxDetectorForEach": "BboxDetectorForEach (Legacy)", diff --git a/impact_config.py b/impact_config.py index d2065d5..da8434d 100644 --- a/impact_config.py +++ b/impact_config.py @@ -1,7 +1,7 @@ import configparser import os -version = "V2.4" +version = "V2.5" dependency_version = 1 diff --git a/impact_core.py b/impact_core.py index 94e1a17..36db0c4 100644 --- a/impact_core.py +++ b/impact_core.py @@ -4,6 +4,7 @@ import mmcv from mmdet.apis import (inference_detector, init_detector) from mmdet.evaluation import get_classes from segment_anything import SamPredictor +import torch.nn.functional as F from impact_utils import * from collections import namedtuple @@ -655,9 +656,14 @@ class KSamplerWrapper: 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): + def sample(self, latent_image, hook): model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise = self.params - return nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise) + + 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 PixelKSampleHook: @@ -821,10 +827,135 @@ def latent_upscale_on_pixel_space_with_model(samples, scale_method, upscale_mode return vae_encode(vae, pixels, use_tile, hook) +class TwoSamplersForMaskUpscaler: + params = None + upscale_model = None + hook_base = None + hook_mask = None + hook_full = None + use_tiled_vae = False + is_tiled = False + + def __init__(self, scale_method, sample_schedule, use_tiled_vae, base_sampler, mask_sampler, mask, vae, + full_sampler_opt=None, upscale_model_opt=None, hook_base_opt=None, hook_mask_opt=None, hook_full_opt=None): + mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])) + + self.params = scale_method, sample_schedule, use_tiled_vae, base_sampler, mask_sampler, mask, vae + self.upscale_model = upscale_model_opt + self.full_sampler = full_sampler_opt + self.hook_base = hook_base_opt + self.hook_mask = hook_mask_opt + self.hook_full = hook_full_opt + self.use_tiled_vae = use_tiled_vae + + def upscale(self, step_info, samples, upscale_factor, save_temp_prefix=None): + scale_method, sample_schedule, use_tiled_vae, base_sampler, mask_sampler, mask, vae = self.params + + self.prepare_hook(step_info) + + # upscale latent + if self.upscale_model is None: + upscaled_latent = latent_upscale_on_pixel_space(samples, scale_method, upscale_factor, vae, + use_tile=self.use_tiled_vae, + save_temp_prefix=save_temp_prefix, hook=self.hook_base) + else: + upscaled_latent = latent_upscale_on_pixel_space_with_model(samples, scale_method, self.upscale_model, upscale_factor, vae, + save_temp_prefix=save_temp_prefix, hook=self.hook_mask) + + return self.do_samples(step_info, base_sampler, mask_sampler, sample_schedule, mask, upscaled_latent) + + def prepare_hook(self, step_info): + if self.hook_base is not None: + self.hook_base.set_steps(step_info) + if self.hook_mask is not None: + self.hook_mask.set_steps(step_info) + if self.hook_full is not None: + self.hook_full.set_steps(step_info) + + def upscale_shape(self, step_info, samples, w, h, save_temp_prefix=None): + scale_method, sample_schedule, use_tiled_vae, base_sampler, mask_sampler, mask, vae = self.params + + self.prepare_hook(step_info) + + # upscale latent + if self.upscale_model is None: + upscaled_latent = latent_upscale_on_pixel_space_shape(samples, scale_method, w, h, vae, + use_tile=self.use_tiled_vae, + save_temp_prefix=save_temp_prefix, hook=self.hook_base) + else: + upscaled_latent = latent_upscale_on_pixel_space_with_model_shape(samples, scale_method, self.upscale_model, w, h, vae, + save_temp_prefix=save_temp_prefix, hook=self.hook_mask) + + return self.do_samples(step_info, base_sampler, mask_sampler, sample_schedule, mask, upscaled_latent) + + def is_full_sample_time(self, step_info, sample_schedule): + cur_step, total_step = step_info + + # make start from 1 instead of zero + cur_step += 1 + total_step += 1 + + if sample_schedule == "none": + return False + + elif sample_schedule == "interleave1": + return cur_step % 2 == 0 + + elif sample_schedule == "interleave2": + return cur_step % 3 == 0 + + elif sample_schedule == "interleave3": + return cur_step % 4 == 0 + + elif sample_schedule == "last1": + return cur_step == total_step + + elif sample_schedule == "last2": + return cur_step >= total_step-1 + + elif sample_schedule == "interleave1+last1": + return cur_step % 2 == 0 or cur_step >= total_step-1 + + elif sample_schedule == "interleave2+last1": + return cur_step % 2 == 0 or cur_step >= total_step-1 + + elif sample_schedule == "interleave3+last1": + return cur_step % 2 == 0 or cur_step >= total_step-1 + + def do_samples(self, step_info, base_sampler, mask_sampler, sample_schedule, mask, upscaled_latent): + if self.is_full_sample_time(step_info, sample_schedule): + print(f"step_info={step_info} / full time") + + upscaled_latent = base_sampler.sample(upscaled_latent, self.hook_base) + sampler = self.full_sampler if self.full_sampler is not None else base_sampler + return sampler.sample(upscaled_latent, self.hook_full) + + else: + print(f"step_info={step_info} / non-full time") + # upscale mask + upscaled_mask = F.interpolate(mask, size=(upscaled_latent['samples'].shape[2], upscaled_latent['samples'].shape[3]), + mode='bilinear', align_corners=True) + upscaled_mask = upscaled_mask[:, :, :upscaled_latent['samples'].shape[2], :upscaled_latent['samples'].shape[3]] + + # base sampler + upscaled_inv_mask = torch.where(upscaled_mask != 1.0, torch.tensor(1.0), torch.tensor(0.0)) + upscaled_latent['noise_mask'] = upscaled_inv_mask + upscaled_latent = base_sampler.sample(upscaled_latent, self.hook_base) + + # mask sampler + upscaled_latent['noise_mask'] = upscaled_mask + upscaled_latent = mask_sampler.sample(upscaled_latent, self.hook_mask) + + # remove mask + del upscaled_latent['noise_mask'] + return upscaled_latent + + class PixelKSampleUpscaler: params = None upscale_model = None hook = None + use_tiled_vae = False is_tiled = False def __init__(self, scale_method, model, vae, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, @@ -854,7 +985,7 @@ class PixelKSampleUpscaler: refined_latent = nodes.KSampler().sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise) - return refined_latent + return refined_latent[0] def upscale_shape(self, step_info, samples, w, h, save_temp_prefix=None): scale_method, model, vae, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise = self.params @@ -882,7 +1013,30 @@ class PixelKSampleUpscaler: # REQUIREMENTS: BlenderNeko/ComfyUI_TiledKSampler try: - class PixelTiledKSampleUpscaler: + class TiledKSamplerWrapper: + params = None + + def __init__(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, + tile_width, tile_height, concurrent_tiles): + self.params = model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, tile_width, tile_height, concurrent_tiles + + def sample(self, latent_image, hook): + from custom_nodes.ComfyUI_TiledKSampler.nodes import TiledKSamplerAdvanced + + model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, tile_width, tile_height, concurrent_tiles = self.params + + steps = int(steps/denoise) + start_at_step = int(steps*(1.0 - denoise)) + end_at_step = steps + + 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 TiledKSamplerAdvanced().sample(model, "enable", seed, tile_width, tile_height, concurrent_tiles, steps, cfg, sampler_name, scheduler, + positive, negative, latent_image, start_at_step, end_at_step, "disable")[0] + + class PixelTiledKSampleUpscaler: params = None upscale_model = None tile_params = None @@ -908,10 +1062,8 @@ try: end_at_step = steps #print(f"steps={steps}, start_at_step={start_at_step}, end_at_step={end_at_step}") - refined_latent = TiledKSamplerAdvanced().sample(model, "enable", seed, tile_width, tile_height, concurrent_tiles, steps, cfg, sampler_name, scheduler, - positive, negative, latent, start_at_step, end_at_step, "disable") - - return refined_latent + return TiledKSamplerAdvanced().sample(model, "enable", seed, tile_width, tile_height, concurrent_tiles, steps, cfg, sampler_name, scheduler, + positive, negative, latent, start_at_step, end_at_step, "disable")[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 @@ -945,7 +1097,7 @@ try: refined_latent = self.emulate_non_advanced(upscaled_latent) - return refined_latent[0] + return refined_latent except: pass diff --git a/impact_pack.py b/impact_pack.py index 4d1ac46..0aaf427 100644 --- a/impact_pack.py +++ b/impact_pack.py @@ -377,16 +377,15 @@ class FaceDetailer: def doit(self, image, model, vae, guide_size, guide_size_for, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, feather, noise_mask, force_inpaint, bbox_threshold, bbox_dilation, bbox_crop_factor, - sam_detection_hint, sam_dilation, sam_threshold, sam_bbox_expansion, sam_mask_hint_threshold, sam_mask_hint_use_negative, - bbox_detector, sam_model_opt=None): + sam_detection_hint, sam_dilation, sam_threshold, sam_bbox_expansion, sam_mask_hint_threshold, + sam_mask_hint_use_negative, bbox_detector, sam_model_opt=None): enhanced_img, cropped_enhanced, mask = FaceDetailer.enhance_face( image, model, vae, guide_size, guide_size_for, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, feather, noise_mask, force_inpaint, bbox_threshold, bbox_dilation, bbox_crop_factor, sam_detection_hint, sam_dilation, sam_threshold, sam_bbox_expansion, sam_mask_hint_threshold, - sam_mask_hint_use_negative, - bbox_detector, sam_model_opt) + sam_mask_hint_use_negative, bbox_detector, sam_model_opt) pipe = (model, vae, positive, negative, bbox_detector, sam_model_opt) return enhanced_img, cropped_enhanced, mask, pipe @@ -490,6 +489,35 @@ class PixelKSampleHookCombine: return (hook, ) +class TiledKSamplerProvider: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "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": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "tile_width": ("INT", {"default": 512, "min": 256, "max": MAX_RESOLUTION, "step": 64}), + "tile_height": ("INT", {"default": 512, "min": 256, "max": MAX_RESOLUTION, "step": 64}), + "concurrent_tiles": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1}), + "basic_pipe": ("BASIC_PIPE", ) + }} + + RETURN_TYPES = ("KSAMPLER",) + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Sampler" + + def doit(self, seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, concurrent_tiles, basic_pipe): + model, _, _, positive, negative = basic_pipe + sampler = core.TiledKSamplerWrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, + tile_width, tile_height, concurrent_tiles) + return (sampler, ) + + class PixelTiledKSampleUpscalerProvider: upscale_methods = ["nearest-exact", "bilinear", "area"] @@ -624,7 +652,7 @@ class PixelKSampleUpscalerProviderPipe(PixelKSampleUpscalerProvider): "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - "use_tiled_vae": (["enabled", "disabled"],), + "use_tiled_vae": (["disabled", "enabled"],), "basic_pipe": ("BASIC_PIPE",) }, "optional": { @@ -646,6 +674,89 @@ class PixelKSampleUpscalerProviderPipe(PixelKSampleUpscalerProvider): return (upscaler, ) +class TwoSamplersForMaskUpscalerProvider: + upscale_methods = ["nearest-exact", "bilinear", "area"] + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "scale_method": (s.upscale_methods,), + "full_sample_schedule": ( + ["none", "interleave1", "interleave2", "interleave3", + "last1", "last2", + "interleave1+last1", "interleave2+last1", "interleave3+last1", + ],), + "use_tiled_vae": (["disabled", "enabled"],), + "base_sampler": ("KSAMPLER", ), + "mask_sampler": ("KSAMPLER", ), + "mask": ("MASK", ), + "vae": ("VAE",), + }, + "optional": { + "full_sampler_opt": ("KSAMPLER",), + "upscale_model_opt": ("UPSCALE_MODEL", ), + "pk_hook_base_opt": ("PK_HOOK", ), + "pk_hook_mask_opt": ("PK_HOOK", ), + "pk_hook_full_opt": ("PK_HOOK", ), + } + } + + RETURN_TYPES = ("UPSCALER", ) + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Upscale" + + def doit(self, scale_method, full_sample_schedule, use_tiled_vae, base_sampler, mask_sampler, mask, vae, + full_sampler_opt=None, upscale_model_opt=None, + pk_hook_base_opt=None, pk_hook_mask_opt=None, pk_hook_full_opt=None): + upscaler = core.TwoSamplersForMaskUpscaler(scale_method, full_sample_schedule, use_tiled_vae == "enabled", + base_sampler, mask_sampler, mask, vae, full_sampler_opt, upscale_model_opt, + pk_hook_base_opt, pk_hook_mask_opt, pk_hook_full_opt) + return (upscaler, ) + + +class TwoSamplersForMaskUpscalerProviderPipe: + upscale_methods = ["nearest-exact", "bilinear", "area"] + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "scale_method": (s.upscale_methods,), + "full_sample_schedule": ( + ["none", "interleave1", "interleave2", "interleave3", + "last1", "last2", + "interleave1+last1", "interleave2+last1", "interleave3+last1", + ],), + "use_tiled_vae": (["disabled", "enabled"],), + "base_sampler": ("KSAMPLER", ), + "mask_sampler": ("KSAMPLER", ), + "mask": ("MASK", ), + "basic_pipe": ("BASIC_PIPE",), + }, + "optional": { + "full_sampler_opt": ("KSAMPLER",), + "upscale_model_opt": ("UPSCALE_MODEL", ), + "pk_hook_base_opt": ("PK_HOOK", ), + "pk_hook_mask_opt": ("PK_HOOK", ), + "pk_hook_full_opt": ("PK_HOOK", ), + } + } + + RETURN_TYPES = ("UPSCALER", ) + FUNCTION = "doit" + + CATEGORY = "ImpactPack/Upscale" + + def doit(self, scale_method, full_sample_schedule, use_tiled_vae, base_sampler, mask_sampler, mask, basic_pipe, + full_sampler_opt=None, upscale_model_opt=None, + pk_hook_base_opt=None, pk_hook_mask_opt=None, pk_hook_full_opt=None): + _, _, vae, _, _ = basic_pipe + upscaler = core.TwoSamplersForMaskUpscaler(scale_method, full_sample_schedule, use_tiled_vae == "enabled", + base_sampler, mask_sampler, mask, vae, full_sampler_opt, upscale_model_opt, + pk_hook_base_opt, pk_hook_mask_opt, pk_hook_full_opt) + return (upscaler, ) + + class IterativeLatentUpscale: @classmethod def INPUT_TYPES(s): @@ -678,14 +789,14 @@ class IterativeLatentUpscale: scale += upscale_factor_unit new_w = w*scale new_h = h*scale - print(f"IterativeLatentUpscale[{i+1}/{steps}]: {new_w}x{new_h} (scale:{scale:.2f}) ") + print(f"IterativeLatentUpscale[{i+1}/{steps}]: {new_w:.1f}x{new_h:.1f} (scale:{scale:.2f}) ") step_info = i, steps current_latent = upscaler.upscale_shape(step_info, current_latent, new_w, new_h, temp_prefix) if scale < upscale_factor: new_w = w*upscale_factor new_h = h*upscale_factor - print(f"IterativeLatentUpscale[Final]: {new_w}x{new_h} (scale:{upscale_factor:.2f}) ") + print(f"IterativeLatentUpscale[Final]: {new_w:.1f}x{new_h:.1f} (scale:{upscale_factor:.2f}) ") step_info = steps, steps current_latent = upscaler.upscale_shape(step_info, current_latent, new_w, new_h, temp_prefix) @@ -1070,57 +1181,34 @@ class SubtractMask: return (mask,) - import nodes - -class MaskPainter(nodes.PreviewImage): +class PreviewBridge(nodes.PreviewImage): @classmethod def INPUT_TYPES(s): - return {"required": {"images": ("IMAGE", ), }, - "hidden": { - "prompt": "PROMPT", - "extra_pnginfo": "EXTRA_PNGINFO", - }, - "optional": {"mask_image": ("IMAGE_PATH", ), }, + return {"required": {"images": ("IMAGE",), }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", }, } - - RETURN_TYPES = ("MASK", ) - - FUNCTION = "save_painted_images" + + RETURN_TYPES = ("IMAGE", "MASK", ) + + FUNCTION = "doit" CATEGORY = "ImpactPack/Util" - def load_mask(self, imagepath): - if imagepath['type'] == "temp": - input_dir = folder_paths.get_temp_directory() - else: - input_dir = folder_paths.get_input_directory() - - image_path = os.path.join(input_dir, imagepath['filename']) - - if os.path.exists(image_path): - i = Image.open(image_path) - - if 'A' in i.getbands(): - mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 - mask = 1. - torch.from_numpy(mask) - else: - mask = torch.zeros((8, 8), dtype=torch.float32, device="cpu") - else: - mask = torch.zeros((8, 8), dtype=torch.float32, device="cpu") - - return (mask, ) - - def save_painted_images(self, images, filename_prefix="impact-mask", - prompt=None, extra_pnginfo=None, mask_image=None): + def doit(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): res = self.save_images(images, filename_prefix, prompt, extra_pnginfo) - if mask_image is not None: - res['result'] = self.load_mask(mask_image) + item = res['ui']['images'][0] + + if not item['filename'].endswith(']'): + filepath = f"{item['filename']} [{item['type']}]" else: - mask = torch.zeros((8, 8), dtype=torch.float32, device="cpu") - res['result'] = (mask, ) + filepath = item['filename'] + + image, mask = nodes.LoadImage().load_image(filepath) + + res['result'] = (image, mask, ) return res @@ -1175,7 +1263,7 @@ class DetailerForEach: enhanced_pil = core.enhance_detail(cropped_image, model, vae, guide_size, guide_size_for, seg.bbox, seed, steps, cfg, sampler_name, scheduler, - positive, negative, denoise, cropped_mask, force_inpaint) + positive, negative, denoise, cropped_mask, force_inpaint == "enabled") if not (enhanced_pil is None): # don't latent composite-> converting to latent caused poor quality diff --git a/install.py b/install.py index 444ba85..cffe0e0 100644 --- a/install.py +++ b/install.py @@ -73,7 +73,7 @@ def ensure_mmdet_package(): subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-U', 'openmim']) subprocess.check_call([sys.executable, '-m', 'mim', 'install', 'mmcv==2.0.0']) subprocess.check_call([sys.executable, '-m', 'mim', 'install', 'mmdet==3.0.0']) - subprocess.check_call([sys.executable, '-m', 'mim', 'install', 'mmengine==0.7.2']) + subprocess.check_call([sys.executable, '-m', 'mim', 'install', 'mmengine==0.7.3']) def install(): diff --git a/legacy_nodes.py b/legacy_nodes.py index 114b3ac..b77d240 100644 --- a/legacy_nodes.py +++ b/legacy_nodes.py @@ -2,7 +2,8 @@ import folder_paths import impact_core as core from impact_utils import * from impact_core import SEG - +import nodes +import os class NO_BBOX_MODEL: pass @@ -200,4 +201,55 @@ class SegsMaskCombine: return torch.from_numpy(mask.astype(np.float32) / 255.0) def doit(self, segs, image): - return (SegsMaskCombine.combine(segs, image), ) \ No newline at end of file + return (SegsMaskCombine.combine(segs, image), ) + + +class MaskPainter(nodes.PreviewImage): + @classmethod + def INPUT_TYPES(s): + return {"required": {"images": ("IMAGE",), }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + "optional": {"mask_image": ("IMAGE_PATH",), }, + } + + RETURN_TYPES = ("MASK",) + + FUNCTION = "save_painted_images" + + CATEGORY = "ImpactPack/Legacy" + + def load_mask(self, imagepath): + if imagepath['type'] == "temp": + input_dir = folder_paths.get_temp_directory() + else: + input_dir = folder_paths.get_input_directory() + + image_path = os.path.join(input_dir, imagepath['filename']) + + if os.path.exists(image_path): + i = Image.open(image_path) + + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + else: + mask = torch.zeros((8, 8), dtype=torch.float32, device="cpu") + else: + mask = torch.zeros((8, 8), dtype=torch.float32, device="cpu") + + return (mask,) + + def save_painted_images(self, images, filename_prefix="impact-mask", + prompt=None, extra_pnginfo=None, mask_image=None): + res = self.save_images(images, filename_prefix, prompt, extra_pnginfo) + + if mask_image is not None: + res['result'] = self.load_mask(mask_image) + else: + mask = torch.zeros((8, 8), dtype=torch.float32, device="cpu") + res['result'] = (mask,) + + return res \ No newline at end of file