From bcdb6fcc80d6d32d390fa821ea162240eadb2566 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Fri, 13 Oct 2023 14:18:25 +0900 Subject: [PATCH] refactor: PreviewBridge --- js/impact-pack.js | 2 +- modules/impact/core.py | 3 + modules/impact/impact_pack.py | 109 ++++++++++++++++---------------- modules/impact/impact_server.py | 10 +++ 4 files changed, 67 insertions(+), 57 deletions(-) diff --git a/js/impact-pack.js b/js/impact-pack.js index 1da4720..b0cc8bd 100644 --- a/js/impact-pack.js +++ b/js/impact-pack.js @@ -636,7 +636,7 @@ app.registerExtension({ populated_text_widget.serializeValue = force_serializeValue; } - if (node.comfyClass == "PreviewBridge" || node.comfyClass == "MaskPainter") { + if (node.comfyClass == "MaskPainter") { node.widgets[0].value = '#placeholder'; Object.defineProperty(node, "images", { diff --git a/modules/impact/core.py b/modules/impact/core.py index 7941b76..400f0dd 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -20,6 +20,9 @@ SEG = namedtuple("SEG", defaults=[None]) +preview_bridge_cache = {} + + def erosion_mask(mask, grow_mask_by): if len(mask.shape) == 3: mask = mask.squeeze(0) diff --git a/modules/impact/impact_pack.py b/modules/impact/impact_pack.py index fda58ba..b81a4dc 100644 --- a/modules/impact/impact_pack.py +++ b/modules/impact/impact_pack.py @@ -1405,81 +1405,78 @@ def get_image_hash(arr): return hash((sum1, sum2, sum3, sum4)) +def get_file_item(base_type, path): + path_type = base_type + + if path == "[output]": + path_type = "output" + path = path[:-9] + elif path == "[input]": + path_type = "input" + path = path[:-8] + elif path == "[temp]": + path_type = "temp" + path = path[:-7] + + subfolder = os.path.dirname(path) + filename = os.path.basename(path) + + return { + "filename": filename, + "subfolder": subfolder, + "type": path_type + } + + class PreviewBridge(nodes.PreviewImage): @classmethod def INPUT_TYPES(s): - return {"required": {"images": ("IMAGE",), }, - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", "unique_id": "UNIQUE_ID"}, - "optional": {"image": (["#placeholder"], )}, + return {"required": { + "images": ("IMAGE",), + "image": ("STRING", {"default": ""}), + }, + "hidden": {"unique_id": "UNIQUE_ID"}, } RETURN_TYPES = ("IMAGE", "MASK", ) FUNCTION = "doit" + OUTPUT_NODE = True + CATEGORY = "ImpactPack/Util" def __init__(self): super().__init__() self.prev_hash = None - def doit(self, images, image, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, unique_id=None): - if image != "#placeholder" and isinstance(image, str): - image_path = folder_paths.get_annotated_filepath(image) - img = Image.open(image_path).convert("RGB") - data = np.array(img) - image_hash = get_image_hash(data) - else: - data = (255. * images[0].cpu().numpy()).astype(int) - image_hash = get_image_hash(data) + def doit(self, images, image, unique_id): + if unique_id not in impact.core.preview_bridge_cache: + image = "" + elif impact.core.preview_bridge_cache[unique_id] is not images: + image = "" - is_changed = False - if self.prev_hash is None or self.prev_hash != image_hash: - self.prev_hash = image_hash - is_changed = True + if image != "": + try: + pixels, mask = nodes.LoadImage().load_image(image) + image = [get_file_item("temp", image)] + except: + image = "" - if is_changed or image == "#placeholder": - # new input image - res = self.save_images(images, filename_prefix, prompt, extra_pnginfo) + if image == "": + impact.core.preview_bridge_cache[unique_id] = images - item = res['ui']['images'][0] + res = nodes.PreviewImage().save_images(images, filename_prefix="PreviewBridge/PB-") + image = res['ui']['images'] + pixels = images + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + image_feedback = f"${image[0]['filename']} [temp]" + PromptServer.instance.send_sync("impact-node-feedback", {"id": unique_id, "widget_name": "image", "type": "text", "value": image_feedback}) - if not item['filename'].endswith(']'): - filepath = f"{item['filename']} [{item['type']}]" - else: - filepath = item['filename'] - - image, mask = nodes.LoadImage().load_image(filepath) - - res['ui']['aux'] = [image_hash, res['ui']['images']] - res['result'] = (image, mask, ) - - return res - - else: - # new mask - if '0' in image: # fallback - image = image['0'] - - forward = {'filename': image['forward_filename'], - 'subfolder': image['forward_subfolder'], - 'type': image['forward_type'], } - - res = {'ui': {'images': [forward]}} - - imgpath = "" - if 'subfolder' in image and image['subfolder'] != "": - imgpath = image['subfolder'] + "/" - - imgpath += f"{image['filename']}" - - if 'type' in image and image['type'] != "": - imgpath += f" [{image['type']}]" - - res['ui']['aux'] = [image_hash, [forward]] - res['result'] = nodes.LoadImage().load_image(imgpath) - - return res + return { + "ui": {"images": image}, + "result": (pixels, mask, ), + } class ImageReceiver(nodes.LoadImage): diff --git a/modules/impact/impact_server.py b/modules/impact/impact_server.py index d6457d5..2465222 100644 --- a/modules/impact/impact_server.py +++ b/modules/impact/impact_server.py @@ -295,9 +295,19 @@ def onprompt_for_pickers(json_data): del segs_picker_map[key] +def gc_preview_bridge_cache(json_data): + prompt_keys = json_data['prompt'].keys() + + for key in list(core.preview_bridge_cache.keys()): + if key not in prompt_keys: + print(f"key deleted: {key}") + del core.preview_bridge_cache[key] + + def onprompt(json_data): json_data = onprompt_for_switch(json_data) onprompt_for_pickers(json_data) + gc_preview_bridge_cache(json_data) return json_data