catch and fix? bad masks

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

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) {

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')
     }