From d8566b2f754f6ad9fb140881dcd890f36c2dcdbe Mon Sep 17 00:00:00 2001
From: Chris <chris.goringe@gmail.com>
Date: Mon, 30 Dec 2024 15:31:01 +1100
Subject: [PATCH] added title, latent and masks, and ontimeout options

---
 README.md          |  3 ---
 image_filter.py    | 21 +++++++++++++++------
 js/filter.css      |  9 +++++++++
 js/image_filter.js |  3 +++
 4 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 67ece64..2a1469e 100644
--- a/README.md
+++ b/README.md
@@ -31,8 +31,5 @@ If you do nothing, and the number of seconds specified in the timeout is exceede
 
 ## Future plans
 
-- allow a choice on timeout of 'cancel' or 'pass all'
-- add a title (showing the title of the node)
-- add optional Latent and Mask inputs and outputs so that latents or masks associated with the images can also be filtered
 - a node that takes a single image as input, and when executed, activates the mask editor, then outputs an image and the on-the-fly mask
 - a node that takes an image and a string, and allows you to edit the string before passing it on
\ No newline at end of file
diff --git a/image_filter.py b/image_filter.py
index ab6e8ea..95cfdc2 100644
--- a/image_filter.py
+++ b/image_filter.py
@@ -20,8 +20,8 @@ async def cg_image_filter_message(request):
     return web.json_response({})
 
 class ImageFilter(PreviewImage):
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("images",)
+    RETURN_TYPES = ("IMAGE","LATENT","MASK")
+    RETURN_NAMES = ("images","latents","masks")
     FUNCTION = "func"
     CATEGORY = "image_filter"
     OUTPUT_NODE = False
@@ -33,7 +33,12 @@ class ImageFilter(PreviewImage):
         return {
             "required": { 
                 "images" : ("IMAGE", ), 
-                "timeout": ("INT", {"default": 60, "tooltip": "Timeout in seconds before all images are passed on."}),
+                "timeout": ("INT", {"default": 60, "tooltip": "Timeout in seconds."}),
+                "ontimeout": (["send none", "send all"]),
+            },
+            "optional": {
+                "latents" : ("LATENT", {"tooltip": "Optional - if provided, will be output"}),
+                "masks" : ("MASK", {"tooltip": "Optional - if provided, will be output"}),
             },
             "hidden": {
                 "prompt": "PROMPT", 
@@ -46,18 +51,22 @@ class ImageFilter(PreviewImage):
     def IS_CHANGED(cls, **kwargs):
         return float("NaN")
     
-    def func(self, images, timeout, uid, **kwargs):
+    def func(self, images, timeout, ontimeout, uid, latents=None, masks=None, **kwargs):
         urls:list[str] = self.save_images(images=images, **kwargs)['ui']['images']
         PromptServer.instance.send_sync("cg-image-filter-images", {"uid": uid, "urls":urls})
 
         ImageFilter.data = None
         end_time = time.monotonic() + timeout
         while(time.monotonic() < end_time and ImageFilter.data is None): time.sleep(1)
-        response = ImageFilter.data or ''
+        response = ImageFilter.data or ('' if ontimeout=='send none' else ",".join(list(str(x) for x in range(len(images)))))
         ImageFilter.data = None
 
         images_to_return = list(int(x) for x in response.split(",") if x)
 
         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
+        masks = torch.stack(list(masks[i] for i in images_to_return)) if masks is not None else None
                 
-        return (torch.stack(list(images[i] for i in images_to_return)), )
\ No newline at end of file
+        return (images, latents, masks)
\ No newline at end of file
diff --git a/js/filter.css b/js/filter.css
index ac17870..48d7501 100644
--- a/js/filter.css
+++ b/js/filter.css
@@ -12,6 +12,15 @@
     left:5%;
     height:90%;
     top:5%;
+    display: grid;
+    align-content: space-evenly;
+}
+
+.cg_popup .title {
+    position: absolute;
+    left: 40%;
+    top: 2%;
+    font-size: larger;
 }
 
 .cg_popup .buttons {
diff --git a/js/image_filter.js b/js/image_filter.js
index 64fbf0d..fe977d1 100644
--- a/js/image_filter.js
+++ b/js/image_filter.js
@@ -35,10 +35,12 @@ var click_sends = null
 var send_button = null
 var cancel_button = null
 var active = false
+var title_bar = null
 
 function create_popup() {
     popup = create('span', 'cg_popup hidden', document.body)
     const grid = create('span', 'grid', popup)
+    title_bar = create('span', 'title', popup)
     const buttons = create('span', 'buttons', popup)
     click_sends = create('input', 'control', buttons, {type:"checkbox", id:"click_sends"})
     create('label', 'control_text', buttons, {for:"click_sends", innerText:"click to send"})
@@ -112,6 +114,7 @@ function receive_images(details) {
     if (app.graph._nodes_by_id[detail.uid].type=="Image Filter") {
         active = true
         if (!popup) create_popup()
+        title_bar.innerText = app.graph._nodes_by_id[detail.uid].title ?? "Image Filter"
 
         picked.clear()
         if (detail.urls.length==1) picked.add('0')