add nodes: TwoSamplersForMaskUpscalerProvider, TwoSamplersForMaskUpscalerProviderPipe, TiledKSamplerProvider

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.

__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)",

impact_config.py

@@ -1,7 +1,7 @@
 import configparser
 import os
 
-version = "V2.4"
+version = "V2.5"
 
 dependency_version = 1
 

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
 

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

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():

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), )
+        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