From 40a24f398e7e7b04107746f4193122f746a89dee Mon Sep 17 00:00:00 2001
From: Chris <chris.goringe@gmail.com>
Date: Thu, 9 Apr 2026 12:21:16 +1000
Subject: [PATCH] catch and fix? bad masks

---
 image_filter_nodes.py |  6 ++++--
 js/mask_utils.js      | 14 ++++++++------
 js/popup.js           | 17 ++++++++++++-----
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/image_filter_nodes.py b/image_filter_nodes.py
index 22f6a0a..476077a 100644
--- a/image_filter_nodes.py
+++ b/image_filter_nodes.py
@@ -281,10 +281,12 @@ class MaskImageFilter(io.ComfyNode, FilterNodeBase):
             data = response.masked_data.split(',',1)[-1]
             mask = mask_from_data(data)
 
-        if if_no_mask == 'cancel' and torch.all(mask==0): raise InterruptProcessingException()
-        if mask is None: mask = torch.zeros_like(image[...,0])  
+        if mask is None: mask = torch.zeros_like(image[...,0]) 
+        if if_no_mask == 'cancel' and torch.all(mask==0): raise InterruptProcessingException() 
 
         InOutStore.last_output = ( image.clone(), mask.clone(), *response.get_extras((extra1, extra2, extra3)) )
+        if (image.shape[0:3] != mask.shape[0:3]):
+            print(f"Mask shape {mask.shape} does not match image shape {image.shape}")
         return io.NodeOutput( *InOutStore.last_output )
 
     @classmethod
diff --git a/js/mask_utils.js b/js/mask_utils.js
index 999439d..f458d7b 100644
--- a/js/mask_utils.js
+++ b/js/mask_utils.js
@@ -19,15 +19,17 @@ export function hide_mask_editor() {
 }
 
 function get_mask_editor_cancel_button() {
-
     try {
         var button = document.getElementById("maskEditor_topBarCancelButton")
         if (button) return button
-        const buttons = Array.from(get_mask_editor_element()?.getElementsByTagName('button'))
-        button = buttons.find((b)=>(b.ariaLabel=='Cancel'))
-        if (button) return button
-        button = buttons.find((b)=>(b.innerText=='Cancel'))
-        if (button) return button
+        const buttonlist = get_mask_editor_element()?.getElementsByTagName('button')
+        if (buttonlist) {
+            const buttons = Array.from(buttonlist)
+            button = buttons.find((b)=>(b.ariaLabel=='Cancel'))
+            if (button) return button
+            button = buttons.find((b)=>(b.innerText=='Cancel'))
+            if (button) return button
+        }
         button = get_mask_editor_element()?.parentElement?.lastChild?.childNodes[2]
         if (button) return button
     } catch (e) {
diff --git a/js/popup.js b/js/popup.js
index 8e3823f..a8bdea4 100644
--- a/js/popup.js
+++ b/js/popup.js
@@ -396,14 +396,21 @@ class Popup extends HTMLElement {
         if (mask_editor_showing()) {
             setTimeout(this.wait_while_mask_editing.bind(this), 100)
         } else {
-            const masked_image = this.node.imgs?.[0]?.src ? this.extract_filename(this.node.imgs[0].src) : this.node.images?.[0]?.filename
-            this._send_response({masked_image:masked_image})
-
-            const the_node = this.node.id
-            setTimeout(remove_preview, 500, [the_node,])
+            setTimeout(this.when_mask_editor_closes.bind(this), 300) // allow a pause to make sure the mask editor has saved the image to the node
         } 
     }
 
+    when_mask_editor_closes() {
+        if (this.node.imgs?.[0]?.src) {
+            this._send_response({masked_image:this.extract_filename(this.node.imgs[0].src)})
+        } else {
+            this._send_response({masked_image:this.node.images?.[0]?.filename})
+        }
+
+        const the_node = this.node.id
+        setTimeout(remove_preview, 200, [the_node,])
+    }
+
     extract_filename(url_string) {
         return (new URL(url_string)).searchParams.get('filename')
     }