From 49b151fc07faba830496151c70ba0bd8657b5d87 Mon Sep 17 00:00:00 2001 From: Chris Date: Wed, 1 Jan 2025 18:19:04 +1100 Subject: [PATCH] image mask --- __init__.py | 5 ++-- image_filter.py | 67 ++++++++++++++++++++++++++++++++++++++----------- js/filter.css | 1 + js/popup.js | 32 +++++++++++++++++++++-- 4 files changed, 86 insertions(+), 19 deletions(-) diff --git a/__init__.py b/__init__.py index ed404a9..136d8b3 100644 --- a/__init__.py +++ b/__init__.py @@ -5,14 +5,15 @@ @description: A custom node that pauses the flow while you choose which image or images to pass on to the rest of the workflow. Simplified and improved version of cg-image-picker. """ -from .image_filter import ImageFilter, TextImageFilter +from .image_filter import ImageFilter, TextImageFilter, MaskEditFilter VERSION = "0.1" WEB_DIRECTORY = "./js" NODE_CLASS_MAPPINGS= { "Image Filter": ImageFilter, - "Text Image Filter": TextImageFilter + "Text Image Filter": TextImageFilter, + "Mask Edit Filter": MaskEditFilter } __all__ = ["NODE_CLASS_MAPPINGS", "WEB_DIRECTORY"] diff --git a/image_filter.py b/image_filter.py index 6e4e8e1..574c3e6 100644 --- a/image_filter.py +++ b/image_filter.py @@ -1,8 +1,8 @@ from server import PromptServer from aiohttp import web -from nodes import PreviewImage +from nodes import PreviewImage, LoadImage from comfy.model_management import InterruptProcessingException -import time +import time, os import torch ## @@ -24,12 +24,20 @@ def wait(secs): end_time = time.monotonic() + secs while(time.monotonic() < end_time and Message.data is None): time.sleep(1) response = Message.data + if response is None: + PromptServer.instance.send_sync("cg-image-filter-images", {"timeout": True}) Message.data = None return response class Message: data:str = None +HIDDEN = { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + "uid":"UNIQUE_ID" + } + class ImageFilter(PreviewImage): RETURN_TYPES = ("IMAGE","LATENT","MASK") RETURN_NAMES = ("images","latents","masks") @@ -49,11 +57,7 @@ class ImageFilter(PreviewImage): "latents" : ("LATENT", {"tooltip": "Optional - if provided, will be output"}), "masks" : ("MASK", {"tooltip": "Optional - if provided, will be output"}), }, - "hidden": { - "prompt": "PROMPT", - "extra_pnginfo": "EXTRA_PNGINFO", - "uid":"UNIQUE_ID" - }, + "hidden": HIDDEN, } @classmethod @@ -67,7 +71,8 @@ class ImageFilter(PreviewImage): response = wait(timeout) or ('' if ontimeout=='send none' else ",".join(list(str(x) for x in range(len(images))))) images_to_return = list(int(x) for x in response.split(",") if x) - if len(images_to_return) == 0: raise InterruptProcessingException() + if len(images_to_return) == 0: + raise InterruptProcessingException() images = torch.stack(list(images[i] for i in images_to_return)) latents = {"samples": torch.stack(list(latents['samples'][i] for i in images_to_return))} if latents is not None else None @@ -90,11 +95,7 @@ class TextImageFilter(PreviewImage): "text" : ("STRING", {"forceInput":True, "default":""}), "timeout": ("INT", {"default": 60, "tooltip": "Timeout in seconds."}), }, - "hidden": { - "prompt": "PROMPT", - "extra_pnginfo": "EXTRA_PNGINFO", - "uid":"UNIQUE_ID" - }, + "hidden": HIDDEN, } @classmethod @@ -105,5 +106,41 @@ class TextImageFilter(PreviewImage): urls:list[str] = self.save_images(images=image, **kwargs)['ui']['images'] PromptServer.instance.send_sync("cg-image-filter-images", {"uid": uid, "urls":urls, "text":text}) - response = wait(timeout) or text - return (image, response) \ No newline at end of file + response = wait(timeout) + + + return (image, response) + +class MaskEditFilter(PreviewImage, LoadImage): + RETURN_TYPES = ("IMAGE","MASK") + RETURN_NAMES = ("image","mask") + FUNCTION = "func" + CATEGORY = "image_filter" + OUTPUT_NODE = False + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image" : ("IMAGE", ), + "timeout": ("INT", {"default": 600, "tooltip": "Timeout in seconds."}), + }, + "hidden": HIDDEN, + } + + @classmethod + def IS_CHANGED(cls, **kwargs): + return float("NaN") + + @classmethod + def VALIDATE_INPUTS(cls, **kwargs): return True + + def func(self, image, timeout, uid, **kwargs): + urls:list[str] = self.save_images(images=image, **kwargs)['ui']['images'] + PromptServer.instance.send_sync("cg-image-filter-images", {"uid": uid, "urls":urls, "maskedit":True}) + src = wait(timeout) + if src: + filename = src.split('=')[1].split('&')[0] + return self.load_image(os.path.join('clipspace', filename)+" [input]") + else: + return self.load_image(urls[0]['filename']+" [temp]") diff --git a/js/filter.css b/js/filter.css index 284115a..589424f 100644 --- a/js/filter.css +++ b/js/filter.css @@ -60,6 +60,7 @@ .cg_popup .text_edit { width: 100%; + height: 150px; } .cg_popup.hidden { diff --git a/js/popup.js b/js/popup.js index d49d82b..02798f5 100644 --- a/js/popup.js +++ b/js/popup.js @@ -1,4 +1,4 @@ -import { app } from "../../scripts/app.js"; +import { app, ComfyApp } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; import { create } from "./utils.js"; @@ -44,8 +44,36 @@ class Popup extends HTMLSpanElement { this.active = false } + respond_after_maskeditor() { + if (document.getElementById('maskEditor').style.display == 'none') { + this._send_response(this.node.imgs[0].src) + } else { + setTimeout(this.respond_after_maskeditor.bind(this), 100) + } + } + show(details) { const detail = details.detail + if (detail.timeout) { + this.hide() + return + } + if (detail.maskedit) { + this.hide() + this.node = app.graph._nodes_by_id[detail.uid] + //this.node.imgs = [...detail.urls] + //this.node.imgs.forEach((url, i)=>{ this.node.imgs[i].src = api.apiURL( `/view?filename=${encodeURIComponent(url.filename)}&type=${url.type}&subfolder=${url.subfolder}`) }) + this.node.imgs = [] + detail.urls.forEach((url, i)=>{ + this.node.imgs.push( new Image() ); + this.node.imgs[i].src = api.apiURL( `/view?filename=${encodeURIComponent(url.filename)}&type=${url.type}&subfolder=${url.subfolder}`) + }) + ComfyApp.copyToClipspace(this.node) + ComfyApp.clipspace_return_node = this.node + ComfyApp.open_maskeditor() + setTimeout(this.respond_after_maskeditor.bind(this), 1000) + return + } this.doing_text = (detail.text != null) this.n_images = detail.urls?.length @@ -76,7 +104,7 @@ class Popup extends HTMLSpanElement { } }) - if (detail.text) { this.text_edit = create('textarea', 'text_edit', this.grid, {"innerText":detail.text}) } + if (detail.text) { this.text_edit = create('textarea', 'text_edit', this.grid, {"innerHTML":detail.text}) } this.layout() this.classList.remove('hidden') }