diff --git a/modules/impact/bridge_nodes.py b/modules/impact/bridge_nodes.py index ff817e4..b3923b5 100644 --- a/modules/impact/bridge_nodes.py +++ b/modules/impact/bridge_nodes.py @@ -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.
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, - } + } \ No newline at end of file