From f7686845c6e99b3797b3cabb660ab82869945ccf Mon Sep 17 00:00:00 2001 From: Jonathan Nogueira Date: Fri, 22 Nov 2024 04:41:55 +0000 Subject: [PATCH] Update core.py - save mask when edited (#817) * Update core.py - save mask when edited In order to facilitate the `restore_mask` functionality better, one needs to save the mask to the cache when the mask is edited. * Update core.py - add missing unsqueeze(0) mask cache was missing unsqueeze function, giving the resulting mask the wrong size. --- modules/impact/core.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/modules/impact/core.py b/modules/impact/core.py index a1cc056a..2abfd11d 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -11,6 +11,7 @@ from impact.utils import * from collections import namedtuple import numpy as np from skimage.measure import label +from PIL import ImageOps import nodes import comfy_extras.nodes_upscale_model as model_upscale @@ -69,6 +70,13 @@ def set_previewbridge_image(node_id, file, item): pb_id = f"${node_id}-{pb_id_cnt}" preview_bridge_image_id_map[pb_id] = (file, item) preview_bridge_image_name_map[node_id, file] = (pb_id, item) + if os.path.isfile(file): + i = Image.open(file) + i = ImageOps.exif_transpose(i) + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + preview_bridge_last_mask_cache[node_id] = mask.unsqueeze(0) pb_id_cnt += 1 return pb_id