From 0e0722ec08c3b1c1637d9294e469bb5efbb8ac81 Mon Sep 17 00:00:00 2001 From: Alex Butler Date: Sat, 18 Jan 2025 18:16:56 +0000 Subject: [PATCH] Add tiled vae encoding/decoding toggle to detailer nodes (#883) * detailers: support optional tiled vae encoding & decoding * Remove start encoding/decoding logging --- modules/impact/core.py | 23 +++++++++------ modules/impact/impact_pack.py | 53 ++++++++++++++++++++++++----------- modules/impact/utils.py | 13 +++++++-- 3 files changed, 62 insertions(+), 27 deletions(-) diff --git a/modules/impact/core.py b/modules/impact/core.py index 654fc23..2f5e7d1 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -244,7 +244,8 @@ def enhance_detail(image, model, clip, vae, guide_size, guide_size_for_bbox, max detailer_hook=None, refiner_ratio=None, refiner_model=None, refiner_clip=None, refiner_positive=None, refiner_negative=None, control_net_wrapper=None, cycle=1, - inpaint_model=False, noise_mask_feather=0, scheduler_func=None): + inpaint_model=False, noise_mask_feather=0, scheduler_func=None, + vae_tiled_encode=False, vae_tiled_decode=False): if noise_mask is not None: noise_mask = utils.tensor_gaussian_blur_mask(noise_mask, noise_mask_feather) @@ -334,7 +335,7 @@ def enhance_detail(image, model, clip, vae, guide_size, guide_size_for_bbox, max print(f"[Impact Pack] ComfyUI is an outdated version.") positive, negative, latent_image = imc_encode(positive, negative, upscaled_image, vae, noise_mask) else: - latent_image = to_latent_image(upscaled_image, vae) + latent_image = to_latent_image(upscaled_image, vae, vae_tiled_encode=vae_tiled_encode) if noise_mask is not None: latent_image['noise_mask'] = noise_mask @@ -369,12 +370,18 @@ def enhance_detail(image, model, clip, vae, guide_size, guide_size_for_bbox, max refined_latent = detailer_hook.pre_decode(refined_latent) # non-latent downscale - latent downscale cause bad quality - try: - # try to decode image normally - refined_image = vae.decode(refined_latent['samples']) - except Exception as e: - #usually an out-of-memory exception from the decode, so try a tiled approach - refined_image = vae.decode_tiled(refined_latent["samples"], tile_x=64, tile_y=64, ) + start = time.time() + if vae_tiled_decode: + (refined_image,) = nodes.VAEDecodeTiled().decode(vae, refined_latent, 512) # using default settings + print(f"[Impact Pack] vae decoded (tiled) in {time.time() - start:.1f}s") + else: + try: + refined_image = vae.decode(refined_latent['samples']) + except Exception as e: + # usually an out-of-memory exception from the decode, so try a tiled approach + print(f"[Impact Pack] failed after {time.time() - start:.1f}s, doing vae.decode_tiled 64...") + refined_image = vae.decode_tiled(refined_latent["samples"], tile_x=64, tile_y=64, ) + print(f"[Impact Pack] vae decoded in {time.time() - start:.1f}s") if detailer_hook is not None: refined_image = detailer_hook.post_decode(refined_image) diff --git a/modules/impact/impact_pack.py b/modules/impact/impact_pack.py index b13ff6f..ef20cb1 100644 --- a/modules/impact/impact_pack.py +++ b/modules/impact/impact_pack.py @@ -218,6 +218,8 @@ class DetailerForEach: "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), "noise_mask_feather": ("INT", {"default": 20, "min": 0, "max": 100, "step": 1}), "scheduler_func_opt": ("SCHEDULER_FUNC",), + "tiled_encode": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + "tiled_decode": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), } } @@ -234,7 +236,7 @@ class DetailerForEach: def do_detail(image, segs, model, clip, vae, guide_size, guide_size_for_bbox, max_size, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, feather, noise_mask, force_inpaint, wildcard_opt=None, detailer_hook=None, refiner_ratio=None, refiner_model=None, refiner_clip=None, refiner_positive=None, refiner_negative=None, - cycle=1, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None): + cycle=1, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None, tiled_encode=False, tiled_decode=False): if len(image) > 1: raise Exception('[Impact Pack] ERROR: DetailerForEach does not allow image batches.\nPlease refer to https://github.com/ltdrdata/ComfyUI-extension-tutorials/blob/Main/ComfyUI-Impact-Pack/tutorial/batching-detailer.md for more information.') @@ -338,7 +340,8 @@ class DetailerForEach: refiner_clip=refiner_clip, refiner_positive=refiner_positive, refiner_negative=refiner_negative, control_net_wrapper=seg.control_net_wrapper, cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, - scheduler_func=scheduler_func_opt) + scheduler_func=scheduler_func_opt, vae_tiled_encode=tiled_encode, + vae_tiled_decode=tiled_decode) else: enhanced_image = cropped_image cnet_pils = None @@ -384,13 +387,15 @@ class DetailerForEach: def doit(self, image, segs, model, clip, vae, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, feather, noise_mask, force_inpaint, wildcard, cycle=1, - detailer_hook=None, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None): + detailer_hook=None, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None, + tiled_encode=False, tiled_decode=False): enhanced_img, *_ = \ DetailerForEach.do_detail(image, segs, model, clip, vae, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, feather, noise_mask, force_inpaint, wildcard, detailer_hook, - cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt) + cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, + scheduler_func_opt=scheduler_func_opt, tiled_encode=tiled_encode, tiled_decode=tiled_decode) return (enhanced_img, ) @@ -425,6 +430,8 @@ class DetailerForEachPipe: "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), "noise_mask_feather": ("INT", {"default": 20, "min": 0, "max": 100, "step": 1}), "scheduler_func_opt": ("SCHEDULER_FUNC",), + "tiled_encode": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + "tiled_decode": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), } } @@ -438,7 +445,8 @@ class DetailerForEachPipe: def doit(self, image, segs, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, denoise, feather, noise_mask, force_inpaint, basic_pipe, wildcard, refiner_ratio=None, detailer_hook=None, refiner_basic_pipe_opt=None, - cycle=1, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None): + cycle=1, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None, + tiled_encode=False, tiled_decode=False): if len(image) > 1: raise Exception('[Impact Pack] ERROR: DetailerForEach does not allow image batches.\nPlease refer to https://github.com/ltdrdata/ComfyUI-extension-tutorials/blob/Main/ComfyUI-Impact-Pack/tutorial/batching-detailer.md for more information.') @@ -456,7 +464,8 @@ class DetailerForEachPipe: force_inpaint, wildcard, detailer_hook, refiner_ratio=refiner_ratio, refiner_model=refiner_model, refiner_clip=refiner_clip, refiner_positive=refiner_positive, refiner_negative=refiner_negative, - cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt) + cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt, + tiled_encode=tiled_encode, tiled_decode=tiled_decode) # set fallback image if len(cnet_pil_list) == 0: @@ -513,6 +522,8 @@ class FaceDetailer: "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), "noise_mask_feather": ("INT", {"default": 20, "min": 0, "max": 100, "step": 1}), "scheduler_func_opt": ("SCHEDULER_FUNC",), + "tiled_encode": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + "tiled_decode": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), }} RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "MASK", "DETAILER_PIPE", "IMAGE") @@ -530,7 +541,7 @@ class FaceDetailer: sam_mask_hint_use_negative, drop_size, bbox_detector, segm_detector=None, sam_model_opt=None, wildcard_opt=None, detailer_hook=None, refiner_ratio=None, refiner_model=None, refiner_clip=None, refiner_positive=None, refiner_negative=None, cycle=1, - inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None): + inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None, tiled_encode=False, tiled_decode=False): # make default prompt as 'face' if empty prompt for CLIPSeg bbox_detector.setAux('face') @@ -562,7 +573,8 @@ class FaceDetailer: refiner_ratio=refiner_ratio, refiner_model=refiner_model, refiner_clip=refiner_clip, refiner_positive=refiner_positive, refiner_negative=refiner_negative, - cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt) + cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, + scheduler_func_opt=scheduler_func_opt, tiled_encode=tiled_encode, tiled_decode=tiled_decode) else: enhanced_img = image cropped_enhanced = [] @@ -588,7 +600,8 @@ class FaceDetailer: 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, drop_size, bbox_detector, wildcard, cycle=1, - sam_model_opt=None, segm_detector_opt=None, detailer_hook=None, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None): + sam_model_opt=None, segm_detector_opt=None, detailer_hook=None, inpaint_model=False, noise_mask_feather=0, + scheduler_func_opt=None, tiled_encode=False, tiled_decode=False): result_img = None result_mask = None @@ -606,7 +619,8 @@ class FaceDetailer: 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, drop_size, bbox_detector, segm_detector_opt, sam_model_opt, wildcard, detailer_hook, - cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt) + cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt, + tiled_encode=tiled_encode, tiled_decode=tiled_decode) result_img = torch.cat((result_img, enhanced_img), dim=0) if result_img is not None else enhanced_img result_mask = torch.cat((result_mask, mask), dim=0) if result_mask is not None else mask @@ -1381,6 +1395,8 @@ class FaceDetailerPipe: "inpaint_model": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), "noise_mask_feather": ("INT", {"default": 20, "min": 0, "max": 100, "step": 1}), "scheduler_func_opt": ("SCHEDULER_FUNC",), + "tiled_encode": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + "tiled_decode": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), } } @@ -1395,7 +1411,8 @@ class FaceDetailerPipe: 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, drop_size, refiner_ratio=None, - cycle=1, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None): + cycle=1, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None, + tiled_encode=False, tiled_decode=False): result_img = None result_mask = None @@ -1418,7 +1435,8 @@ class FaceDetailerPipe: sam_mask_hint_use_negative, drop_size, bbox_detector, segm_detector, sam_model_opt, wildcard, detailer_hook, refiner_ratio=refiner_ratio, refiner_model=refiner_model, refiner_clip=refiner_clip, refiner_positive=refiner_positive, refiner_negative=refiner_negative, - cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt) + cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt, + tiled_encode=tiled_encode, tiled_decode=tiled_decode) result_img = torch.cat((result_img, enhanced_img), dim=0) if result_img is not None else enhanced_img result_mask = torch.cat((result_mask, mask), dim=0) if result_mask is not None else mask @@ -1552,7 +1570,7 @@ class DetailerForEachTest(DetailerForEach): def doit(self, image, segs, model, clip, vae, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, feather, noise_mask, force_inpaint, wildcard, detailer_hook=None, - cycle=1, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None): + cycle=1, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None, tiled_encode=False, tiled_decode=False): if len(image) > 1: raise Exception('[Impact Pack] ERROR: DetailerForEach does not allow image batches.\nPlease refer to https://github.com/ltdrdata/ComfyUI-extension-tutorials/blob/Main/ComfyUI-Impact-Pack/tutorial/batching-detailer.md for more information.') @@ -1561,7 +1579,8 @@ class DetailerForEachTest(DetailerForEach): DetailerForEach.do_detail(image, segs, model, clip, vae, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, positive, negative, denoise, feather, noise_mask, force_inpaint, wildcard, detailer_hook, - cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt) + cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, + scheduler_func_opt=scheduler_func_opt, tiled_encode=tiled_encode, tiled_decode=tiled_decode) # set fallback image if len(cropped) == 0: @@ -1590,7 +1609,8 @@ class DetailerForEachTestPipe(DetailerForEachPipe): def doit(self, image, segs, guide_size, guide_size_for, max_size, seed, steps, cfg, sampler_name, scheduler, denoise, feather, noise_mask, force_inpaint, basic_pipe, wildcard, cycle=1, - refiner_ratio=None, detailer_hook=None, refiner_basic_pipe_opt=None, inpaint_model=False, noise_mask_feather=0, scheduler_func_opt=None): + refiner_ratio=None, detailer_hook=None, refiner_basic_pipe_opt=None, inpaint_model=False, noise_mask_feather=0, + scheduler_func_opt=None, tiled_encode=False, tiled_decode=False): if len(image) > 1: raise Exception('[Impact Pack] ERROR: DetailerForEach does not allow image batches.\nPlease refer to https://github.com/ltdrdata/ComfyUI-extension-tutorials/blob/Main/ComfyUI-Impact-Pack/tutorial/batching-detailer.md for more information.') @@ -1609,7 +1629,8 @@ class DetailerForEachTestPipe(DetailerForEachPipe): refiner_ratio=refiner_ratio, refiner_model=refiner_model, refiner_clip=refiner_clip, refiner_positive=refiner_positive, refiner_negative=refiner_negative, - cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, scheduler_func_opt=scheduler_func_opt) + cycle=cycle, inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather, + scheduler_func_opt=scheduler_func_opt, tiled_encode=tiled_encode, tiled_decode=tiled_decode) # set fallback image if len(cropped) == 0: diff --git a/modules/impact/utils.py b/modules/impact/utils.py index bd583cb..0e70fb5 100644 --- a/modules/impact/utils.py +++ b/modules/impact/utils.py @@ -7,6 +7,7 @@ import nodes from . import config from PIL import Image import comfy +import time class TensorBatchBuilder: @@ -501,15 +502,21 @@ def crop_image(image, crop_region): return crop_tensor4(image, crop_region) -def to_latent_image(pixels, vae): +def to_latent_image(pixels, vae, vae_tiled_encode=False): x = pixels.shape[1] y = pixels.shape[2] if pixels.shape[1] != x or pixels.shape[2] != y: pixels = pixels[:, :x, :y, :] - vae_encode = nodes.VAEEncode() + start = time.time() + if vae_tiled_encode: + encoded = nodes.VAEEncodeTiled().encode(vae, pixels, 512, overlap=64)[0] # using default settings + print(f"[Impact Pack] vae encoded (tiled) in {time.time() - start:.1f}s") + else: + encoded = nodes.VAEEncode().encode(vae, pixels)[0] + print(f"[Impact Pack] vae encoded in {time.time() - start:.1f}s") - return vae_encode.encode(vae, pixels)[0] + return encoded def empty_pil_tensor(w=64, h=64):