diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 058ecea..9c58338 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -4576,9 +4576,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] @@ -10817,7 +10818,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"], ), } } @@ -10830,15 +10831,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] @@ -10865,8 +10866,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, ) @@ -11080,7 +11089,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.shape[0] > 1: + 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 @@ -11089,38 +11103,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,) @@ -11182,30 +11199,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,)