nodes/ltdrdata-was-node-suite-comfyui
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge branch 'main' of https://github.com/WASasquatch/was-node-suite-comfyui

Browse Source
This commit is contained in:
Jordan Thompson
2023-08-21 15:33:13 -07:00
parent f159ae38bc c6cadcca87
commit fd7b210597
1 changed files with 75 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
WAS_Node_Suite.py
+75 -50
View File
@@ -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,)