Spaces:
Paused
Paused
| 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: | |
| 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", "IMAGE") | |
| RETURN_NAMES = ("segs", "cnet_images") | |
| OUTPUT_IS_LIST = (False, True) | |
| FUNCTION = "doit" | |
| CATEGORY = "ImpactPack/Detailer" | |
| 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 = [] | |
| cnet_image_list = [] | |
| 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.cpu().numpy() | |
| enhanced_image_tensor, cnet_images = 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, control_net_wrapper=seg.control_net_wrapper, | |
| inpaint_model=inpaint_model, noise_mask_feather=noise_mask_feather) | |
| if cnet_images is not None: | |
| cnet_image_list.extend(cnet_images) | |
| if enhanced_image_tensor is None: | |
| new_cropped_image = cropped_image_frames | |
| else: | |
| new_cropped_image = enhanced_image_tensor.cpu().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), cnet_image_list | |
| 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, cnet_images = 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) | |
| if len(cnet_images) == 0: | |
| cnet_images = [empty_pil_tensor()] | |
| return (segs, cnet_images) | |
| class DetailerForEachPipeForAnimateDiff: | |
| 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", "IMAGE") | |
| RETURN_NAMES = ("image", "segs", "basic_pipe", "cnet_images") | |
| OUTPUT_IS_LIST = (False, False, False, True) | |
| FUNCTION = "doit" | |
| CATEGORY = "ImpactPack/Detailer" | |
| 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 = [] | |
| cnet_image_list = [] | |
| for sub_seg in segs[1]: | |
| single_seg = segs[0], [sub_seg] | |
| enhanced_seg, cnet_images = 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] | |
| if cnet_images is not None: | |
| cnet_image_list.extend(cnet_images) | |
| if detailer_hook is not None: | |
| detailer_hook.post_paste(image_frames) | |
| enhanced_segs += enhanced_seg[1] | |
| new_segs = segs[0], enhanced_segs | |
| return image_frames, new_segs, basic_pipe, cnet_image_list | |