Fix: PreviewBridge mask editor KeyError with clipspace files + restore fresh image behavior (#1009)
* fix: Handle clipspace files in PreviewBridge mask editor and restore fresh image behavior - Fix KeyError when using mask editor multiple times - Add register_clipspace_image() method to detect and register clipspace files - Resolve timing issue between frontend JS conversion and backend processing - Restore original "fresh start" behavior for new image generations - Clear mask cache when images/latents change to ensure clean slate for new content - Maintain backwards compatibility with existing preview bridge functionality - Add dual registration for clipspace paths and preview IDs Fixes issue where second+ mask saves would fail with: KeyError: 'clipspace/clipspace-mask-XXXXX.png [input]' Also fixed regression where new images would retain previous masks instead of starting fresh. * Fix restore_mask 'always' and 'if_same_size' modes in PreviewBridge - Fixed 'always' mode to properly preserve masks when changing input images/latents - Fixed 'if_same_size' mode to preserve masks when switching between same-sized images - Modified cache clearing logic to preserve cache for both 'always' and 'if_same_size' modes - Updated restoration logic to handle size comparison correctly for 'if_same_size' This addresses the issue where masks were being cleared even when restore_mask was set to 'always' or when switching between same-sized images with 'if_same_size'. Fixes the issue discussed in PR #1009
This commit is contained in:
Dustin
+111
-15
@@ -52,10 +52,10 @@ class PreviewBridge:
|
||||
if pb_id not in core.preview_bridge_image_id_map:
|
||||
is_fail = True
|
||||
|
||||
image_path, ui_item = core.preview_bridge_image_id_map[pb_id]
|
||||
|
||||
if not os.path.isfile(image_path):
|
||||
is_fail = True
|
||||
if not is_fail:
|
||||
image_path, ui_item = core.preview_bridge_image_id_map[pb_id]
|
||||
if not os.path.isfile(image_path):
|
||||
is_fail = True
|
||||
|
||||
if not is_fail:
|
||||
i = Image.open(image_path)
|
||||
@@ -80,23 +80,93 @@ class PreviewBridge:
|
||||
|
||||
return image, mask.unsqueeze(0), ui_item
|
||||
|
||||
@staticmethod
|
||||
def register_clipspace_image(clipspace_path, node_id):
|
||||
"""Register a clipspace image file in the preview bridge system.
|
||||
|
||||
This handles the case where ComfyUI's mask editor creates clipspace files
|
||||
that need to be integrated with the preview bridge system.
|
||||
"""
|
||||
# Remove [input] suffix if present
|
||||
clean_path = clipspace_path.replace(" [input]", "").replace("[input]", "")
|
||||
|
||||
# Try to find the actual clipspace file
|
||||
input_dir = folder_paths.get_input_directory()
|
||||
potential_paths = [
|
||||
clean_path,
|
||||
os.path.join(input_dir, clean_path),
|
||||
os.path.join(input_dir, "clipspace", os.path.basename(clean_path)),
|
||||
os.path.abspath(clean_path),
|
||||
]
|
||||
|
||||
actual_file = None
|
||||
for path in potential_paths:
|
||||
if os.path.isfile(path):
|
||||
actual_file = path
|
||||
break
|
||||
|
||||
if not actual_file:
|
||||
return False
|
||||
|
||||
# Create ui_item for the clipspace file
|
||||
ui_item = {
|
||||
'filename': os.path.basename(actual_file),
|
||||
'subfolder': 'clipspace',
|
||||
'type': 'input'
|
||||
}
|
||||
|
||||
# Register it using the preview bridge system
|
||||
core.set_previewbridge_image(node_id, actual_file, ui_item)
|
||||
# Also register under the original clipspace path for compatibility
|
||||
core.preview_bridge_image_id_map[clipspace_path] = (actual_file, ui_item)
|
||||
|
||||
return True
|
||||
|
||||
def doit(self, images, image, unique_id, block=False, restore_mask="never", prompt=None, extra_pnginfo=None):
|
||||
need_refresh = False
|
||||
images_changed = False
|
||||
|
||||
# Check if images have changed (this determines if we start fresh)
|
||||
if unique_id not in core.preview_bridge_cache:
|
||||
need_refresh = True
|
||||
|
||||
images_changed = True
|
||||
elif core.preview_bridge_cache[unique_id][0] is not images:
|
||||
need_refresh = True
|
||||
images_changed = True
|
||||
|
||||
# If images changed, clear the mask cache to ensure fresh start behavior
|
||||
# This restores the original behavior where new images start with empty masks
|
||||
# unless restore_mask is set to "always" or "if_same_size"
|
||||
if images_changed and restore_mask not in ["always", "if_same_size"] and unique_id in core.preview_bridge_last_mask_cache:
|
||||
del core.preview_bridge_last_mask_cache[unique_id]
|
||||
|
||||
# Handle clipspace files that aren't registered in the preview bridge system
|
||||
# This only applies when images haven't changed (same image, new mask scenario)
|
||||
if not need_refresh and image not in core.preview_bridge_image_id_map:
|
||||
# Check if this is a clipspace file that needs to be registered
|
||||
is_clipspace = image and ("clipspace" in image.lower() or "[input]" in image)
|
||||
if is_clipspace:
|
||||
if not PreviewBridge.register_clipspace_image(image, unique_id):
|
||||
need_refresh = True
|
||||
else:
|
||||
need_refresh = True
|
||||
|
||||
if not need_refresh:
|
||||
pixels, mask, path_item = PreviewBridge.load_image(image)
|
||||
image = [path_item]
|
||||
else:
|
||||
if restore_mask != "never":
|
||||
# For new images (images_changed=True), we want to start fresh regardless of restore_mask
|
||||
# For same image with refresh needed, respect the restore_mask setting
|
||||
# Exception: when restore_mask is "always", restore even with new images
|
||||
# Exception: when restore_mask is "if_same_size", allow restoration to check size compatibility
|
||||
if restore_mask != "never" and (not images_changed or restore_mask in ["always", "if_same_size"]):
|
||||
mask = core.preview_bridge_last_mask_cache.get(unique_id)
|
||||
if mask is None or (restore_mask != "always" and mask.shape[1:] != images.shape[1:3]):
|
||||
if mask is None:
|
||||
mask = None
|
||||
elif restore_mask == "if_same_size" and mask.shape[1:] != images.shape[1:3]:
|
||||
# For if_same_size, clear mask if dimensions don't match
|
||||
mask = None
|
||||
# For "always", keep the mask regardless of size
|
||||
else:
|
||||
mask = None
|
||||
|
||||
@@ -252,10 +322,10 @@ class PreviewBridgeLatent:
|
||||
if pb_id not in core.preview_bridge_image_id_map:
|
||||
is_fail = True
|
||||
|
||||
image_path, ui_item = core.preview_bridge_image_id_map[pb_id]
|
||||
|
||||
if not os.path.isfile(image_path):
|
||||
is_fail = True
|
||||
if not is_fail:
|
||||
image_path, ui_item = core.preview_bridge_image_id_map[pb_id]
|
||||
if not os.path.isfile(image_path):
|
||||
is_fail = True
|
||||
|
||||
if not is_fail:
|
||||
i = Image.open(image_path)
|
||||
@@ -295,15 +365,33 @@ class PreviewBridgeLatent:
|
||||
raise Exception("The version of latent is not compatible with preview_method.<BR>SD3, SD1/SD2, SDXL, SC-Prior, SC-B and FLUX.1 are not compatible with each other.")
|
||||
|
||||
need_refresh = False
|
||||
latent_changed = False
|
||||
|
||||
# Check if latent has changed
|
||||
if unique_id not in core.preview_bridge_cache:
|
||||
need_refresh = True
|
||||
|
||||
latent_changed = True
|
||||
elif (core.preview_bridge_cache[unique_id][0] is not latent
|
||||
or (vae_opt is None and core.preview_bridge_cache[unique_id][2] is not None)
|
||||
or (vae_opt is None and core.preview_bridge_cache[unique_id][1] != preview_method)
|
||||
or (vae_opt is not None and core.preview_bridge_cache[unique_id][2] is not vae_opt)):
|
||||
need_refresh = True
|
||||
latent_changed = True
|
||||
|
||||
# If latent changed, clear the mask cache to ensure fresh start behavior
|
||||
# unless restore_mask is set to "always" or "if_same_size"
|
||||
if latent_changed and restore_mask not in ["always", "if_same_size"] and unique_id in core.preview_bridge_last_mask_cache:
|
||||
del core.preview_bridge_last_mask_cache[unique_id]
|
||||
|
||||
# Handle clipspace files that aren't registered in the preview bridge system
|
||||
# This only applies when latent hasn't changed (same latent, new mask scenario)
|
||||
if not need_refresh and image not in core.preview_bridge_image_id_map:
|
||||
is_clipspace = image and ("clipspace" in image.lower() or "[input]" in image)
|
||||
if is_clipspace:
|
||||
if not PreviewBridge.register_clipspace_image(image, unique_id):
|
||||
need_refresh = True
|
||||
else:
|
||||
need_refresh = True
|
||||
|
||||
if not need_refresh:
|
||||
pixels, mask, path_item = PreviewBridge.load_image(image)
|
||||
@@ -347,10 +435,18 @@ class PreviewBridgeLatent:
|
||||
|
||||
is_empty_mask = False
|
||||
else:
|
||||
if restore_mask != "never":
|
||||
# For new latents (latent_changed=True), start fresh regardless of restore_mask
|
||||
# For same latent with refresh needed, respect the restore_mask setting
|
||||
# Exception: when restore_mask is "always", restore even with new latents
|
||||
# Exception: when restore_mask is "if_same_size", allow restoration to check size compatibility
|
||||
if restore_mask != "never" and (not latent_changed or restore_mask in ["always", "if_same_size"]):
|
||||
mask = core.preview_bridge_last_mask_cache.get(unique_id)
|
||||
if mask is None or (restore_mask != "always" and mask.shape[1:] != decoded_image.shape[1:3]):
|
||||
if mask is None:
|
||||
mask = None
|
||||
elif restore_mask == "if_same_size" and mask.shape[1:] != decoded_image.shape[1:3]:
|
||||
# For if_same_size, clear mask if dimensions don't match
|
||||
mask = None
|
||||
# For "always", keep the mask regardless of size
|
||||
else:
|
||||
mask = None
|
||||
|
||||
@@ -391,4 +487,4 @@ class PreviewBridgeLatent:
|
||||
return {
|
||||
"ui": {"images": res_image},
|
||||
"result": result,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user