From 6c859ff24054a2faad7dbc847bd6a6390d9f2494 Mon Sep 17 00:00:00 2001 From: HermitCrabe Date: Thu, 17 Aug 2023 16:43:55 -0500 Subject: [PATCH 1/4] Updated WAS_Bounded_Image_Blend and WAS_Bounded_Image_Blend_With_Mas Updated WAS_Bounded_Image_Blend and WAS_Bounded_Image_Blend_With_Mask to support multiple `source` inputs (still only supports a single `target` input) --- WAS_Node_Suite.py | 94 +++++++++++++++++++++++++++-------------------- 1 file changed, 55 insertions(+), 39 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 830ec89..6c08098 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -11065,7 +11065,12 @@ class WAS_Bounded_Image_Blend: def bounded_image_blend(self, target, target_bounds, source, blend_factor, feathering): # Convert PyTorch tensors to PIL images target_pil = Image.fromarray((target.squeeze(0).cpu().numpy() * 255).clip(0, 255).astype(np.uint8)) - source_pil = Image.fromarray((source.squeeze(0).cpu().numpy() * 255).astype(np.uint8)) + source_pils = [] + if source.ndim > 3: + for source_img in source: + source_pils.append(Image.fromarray((source_img.squeeze(0).cpu().numpy() * 255).astype(np.uint8))) + else: + source_pils.append(Image.fromarray((source.squeeze(0).cpu().numpy() * 255).astype(np.uint8))) # Extract the target bounds rmin, rmax, cmin, cmax = target_bounds @@ -11074,38 +11079,41 @@ class WAS_Bounded_Image_Blend: width = cmax - cmin + 1 height = rmax - rmin + 1 - # Resize the source image to match the dimensions of the target bounds - source_resized = source_pil.resize((width, height), Image.ANTIALIAS) + result_tensors = [] + for source_pil in source_pils: + # Resize the source image to match the dimensions of the target bounds + source_resized = source_pil.resize((width, height), Image.ANTIALIAS) - # Create the blend mask with the same size as the target image - blend_mask = Image.new('L', target_pil.size, 0) + # Create the blend mask with the same size as the target image + blend_mask = Image.new('L', target_pil.size, 0) - # Create the feathered mask portion the size of the target bounds - if feathering > 0: - inner_mask = Image.new('L', (width - (2 * feathering), height - (2 * feathering)), 255) - inner_mask = ImageOps.expand(inner_mask, border=feathering, fill=0) - inner_mask = inner_mask.filter(ImageFilter.GaussianBlur(radius=feathering)) - else: - inner_mask = Image.new('L', (width, height), 255) + # Create the feathered mask portion the size of the target bounds + if feathering > 0: + inner_mask = Image.new('L', (width - (2 * feathering), height - (2 * feathering)), 255) + inner_mask = ImageOps.expand(inner_mask, border=feathering, fill=0) + inner_mask = inner_mask.filter(ImageFilter.GaussianBlur(radius=feathering)) + else: + inner_mask = Image.new('L', (width, height), 255) - # Paste the feathered mask portion into the blend mask at the target bounds position - blend_mask.paste(inner_mask, (cmin, rmin)) + # Paste the feathered mask portion into the blend mask at the target bounds position + blend_mask.paste(inner_mask, (cmin, rmin)) - # Create a blank image with the same size and mode as the target - source_positioned = Image.new(target_pil.mode, target_pil.size) + # Create a blank image with the same size and mode as the target + source_positioned = Image.new(target_pil.mode, target_pil.size) - # Paste the source image onto the blank image using the target bounds - source_positioned.paste(source_resized, (cmin, rmin)) + # Paste the source image onto the blank image using the target bounds + source_positioned.paste(source_resized, (cmin, rmin)) - # Create a blend mask using the blend_mask and blend factor - blend_mask = blend_mask.point(lambda p: p * blend_factor).convert('L') + # Create a blend mask using the blend_mask and blend factor + blend_mask = blend_mask.point(lambda p: p * blend_factor).convert('L') - # Blend the source and target images using the blend mask - result = Image.composite(source_positioned, target_pil, blend_mask) + # Blend the source and target images using the blend mask + result = Image.composite(source_positioned, target_pil, blend_mask) - # Convert the result back to a PyTorch tensor - result = torch.from_numpy(np.array(result).astype(np.float32) / 255).unsqueeze(0) - + # Convert the result back to a PyTorch tensor + result_tensors.append(torch.from_numpy(np.array(result).astype(np.float32) / 255).unsqueeze(0)) + + result = torch.cat(result_tensors, dim=0) return (result,) @@ -11167,30 +11175,38 @@ class WAS_Bounded_Image_Blend_With_Mask: # Convert PyTorch tensors to PIL images target_pil = Image.fromarray((target.squeeze(0).cpu().numpy() * 255).clip(0, 255).astype(np.uint8)) target_mask_pil = Image.fromarray((target_mask.cpu().numpy() * 255).astype(np.uint8), mode='L') - source_pil = Image.fromarray((source.squeeze(0).cpu().numpy() * 255).astype(np.uint8)) + source_pils = [] + if source.ndim > 3: + for source_img in source: + source_pils.append(Image.fromarray((source_img.squeeze(0).cpu().numpy() * 255).astype(np.uint8))) + else: + source_pils.append(Image.fromarray((source.squeeze(0).cpu().numpy() * 255).astype(np.uint8))) # Extract the target bounds rmin, rmax, cmin, cmax = target_bounds - # Create a blank image with the same size and mode as the target - source_positioned = Image.new(target_pil.mode, target_pil.size) + result_tensors = [] + for source_pil in source_pils: + # Create a blank image with the same size and mode as the target + source_positioned = Image.new(target_pil.mode, target_pil.size) - # Paste the source image onto the blank image using the target bounds - source_positioned.paste(source_pil, (cmin, rmin)) + # Paste the source image onto the blank image using the target bounds + source_positioned.paste(source_pil, (cmin, rmin)) - # Create a blend mask using the target mask and blend factor - blend_mask = target_mask_pil.point(lambda p: p * blend_factor).convert('L') + # Create a blend mask using the target mask and blend factor + blend_mask = target_mask_pil.point(lambda p: p * blend_factor).convert('L') - # Apply feathering (Gaussian blur) to the blend mask if feather_amount is greater than 0 - if feathering > 0: - blend_mask = blend_mask.filter(ImageFilter.GaussianBlur(radius=feathering)) + # Apply feathering (Gaussian blur) to the blend mask if feather_amount is greater than 0 + if feathering > 0: + blend_mask = blend_mask.filter(ImageFilter.GaussianBlur(radius=feathering)) - # Blend the source and target images using the blend mask - result = Image.composite(source_positioned, target_pil, blend_mask) + # Blend the source and target images using the blend mask + result = Image.composite(source_positioned, target_pil, blend_mask) - # Convert the result back to a PyTorch tensor - result_tensor = torch.from_numpy(np.array(result).astype(np.float32) / 255).unsqueeze(0) + # Convert the result back to a PyTorch tensor + result_tensors.append(torch.from_numpy(np.array(result).astype(np.float32) / 255).unsqueeze(0)) + result_tensor = torch.cat(result_tensors, dim=0) return (result_tensor,) From bf7af541fc12d582815f4327e5fda145acd5705f Mon Sep 17 00:00:00 2001 From: HermitCrabe Date: Thu, 17 Aug 2023 18:56:44 -0500 Subject: [PATCH 2/4] Allow WAS_Image_Batch to combine individual images with sets of images Does not affect error handling for different image dimensions --- WAS_Node_Suite.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 6c08098..3a58d59 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -4561,9 +4561,10 @@ class WAS_Image_Batch: def _check_image_dimensions(self, tensors, names): dimensions = [tensor.shape for tensor in tensors] if len(set(dimensions)) > 1: - mismatched_indices = [i for i, dim in enumerate(dimensions) if dim != dimensions[0]] + mismatched_indices = [i for i, dim in enumerate(dimensions) if dim[1:] != dimensions[0][1:]] mismatched_images = [names[i] for i in mismatched_indices] - raise ValueError(f"WAS Image Batch Warning: Input image dimensions do not match for images: {mismatched_images}") + if mismatched_images: + raise ValueError(f"WAS Image Batch Warning: Input image dimensions do not match for images: {mismatched_images}") def image_batch(self, **kwargs): batched_tensors = [kwargs[key] for key in kwargs if kwargs[key] is not None] From 1a6aa4fab62323df66df32c6533b38dbf8ea078c Mon Sep 17 00:00:00 2001 From: HermitCrabe Date: Thu, 17 Aug 2023 19:11:17 -0500 Subject: [PATCH 3/4] Replaced `ndim` check with `shape[0]` check This doesn't appear to change any functionality, but I think it's a more correct way of checking for the number of images --- WAS_Node_Suite.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 3a58d59..d9cbcb5 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -11067,7 +11067,7 @@ class WAS_Bounded_Image_Blend: # Convert PyTorch tensors to PIL images target_pil = Image.fromarray((target.squeeze(0).cpu().numpy() * 255).clip(0, 255).astype(np.uint8)) source_pils = [] - if source.ndim > 3: + if source.shape[0] > 1: for source_img in source: source_pils.append(Image.fromarray((source_img.squeeze(0).cpu().numpy() * 255).astype(np.uint8))) else: From 2b7fd659d7d23fd6dd3c2d14eff32d1a51056061 Mon Sep 17 00:00:00 2001 From: Tengyang Chen Date: Fri, 18 Aug 2023 17:06:44 +0800 Subject: [PATCH 4/4] small fix for sam node --- WAS_Node_Suite.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 830ec89..2b6d7e8 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -10802,7 +10802,7 @@ class WAS_SAM_Model_Loader: def INPUT_TYPES(self): return { "required": { - "model_size": (["ViT-H (91M)", "ViT-L (308M)", "ViT-B (636M)"], ), + "model_size": (["ViT-H", "ViT-L", "ViT-B"], ), } } @@ -10815,15 +10815,15 @@ class WAS_SAM_Model_Loader: conf = getSuiteConfig() model_filename_mapping = { - "ViT-H (91M)": "sam_vit_h_4b8939.pth", - "ViT-L (308M)": "sam_vit_l_0b3195.pth", - "ViT-B (636M)": "sam_vit_b_01ec64.pth", + "ViT-H": "sam_vit_h_4b8939.pth", + "ViT-L": "sam_vit_l_0b3195.pth", + "ViT-B": "sam_vit_b_01ec64.pth", } model_url_mapping = { - "ViT-H (91M)": conf['sam_model_vith_url'] if conf.__contains__('sam_model_vith_url') else r"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth", - "ViT-L (308M)": conf['sam_model_vitl_url'] if conf.__contains__('sam_model_vitl_url') else r"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth", - "ViT-B (636M)": conf['sam_model_vitb_url'] if conf.__contains__('sam_model_vitb_url') else r"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth", + "ViT-H": conf['sam_model_vith_url'] if conf.__contains__('sam_model_vith_url') else r"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth", + "ViT-L": conf['sam_model_vitl_url'] if conf.__contains__('sam_model_vitl_url') else r"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth", + "ViT-B": conf['sam_model_vitb_url'] if conf.__contains__('sam_model_vitb_url') else r"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth", } model_url = model_url_mapping[model_size] @@ -10850,8 +10850,16 @@ class WAS_SAM_Model_Loader: r = requests.get(model_url, allow_redirects=True) open(sam_file, 'wb').write(r.content) - from segment_anything import build_sam - sam_model = build_sam(checkpoint=sam_file) + from segment_anything import build_sam_vit_h, build_sam_vit_l, build_sam_vit_b + + if model_size == 'ViT-H': + sam_model = build_sam_vit_h(sam_file) + elif model_size == 'ViT-L': + sam_model = build_sam_vit_l(sam_file) + elif model_size == 'ViT-B': + sam_model = build_sam_vit_b(sam_file) + else: + raise ValueError(f"SAM model does not match the model_size: '{model_size}'.") return (sam_model, )