Spaces:
Running
Running
| import copy | |
| import torch | |
| import nodes | |
| from impact import utils | |
| from . import segs_nodes | |
| from thirdparty import noise_nodes | |
| from server import PromptServer | |
| import asyncio | |
| import folder_paths | |
| import os | |
| from comfy_extras import nodes_custom_sampler | |
| import math | |
| class PixelKSampleHook: | |
| cur_step = 0 | |
| total_step = 0 | |
| def __init__(self): | |
| pass | |
| def set_steps(self, info): | |
| self.cur_step, self.total_step = info | |
| def post_decode(self, pixels): | |
| return pixels | |
| def post_upscale(self, pixels): | |
| return pixels | |
| def post_encode(self, samples): | |
| return samples | |
| def pre_decode(self, samples): | |
| return samples | |
| def pre_ksample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, | |
| denoise): | |
| return model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise | |
| def post_crop_region(self, w, h, item_bbox, crop_region): | |
| return crop_region | |
| def touch_scaled_size(self, w, h): | |
| return w, h | |
| class PixelKSampleHookCombine(PixelKSampleHook): | |
| hook1 = None | |
| hook2 = None | |
| def __init__(self, hook1, hook2): | |
| super().__init__() | |
| self.hook1 = hook1 | |
| self.hook2 = hook2 | |
| def set_steps(self, info): | |
| self.hook1.set_steps(info) | |
| self.hook2.set_steps(info) | |
| def pre_decode(self, samples): | |
| return self.hook2.pre_decode(self.hook1.pre_decode(samples)) | |
| def post_decode(self, pixels): | |
| return self.hook2.post_decode(self.hook1.post_decode(pixels)) | |
| def post_upscale(self, pixels): | |
| return self.hook2.post_upscale(self.hook1.post_upscale(pixels)) | |
| def post_encode(self, samples): | |
| return self.hook2.post_encode(self.hook1.post_encode(samples)) | |
| def post_crop_region(self, w, h, item_bbox, crop_region): | |
| crop_region = self.hook1.post_crop_region(w, h, item_bbox, crop_region) | |
| return self.hook2.post_crop_region(w, h, item_bbox, crop_region) | |
| def touch_scaled_size(self, w, h): | |
| w, h = self.hook1.touch_scaled_size(w, h) | |
| return self.hook2.touch_scaled_size(w, h) | |
| def pre_ksample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, | |
| denoise): | |
| model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise = \ | |
| self.hook1.pre_ksample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, | |
| upscaled_latent, denoise) | |
| return self.hook2.pre_ksample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, | |
| upscaled_latent, denoise) | |
| class DetailerHookCombine(PixelKSampleHookCombine): | |
| def cycle_latent(self, latent): | |
| latent = self.hook1.cycle_latent(latent) | |
| latent = self.hook2.cycle_latent(latent) | |
| return latent | |
| def post_detection(self, segs): | |
| segs = self.hook1.post_detection(segs) | |
| segs = self.hook2.post_detection(segs) | |
| return segs | |
| def post_paste(self, image): | |
| image = self.hook1.post_paste(image) | |
| image = self.hook2.post_paste(image) | |
| return image | |
| def get_custom_noise(self, seed, noise, is_touched): | |
| noise_1st, is_touched = self.hook1.get_custom_noise(seed, noise, is_touched) | |
| noise_2nd, is_touched = self.hook2.get_custom_noise(seed, noise, is_touched) | |
| return noise, is_touched | |
| class SimpleCfgScheduleHook(PixelKSampleHook): | |
| target_cfg = 0 | |
| def __init__(self, target_cfg): | |
| super().__init__() | |
| self.target_cfg = target_cfg | |
| def pre_ksample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise): | |
| if self.total_step > 1: | |
| progress = self.cur_step / (self.total_step - 1) | |
| gap = self.target_cfg - cfg | |
| current_cfg = int(cfg + gap * progress) | |
| else: | |
| current_cfg = self.target_cfg | |
| return model, seed, steps, current_cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise | |
| class SimpleDenoiseScheduleHook(PixelKSampleHook): | |
| def __init__(self, target_denoise): | |
| super().__init__() | |
| self.target_denoise = target_denoise | |
| def pre_ksample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise): | |
| if self.total_step > 1: | |
| progress = self.cur_step / (self.total_step - 1) | |
| gap = self.target_denoise - denoise | |
| current_denoise = denoise + gap * progress | |
| else: | |
| current_denoise = self.target_denoise | |
| return model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, current_denoise | |
| class SimpleStepsScheduleHook(PixelKSampleHook): | |
| def __init__(self, target_steps): | |
| super().__init__() | |
| self.target_steps = target_steps | |
| def pre_ksample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise): | |
| if self.total_step > 1: | |
| progress = self.cur_step / (self.total_step - 1) | |
| gap = self.target_steps - steps | |
| current_steps = int(steps + gap * progress) | |
| else: | |
| current_steps = self.target_steps | |
| return model, seed, current_steps, cfg, sampler_name, scheduler, positive, negative, upscaled_latent, denoise | |
| class DetailerHook(PixelKSampleHook): | |
| def cycle_latent(self, latent): | |
| return latent | |
| def post_detection(self, segs): | |
| return segs | |
| def post_paste(self, image): | |
| return image | |
| def get_custom_noise(self, seed, noise, is_touched): | |
| return noise, is_touched | |
| # class CustomNoiseDetailerHookProvider(DetailerHook): | |
| # def __init__(self, noise): | |
| # super().__init__() | |
| # self.noise = noise | |
| # | |
| # def get_custom_noise(self, seed, noise, is_start): | |
| # return self.noise | |
| class VariationNoiseDetailerHookProvider(DetailerHook): | |
| def __init__(self, variation_seed, variation_strength): | |
| super().__init__() | |
| self.variation_seed = variation_seed | |
| self.variation_strength = variation_strength | |
| def get_custom_noise(self, seed, noise, is_touched): | |
| empty_noise = {'samples': torch.zeros(noise.size())} | |
| if not is_touched: | |
| noise = nodes_custom_sampler.Noise_RandomNoise(seed).generate_noise(empty_noise) | |
| noise_2nd = nodes_custom_sampler.Noise_RandomNoise(self.variation_seed).generate_noise(empty_noise) | |
| mixed_noise = ((1 - self.variation_strength) * noise + self.variation_strength * noise_2nd) | |
| # NOTE: Since the variance of the Gaussian noise in mixed_noise has changed, it must be corrected through scaling. | |
| scale_factor = math.sqrt((1 - self.variation_strength) ** 2 + self.variation_strength ** 2) | |
| corrected_noise = mixed_noise / scale_factor # Scale the noise to maintain variance of 1 | |
| return corrected_noise, True | |
| class SimpleDetailerDenoiseSchedulerHook(DetailerHook): | |
| def __init__(self, target_denoise): | |
| super().__init__() | |
| self.target_denoise = target_denoise | |
| def pre_ksample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise): | |
| if self.total_step > 1: | |
| progress = self.cur_step / (self.total_step - 1) | |
| gap = self.target_denoise - denoise | |
| current_denoise = denoise + gap * progress | |
| else: | |
| # ignore hook if total cycle <= 1 | |
| current_denoise = denoise | |
| return model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, current_denoise | |
| class CoreMLHook(DetailerHook): | |
| def __init__(self, mode): | |
| super().__init__() | |
| resolution = mode.split('x') | |
| self.w = int(resolution[0]) | |
| self.h = int(resolution[1]) | |
| self.override_bbox_by_segm = False | |
| def pre_decode(self, samples): | |
| new_samples = copy.deepcopy(samples) | |
| new_samples['samples'] = samples['samples'][0].unsqueeze(0) | |
| return new_samples | |
| def post_encode(self, samples): | |
| new_samples = copy.deepcopy(samples) | |
| new_samples['samples'] = samples['samples'].repeat(2, 1, 1, 1) | |
| return new_samples | |
| def post_crop_region(self, w, h, item_bbox, crop_region): | |
| x1, y1, x2, y2 = crop_region | |
| bx1, by1, bx2, by2 = item_bbox | |
| crop_w = x2-x1 | |
| crop_h = y2-y1 | |
| crop_ratio = crop_w/crop_h | |
| target_ratio = self.w/self.h | |
| if crop_ratio < target_ratio: | |
| # shrink height | |
| top_gap = by1 - y1 | |
| bottom_gap = y2 - by2 | |
| gap_ratio = top_gap / bottom_gap | |
| target_height = 1/target_ratio*crop_w | |
| delta_height = crop_h - target_height | |
| new_y1 = int(y1 + delta_height*gap_ratio) | |
| new_y2 = int(new_y1 + target_height) | |
| crop_region = x1, new_y1, x2, new_y2 | |
| elif crop_ratio > target_ratio: | |
| # shrink width | |
| left_gap = bx1 - x1 | |
| right_gap = x2 - bx2 | |
| gap_ratio = left_gap / right_gap | |
| target_width = target_ratio*crop_h | |
| delta_width = crop_w - target_width | |
| new_x1 = int(x1 + delta_width*gap_ratio) | |
| new_x2 = int(new_x1 + target_width) | |
| crop_region = new_x1, y1, new_x2, y2 | |
| return crop_region | |
| def touch_scaled_size(self, w, h): | |
| return self.w, self.h | |
| # REQUIREMENTS: BlenderNeko/ComfyUI Noise | |
| class InjectNoiseHook(PixelKSampleHook): | |
| def __init__(self, source, seed, start_strength, end_strength): | |
| super().__init__() | |
| self.source = source | |
| self.seed = seed | |
| self.start_strength = start_strength | |
| self.end_strength = end_strength | |
| def post_encode(self, samples): | |
| cur_step = self.cur_step | |
| size = samples['samples'].shape | |
| seed = cur_step + self.seed + cur_step | |
| if "BNK_NoisyLatentImage" in nodes.NODE_CLASS_MAPPINGS and "BNK_InjectNoise" in nodes.NODE_CLASS_MAPPINGS: | |
| NoisyLatentImage = nodes.NODE_CLASS_MAPPINGS["BNK_NoisyLatentImage"] | |
| InjectNoise = nodes.NODE_CLASS_MAPPINGS["BNK_InjectNoise"] | |
| else: | |
| utils.try_install_custom_node('https://github.com/BlenderNeko/ComfyUI_Noise', | |
| "To use 'NoiseInjectionHookProvider', 'ComfyUI Noise' extension is required.") | |
| raise Exception("'BNK_NoisyLatentImage', 'BNK_InjectNoise' nodes are not installed.") | |
| noise = NoisyLatentImage().create_noisy_latents(self.source, seed, size[3] * 8, size[2] * 8, size[0])[0] | |
| # inj noise | |
| mask = None | |
| if 'noise_mask' in samples: | |
| mask = samples['noise_mask'] | |
| strength = self.start_strength + (self.end_strength - self.start_strength) * cur_step / self.total_step | |
| samples = InjectNoise().inject_noise(samples, strength, noise, mask)[0] | |
| print(f"[Impact Pack] InjectNoiseHook: strength = {strength}") | |
| if mask is not None: | |
| samples['noise_mask'] = mask | |
| return samples | |
| class UnsamplerHook(PixelKSampleHook): | |
| def __init__(self, model, steps, start_end_at_step, end_end_at_step, cfg, sampler_name, | |
| scheduler, normalize, positive, negative): | |
| super().__init__() | |
| self.model = model | |
| self.cfg = cfg | |
| self.sampler_name = sampler_name | |
| self.steps = steps | |
| self.start_end_at_step = start_end_at_step | |
| self.end_end_at_step = end_end_at_step | |
| self.scheduler = scheduler | |
| self.normalize = normalize | |
| self.positive = positive | |
| self.negative = negative | |
| def post_encode(self, samples): | |
| cur_step = self.cur_step | |
| Unsampler = noise_nodes.Unsampler | |
| end_at_step = self.start_end_at_step + (self.end_end_at_step - self.start_end_at_step) * cur_step / self.total_step | |
| end_at_step = int(end_at_step) | |
| print(f"[Impact Pack] UnsamplerHook: end_at_step = {end_at_step}") | |
| # inj noise | |
| mask = None | |
| if 'noise_mask' in samples: | |
| mask = samples['noise_mask'] | |
| samples = Unsampler().unsampler(self.model, self.cfg, self.sampler_name, self.steps, end_at_step, | |
| self.scheduler, self.normalize, self.positive, self.negative, samples)[0] | |
| if mask is not None: | |
| samples['noise_mask'] = mask | |
| return samples | |
| class InjectNoiseHookForDetailer(DetailerHook): | |
| def __init__(self, source, seed, start_strength, end_strength, from_start=False): | |
| super().__init__() | |
| self.source = source | |
| self.seed = seed | |
| self.start_strength = start_strength | |
| self.end_strength = end_strength | |
| self.from_start = from_start | |
| def inject_noise(self, samples): | |
| cur_step = self.cur_step if self.from_start else self.cur_step - 1 | |
| total_step = self.total_step if self.from_start else self.total_step - 1 | |
| size = samples['samples'].shape | |
| seed = cur_step + self.seed + cur_step | |
| if "BNK_NoisyLatentImage" in nodes.NODE_CLASS_MAPPINGS and "BNK_InjectNoise" in nodes.NODE_CLASS_MAPPINGS: | |
| NoisyLatentImage = nodes.NODE_CLASS_MAPPINGS["BNK_NoisyLatentImage"] | |
| InjectNoise = nodes.NODE_CLASS_MAPPINGS["BNK_InjectNoise"] | |
| else: | |
| utils.try_install_custom_node('https://github.com/BlenderNeko/ComfyUI_Noise', | |
| "To use 'NoiseInjectionDetailerHookProvider', 'ComfyUI Noise' extension is required.") | |
| raise Exception("'BNK_NoisyLatentImage', 'BNK_InjectNoise' nodes are not installed.") | |
| noise = NoisyLatentImage().create_noisy_latents(self.source, seed, size[3] * 8, size[2] * 8, size[0])[0] | |
| # inj noise | |
| mask = None | |
| if 'noise_mask' in samples: | |
| mask = samples['noise_mask'] | |
| strength = self.start_strength + (self.end_strength - self.start_strength) * cur_step / total_step | |
| samples = InjectNoise().inject_noise(samples, strength, noise, mask)[0] | |
| if mask is not None: | |
| samples['noise_mask'] = mask | |
| return samples | |
| def cycle_latent(self, latent): | |
| if self.cur_step == 0 and not self.from_start: | |
| return latent | |
| else: | |
| return self.inject_noise(latent) | |
| class UnsamplerDetailerHook(DetailerHook): | |
| def __init__(self, model, steps, start_end_at_step, end_end_at_step, cfg, sampler_name, | |
| scheduler, normalize, positive, negative, from_start=False): | |
| super().__init__() | |
| self.model = model | |
| self.cfg = cfg | |
| self.sampler_name = sampler_name | |
| self.steps = steps | |
| self.start_end_at_step = start_end_at_step | |
| self.end_end_at_step = end_end_at_step | |
| self.scheduler = scheduler | |
| self.normalize = normalize | |
| self.positive = positive | |
| self.negative = negative | |
| self.from_start = from_start | |
| def unsample(self, samples): | |
| cur_step = self.cur_step if self.from_start else self.cur_step - 1 | |
| total_step = self.total_step if self.from_start else self.total_step - 1 | |
| Unsampler = noise_nodes.Unsampler | |
| end_at_step = self.start_end_at_step + (self.end_end_at_step - self.start_end_at_step) * cur_step / total_step | |
| end_at_step = int(end_at_step) | |
| # inj noise | |
| mask = None | |
| if 'noise_mask' in samples: | |
| mask = samples['noise_mask'] | |
| samples = Unsampler().unsampler(self.model, self.cfg, self.sampler_name, self.steps, end_at_step, | |
| self.scheduler, self.normalize, self.positive, self.negative, samples)[0] | |
| if mask is not None: | |
| samples['noise_mask'] = mask | |
| return samples | |
| def cycle_latent(self, latent): | |
| if self.cur_step == 0 and not self.from_start: | |
| return latent | |
| else: | |
| return self.unsample(latent) | |
| class SEGSOrderedFilterDetailerHook(DetailerHook): | |
| def __init__(self, target, order, take_start, take_count): | |
| super().__init__() | |
| self.target = target | |
| self.order = order | |
| self.take_start = take_start | |
| self.take_count = take_count | |
| def post_detection(self, segs): | |
| return segs_nodes.SEGSOrderedFilter().doit(segs, self.target, self.order, self.take_start, self.take_count)[0] | |
| class SEGSRangeFilterDetailerHook(DetailerHook): | |
| def __init__(self, target, mode, min_value, max_value): | |
| super().__init__() | |
| self.target = target | |
| self.mode = mode | |
| self.min_value = min_value | |
| self.max_value = max_value | |
| def post_detection(self, segs): | |
| return segs_nodes.SEGSRangeFilter().doit(segs, self.target, self.mode, self.min_value, self.max_value)[0] | |
| class SEGSLabelFilterDetailerHook(DetailerHook): | |
| def __init__(self, labels): | |
| super().__init__() | |
| self.labels = labels | |
| def post_detection(self, segs): | |
| return segs_nodes.SEGSLabelFilter().doit(segs, "", self.labels)[0] | |
| class PreviewDetailerHook(DetailerHook): | |
| def __init__(self, node_id, quality): | |
| super().__init__() | |
| self.node_id = node_id | |
| self.quality = quality | |
| async def send(self, image): | |
| if len(image) > 0: | |
| image = image[0].unsqueeze(0) | |
| img = utils.tensor2pil(image) | |
| temp_path = os.path.join(folder_paths.get_temp_directory(), 'pvhook') | |
| if not os.path.exists(temp_path): | |
| os.makedirs(temp_path) | |
| fullpath = os.path.join(temp_path, f"{self.node_id}.webp") | |
| img.save(fullpath, quality=self.quality) | |
| item = { | |
| "filename": f"{self.node_id}.webp", | |
| "subfolder": 'pvhook', | |
| "type": 'temp' | |
| } | |
| PromptServer.instance.send_sync("impact-preview", {'node_id': self.node_id, 'item': item}) | |
| def post_paste(self, image): | |
| asyncio.run(self.send(image)) | |
| return image | |