nodes/ManglerFTW-ComfyI2I
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
daa666bbb756bcdcc26e14feddff9721fca8ee92
ManglerFTW-ComfyI2I/ComfyI2I.py
T
Carlo Haas 198f721c3f Update ComfyI2I.py 
mask_mapping_optional should be defined as .long()
2023-09-21 12:26:25 +02:00

1501 lines
58 KiB
Python
Raw Blame History

# By ManglerFTW (Discord: ManglerFTW)
#
# Copyright 2023 Peter Mango (ManglerFTW)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to
# deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
import numpy as np
from collections import namedtuple
import cv2
import torch
import sys
import os
import folder_paths as comfy_paths
from torchvision.ops import masks_to_boxes
import torchvision.transforms.functional as TF
import torch.nn.functional as F
from PIL import Image, ImageFilter, ImageOps
import subprocess
import math
# Check for CUDA availability
device = 'cuda' if torch.cuda.is_available() else 'cpu'
ARRAY_DATATYPE = torch.int32 # Corresponding to 'l'
Rgb = namedtuple('Rgb', ('r', 'g', 'b'))
Hsl = namedtuple('Hsl', ('h', 's', 'l'))
VERY_BIG_SIZE = 1024 * 1024
MAX_RESOLUTION=8192
MODELS_DIR = comfy_paths.models_dir
class cstr(str):
class color:
END = '\33[0m'
BOLD = '\33[1m'
ITALIC = '\33[3m'
UNDERLINE = '\33[4m'
BLINK = '\33[5m'
BLINK2 = '\33[6m'
SELECTED = '\33[7m'
BLACK = '\33[30m'
RED = '\33[31m'
GREEN = '\33[32m'
YELLOW = '\33[33m'
BLUE = '\33[34m'
VIOLET = '\33[35m'
BEIGE = '\33[36m'
WHITE = '\33[37m'
BLACKBG = '\33[40m'
REDBG = '\33[41m'
GREENBG = '\33[42m'
YELLOWBG = '\33[43m'
BLUEBG = '\33[44m'
VIOLETBG = '\33[45m'
BEIGEBG = '\33[46m'
WHITEBG = '\33[47m'
GREY = '\33[90m'
LIGHTRED = '\33[91m'
LIGHTGREEN = '\33[92m'
LIGHTYELLOW = '\33[93m'
LIGHTBLUE = '\33[94m'
LIGHTVIOLET = '\33[95m'
LIGHTBEIGE = '\33[96m'
LIGHTWHITE = '\33[97m'
GREYBG = '\33[100m'
LIGHTREDBG = '\33[101m'
LIGHTGREENBG = '\33[102m'
LIGHTYELLOWBG = '\33[103m'
LIGHTBLUEBG = '\33[104m'
LIGHTVIOLETBG = '\33[105m'
LIGHTBEIGEBG = '\33[106m'
LIGHTWHITEBG = '\33[107m'
@staticmethod
def add_code(name, code):
if not hasattr(cstr.color, name.upper()):
setattr(cstr.color, name.upper(), code)
else:
raise ValueError(f"'cstr' object already contains a code with the name '{name}'.")
def __new__(cls, text):
return super().__new__(cls, text)
def __getattr__(self, attr):
if attr.lower().startswith("_cstr"):
code = getattr(self.color, attr.upper().lstrip("_cstr"))
modified_text = self.replace(f"__{attr[1:]}__", f"{code}")
return cstr(modified_text)
elif attr.upper() in dir(self.color):
code = getattr(self.color, attr.upper())
modified_text = f"{code}{self}{self.color.END}"
return cstr(modified_text)
elif attr.lower() in dir(cstr):
return getattr(cstr, attr.lower())
else:
raise AttributeError(f"'cstr' object has no attribute '{attr}'")
def print(self, **kwargs):
print(self, **kwargs)
def tensor2rgb(t: torch.Tensor) -> torch.Tensor:
size = t.size()
if (len(size) < 4):
return t.unsqueeze(3).repeat(1, 1, 1, 3)
if size[3] == 1:
return t.repeat(1, 1, 1, 3)
elif size[3] == 4:
return t[:, :, :, :3]
else:
return t
def tensor2rgba(t: torch.Tensor) -> torch.Tensor:
size = t.size()
if (len(size) < 4):
return t.unsqueeze(3).repeat(1, 1, 1, 4)
elif size[3] == 1:
return t.repeat(1, 1, 1, 4)
elif size[3] == 3:
alpha_tensor = torch.ones((size[0], size[1], size[2], 1))
return torch.cat((t, alpha_tensor), dim=3)
else:
return t
def tensor2mask(t: torch.Tensor) -> torch.Tensor:
size = t.size()
if (len(size) < 4):
return t
if size[3] == 1:
return t[:,:,:,0]
elif size[3] == 4:
# Not sure what the right thing to do here is. Going to try to be a little smart and use alpha unless all alpha is 1 in case we'll fallback to RGB behavior
if torch.min(t[:, :, :, 3]).item() != 1.:
return t[:,:,:,3]
return TF.rgb_to_grayscale(tensor2rgb(t).permute(0,3,1,2), num_output_channels=1)[:,0,:,:]
def tensor2batch(t: torch.Tensor, bs: torch.Size) -> torch.Tensor:
if len(t.size()) < len(bs):
t = t.unsqueeze(3)
if t.size()[0] < bs[0]:
t.repeat(bs[0], 1, 1, 1)
dim = bs[3]
if dim == 1:
return tensor2mask(t)
elif dim == 3:
return tensor2rgb(t)
elif dim == 4:
return tensor2rgba(t)
def tensors2common(t1: torch.Tensor, t2: torch.Tensor) -> (torch.Tensor, torch.Tensor):
t1s = t1.size()
t2s = t2.size()
if len(t1s) < len(t2s):
t1 = t1.unsqueeze(3)
elif len(t1s) > len(t2s):
t2 = t2.unsqueeze(3)
if len(t1.size()) == 3:
if t1s[0] < t2s[0]:
t1 = t1.repeat(t2s[0], 1, 1)
elif t1s[0] > t2s[0]:
t2 = t2.repeat(t1s[0], 1, 1)
else:
if t1s[0] < t2s[0]:
t1 = t1.repeat(t2s[0], 1, 1, 1)
elif t1s[0] > t2s[0]:
t2 = t2.repeat(t1s[0], 1, 1, 1)
t1s = t1.size()
t2s = t2.size()
if len(t1s) > 3 and t1s[3] < t2s[3]:
return tensor2batch(t1, t2s), t2
elif len(t1s) > 3 and t1s[3] > t2s[3]:
return t1, tensor2batch(t2, t1s)
else:
return t1, t2
# Tensor to PIL
def tensor2pil(image):
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
# PIL to Tensor
def pil2tensor(image):
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
# PIL to Tensor
def pil2tensor_stacked(image):
if isinstance(image, Image.Image):
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
elif isinstance(image, torch.Tensor):
return image
else:
raise ValueError(f"Unexpected datatype for input to 'pil2tensor_stacked'. Expected a PIL Image or tensor, but received type: {type(image)}")
class Color(object):
def __init__(self, r, g, b, proportion):
self.rgb = Rgb(r, g, b)
self.proportion = proportion
def __repr__(self):
return "<colorgram.py Color: {}, {}%>".format(
str(self.rgb), str(self.proportion * 100))
@property
def hsl(self):
try:
return self._hsl
except AttributeError:
self._hsl = Hsl(*hsl(*self.rgb))
return self._hsl
def extract(image_np, number_of_colors, mask_np=None):
# Check and convert the image if needed
if len(image_np.shape) == 2 or image_np.shape[2] != 3: # If grayscale or not RGB
image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
samples = sample(image_np, mask_np)
used = pick_used(samples)
used.sort(key=lambda x: x[0], reverse=True)
return get_colors(samples, used, number_of_colors)
def sample(image, mask=None):
top_two_bits = 0b11000000
sides = 1 << 2
cubes = sides ** 7
samples = torch.zeros((cubes,), dtype=torch.float32, device=device) # Make sure samples is of float32 type
# Handle mask
if mask is not None:
mask_values = (torch.rand_like(mask, dtype=torch.float32) * 255).int()
active_pixels = mask_values > mask
else:
active_pixels = torch.ones_like(image[:, :, 0], dtype=torch.bool)
# Calculate RGB, HSL, and Y
r, g, b = image[:, :, 0], image[:, :, 1], image[:, :, 2]
h, s, l = hsl(r, g, b) # We need to convert the hsl function to use PyTorch
Y = (r * 0.2126 + g * 0.7152 + b * 0.0722).int()
# Packing
packed = ((Y & top_two_bits) << 4) | ((h & top_two_bits) << 2) | (l & top_two_bits)
packed *= 4
# Accumulate samples
packed_active = packed[active_pixels]
r_active, g_active, b_active = r[active_pixels], g[active_pixels], b[active_pixels]
samples.index_add_(0, packed_active, r_active)
samples.index_add_(0, packed_active + 1, g_active)
samples.index_add_(0, packed_active + 2, b_active)
samples.index_add_(0, packed_active + 3, torch.ones_like(packed_active, dtype=torch.float32))
return samples
def pick_used(samples):
# Find indices where count (every 4th value) is non-zero
non_zero_indices = torch.arange(0, samples.size(0), 4, device=samples.device)[samples[3::4] > 0]
# Get counts for non-zero indices
counts = samples[non_zero_indices + 3]
# Combine counts and indices
used = torch.stack((counts, non_zero_indices), dim=-1)
# Convert torch tensors to list of tuples on CPU
used_tuples = [(int(count.item()), int(idx.item())) for count, idx in zip(used[:, 0], used[:, 1])]
return used_tuples
def get_colors(samples, used, number_of_colors):
number_of_colors = min(number_of_colors, len(used))
used = used[:number_of_colors]
# Extract counts and indices
counts, indices = zip(*used)
counts = torch.tensor(counts, dtype=torch.long, device=device)
indices = torch.tensor(indices, dtype=torch.long, device=device)
# Calculate total pixels
total_pixels = torch.sum(counts)
# Get RGB values
r_vals = samples[indices] // counts
g_vals = samples[indices + 1] // counts
b_vals = samples[indices + 2] // counts
# Convert Torch tensors to lists
r_vals_list = r_vals.tolist()
g_vals_list = g_vals.tolist()
b_vals_list = b_vals.tolist()
counts_list = counts.tolist()
# Create Color objects
colors = [Color(r, g, b, count) for r, g, b, count in zip(r_vals_list, g_vals_list, b_vals_list, counts_list)]
# Update proportions
for color in colors:
color.proportion /= total_pixels.item()
return colors
def hsl(r, g, b):
r, g, b = r / 255.0, g / 255.0, b / 255.0
max_val, _ = torch.max(torch.stack([r, g, b]), dim=0)
min_val, _ = torch.min(torch.stack([r, g, b]), dim=0)
diff = max_val - min_val
# Luminance
l = (max_val + min_val) / 2.0
# Saturation
s = torch.where(
(max_val == min_val) | (l == 0),
torch.zeros_like(l),
torch.where(l < 0.5, diff / (max_val + min_val), diff / (2.0 - max_val - min_val))
)
# Hue
conditions = [
max_val == r,
max_val == g,
max_val == b
]
values = [
((g - b) / diff) % 6,
((b - r) / diff) + 2,
((r - g) / diff) + 4
]
h = torch.zeros_like(r)
for condition, value in zip(conditions, values):
h = torch.where(condition, value, h)
h /= 6.0
return (h * 255).int(), (s * 255).int(), (l * 255).int()
def color_distance(pixel_color, palette_color):
return torch.norm(pixel_color - palette_color)
def segment_image(image_torch, palette_colors, mask_torch=None, threshold=128):
"""
Segment the image based on the color similarity of each color in the palette using PyTorch.
"""
if mask_torch is None:
mask_torch = torch.ones(image_torch.shape[:2], device='cuda') * 255
output_image_torch = torch.zeros_like(image_torch)
# Convert palette colors to PyTorch tensor
palette_torch = torch.tensor([list(color.rgb) for color in palette_colors], device='cuda').float()
distances = torch.norm(image_torch.unsqueeze(-2) - palette_torch, dim=-1)
closest_color_indices = torch.argmin(distances, dim=-1)
for idx, palette_color in enumerate(palette_torch):
output_image_torch[closest_color_indices == idx] = palette_color
output_image_torch[mask_torch < threshold] = image_torch[mask_torch < threshold]
# Convert the PyTorch tensor back to a numpy array for saving or further operations
output_image_np = output_image_torch.cpu().numpy().astype('uint8')
return output_image_np
def calculate_luminance_vectorized(colors):
"""Calculate the luminance of an array of RGB colors using PyTorch."""
R, G, B = colors[:, 0], colors[:, 1], colors[:, 2]
return 0.299 * R + 0.587 * G + 0.114 * B
def luminance_match(palette1, palette2):
# Convert palettes to PyTorch tensors
palette1_rgb = torch.tensor([color.rgb for color in palette1], device='cuda').float()
palette2_rgb = torch.tensor([color.rgb for color in palette2], device='cuda').float()
luminance1 = calculate_luminance_vectorized(palette1_rgb)
luminance2 = calculate_luminance_vectorized(palette2_rgb)
# Sort luminances and get the sorted indices
sorted_indices1 = torch.argsort(luminance1)
sorted_indices2 = torch.argsort(luminance2)
reordered_palette2 = [None] * len(palette2)
# Match colors based on sorted luminance order
for idx1, idx2 in zip(sorted_indices1.cpu().numpy(), sorted_indices2.cpu().numpy()):
print(f"idx1: {idx1}, idx2: {idx2}") # Add this to debug
reordered_palette2[idx1] = palette2[idx2]
return reordered_palette2
def apply_blur(image_torch, blur_radius, blur_amount):
image_torch = image_torch.float().div(255.0)
channels = image_torch.shape[2]
kernel_size = int(6 * blur_radius + 1)
kernel_size += 1 if kernel_size % 2 == 0 else 0
# Calculate the padding required to keep the output size the same
padding = kernel_size // 2
# Create a Gaussian kernel
x = torch.linspace(-blur_amount, blur_amount, kernel_size).to(image_torch.device)
x = torch.exp(-x**2 / (2 * blur_radius**2))
x /= x.sum()
kernel = x[:, None] * x[None, :]
# Apply the kernel using depthwise convolution
channels = image_torch.shape[-1]
kernel = kernel[None, None, ...].repeat(channels, 1, 1, 1)
blurred = F.conv2d(image_torch.permute(2, 0, 1)[None, ...], kernel, groups=channels, padding=padding)
# Convert the tensor back to byte and de-normalize
blurred = (blurred * 255.0).byte().squeeze(0).permute(1, 2, 0)
return blurred
def refined_replace_and_blend_colors(Source_np, img_np, palette1, modified_palette2, blur_radius=0, blur_amount=0, mask_torch=None):
# Convert numpy arrays to torch tensors on GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
Source = torch.from_numpy(Source_np).float().to(device)
img_torch = torch.tensor(img_np, device=device).float()
palette1_rgb = torch.stack([torch.tensor(color.rgb, device=device).float() if hasattr(color, 'rgb') else torch.tensor(color, device=device).float() for color in palette1])
modified_palette2_rgb = torch.stack([torch.tensor(color.rgb, device=device).float() if hasattr(color, 'rgb') else torch.tensor(color, device=device).float() for color in modified_palette2])
# Direct color replacement using broadcasting
distances = torch.norm(img_torch[:, :, None] - palette1_rgb, dim=-1)
closest_indices = torch.argmin(distances, dim=-1)
intermediate_output = modified_palette2_rgb[closest_indices]
# Convert to uint8 if not already
intermediate_output = torch.clamp(intermediate_output, 0, 255).byte()
# Apply blur if needed
if blur_radius > 0 and blur_amount > 0:
blurred_output = apply_blur(intermediate_output, blur_radius, blur_amount)
else:
blurred_output = intermediate_output
# Blend based on the mask's intensity values if provided
if mask_torch is not None:
three_channel_mask = mask_torch[:, :, None].expand_as(Source)
output_torch = Source * (1 - three_channel_mask) + blurred_output.float() * three_channel_mask
else:
output_torch = blurred_output
output_np = output_torch.cpu().numpy().astype(np.uint8)
return output_np
def torch_rgb_to_hsv(rgb):
"""
Convert an RGB image to HSV.
Assumes rgb is a PyTorch tensor with values in [0, 1].
"""
# Get R, G, B values
r = rgb[..., 0]
g = rgb[..., 1]
b = rgb[..., 2]
max_val, _ = torch.max(rgb, dim=-1)
min_val, _ = torch.min(rgb, dim=-1)
diff = max_val - min_val
# Calculate Hue
h = torch.zeros_like(r)
h[diff == 0] = 0
mask = (max_val == r) & (diff != 0)
h[mask] = (60 * ((g[mask] - b[mask]) / diff[mask]) + 360) % 360
mask = max_val == g
h[mask] = (60 * ((b[mask] - r[mask]) / diff[mask]) + 120) % 360
mask = max_val == b
h[mask] = (60 * ((r[mask] - g[mask]) / diff[mask]) + 240) % 360
h = h / 360. # Normalize to [0, 1]
# Calculate Saturation
s = torch.zeros_like(r)
s[max_val != 0] = diff[max_val != 0] / max_val[max_val != 0]
# Value
v = max_val
hsv = torch.stack([h, s, v], dim=-1)
return hsv
def torch_hsv_to_rgb(hsv):
"""
Convert an HSV image to RGB.
Assumes hsv is a PyTorch tensor with values in [0, 1] for hue and [0, 1] for saturation/value.
"""
h = hsv[..., 0] * 360.
s = hsv[..., 1]
v = hsv[..., 2]
c = v * s
hh = h / 60.
x = c * (1 - torch.abs(hh % 2 - 1))
m = v - c
r, g, b = v, v, v # Initialize with value
mask = (hh >= 0) & (hh < 1)
r[mask] = c[mask]
g[mask] = x[mask]
mask = (hh >= 1) & (hh < 2)
r[mask] = x[mask]
g[mask] = c[mask]
mask = (hh >= 2) & (hh < 3)
g[mask] = c[mask]
b[mask] = x[mask]
mask = (hh >= 3) & (hh < 4)
g[mask] = x[mask]
b[mask] = c[mask]
mask = (hh >= 4) & (hh < 5)
r[mask] = x[mask]
b[mask] = c[mask]
mask = (hh >= 5) & (hh < 6)
r[mask] = c[mask]
b[mask] = x[mask]
r += m
g += m
b += m
rgb = torch.stack([r, g, b], dim=-1)
return rgb
def retain_luminance_hsv_swap(img1_np, img2_np, strength):
"""
Blend two images while retaining the luminance of the first.
The blending is controlled by the strength parameter.
Assumes img1_np and img2_np are numpy arrays in BGR format.
"""
# Convert BGR to RGB
img1_rgb_np = cv2.cvtColor(img1_np, cv2.COLOR_BGR2RGB).astype(float) / 255.0
img2_rgb_np = cv2.cvtColor(img2_np, cv2.COLOR_BGR2RGB).astype(float) / 255.0
# Blend the two RGB images linearly based on the strength
blended_rgb_np = (1 - strength) * img1_rgb_np + strength * img2_rgb_np
# Convert the blended RGB image and the original RGB image to YUV
blended_yuv_np = cv2.cvtColor((blended_rgb_np * 255).astype(np.uint8), cv2.COLOR_RGB2YUV)
img1_yuv_np = cv2.cvtColor(img1_np, cv2.COLOR_BGR2YUV)
# Replace the Y channel (luminance) of the blended image with the original image's luminance
blended_yuv_np[:,:,0] = img1_yuv_np[:,:,0]
# Convert back to BGR
result_bgr_np = cv2.cvtColor(blended_yuv_np, cv2.COLOR_YUV2BGR)
return result_bgr_np
def adjust_gamma_contrast(image_np, gamma, contrast, brightness, mask_np=None):
# Ensure CUDA is available
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Transfer data to PyTorch tensors and move to the appropriate device
image_torch = torch.tensor(image_np, dtype=torch.float32).to(device)
# Gamma correction using a lookup table
inv_gamma = 1.0 / gamma
table = torch.tensor([(i / 255.0) ** inv_gamma * 255 for i in range(256)], device=device).float()
gamma_corrected = torch.index_select(table, 0, image_torch.long().flatten()).reshape_as(image_torch)
# Contrast and brightness adjustment
contrast_adjusted = contrast * gamma_corrected + brightness
contrast_adjusted = torch.clamp(contrast_adjusted, 0, 255).byte()
# If mask is provided, blend the original and adjusted images
if mask_np is not None:
mask_torch = torch.tensor(mask_np, device=device).float() / 255.0
three_channel_mask = mask_torch.unsqueeze(-1).expand_as(image_torch)
contrast_adjusted = image_torch * (1 - three_channel_mask) + contrast_adjusted.float() * three_channel_mask
# Transfer data back to numpy array
result_np = contrast_adjusted.cpu().numpy()
return result_np
def CutByMask(image, mask, force_resize_width, force_resize_height, mask_mapping_optional):
if len(image.shape) < 4:
C = 1
else:
C = image.shape[3]
# We operate on RGBA to keep the code clean and then convert back after
image = tensor2rgba(image)
mask = tensor2mask(mask)
if mask_mapping_optional is not None:
mask_mapping_optional = mask_mapping_optional.long()
image = image[mask_mapping_optional]
# Scale the mask to match the image size if it isn't
B, H, W, _ = image.shape
mask = F.interpolate(mask.unsqueeze(1), size=(H, W), mode='nearest')[:,0,:,:]
MB, _, _ = mask.shape
if MB < B:
assert(B % MB == 0)
mask = mask.repeat(B // MB, 1, 1)
# Masks to boxes
is_empty = ~torch.gt(torch.max(torch.reshape(mask, [B, H * W]), dim=1).values, 0.)
mask[is_empty,0,0] = 1.
boxes = masks_to_boxes(mask)
mask[is_empty,0,0] = 0.
min_x = boxes[:,0]
min_y = boxes[:,1]
max_x = boxes[:,2]
max_y = boxes[:,3]
width = max_x - min_x + 1
height = max_y - min_y + 1
use_width = int(torch.max(width).item())
use_height = int(torch.max(height).item())
if force_resize_width > 0:
use_width = force_resize_width
if force_resize_height > 0:
use_height = force_resize_height
print("use_width: ", use_width)
print("use_height: ", use_height)
alpha_mask = torch.ones((B, H, W, 4))
alpha_mask[:,:,:,3] = mask
image = image * alpha_mask
result = torch.zeros((B, use_height, use_width, 4))
for i in range(0, B):
if not is_empty[i]:
ymin = int(min_y[i].item())
ymax = int(max_y[i].item())
xmin = int(min_x[i].item())
xmax = int(max_x[i].item())
single = (image[i, ymin:ymax+1, xmin:xmax+1,:]).unsqueeze(0)
resized = F.interpolate(single.permute(0, 3, 1, 2), size=(use_height, use_width), mode='bicubic').permute(0, 2, 3, 1)
result[i] = resized[0]
# Preserve our type unless we were previously RGB and added non-opaque alpha due to the mask size
if C == 1:
print("C == 1 output image shape: ", tensor2mask(result).shape)
return tensor2mask(result)
elif C == 3 and torch.min(result[:,:,:,3]) == 1:
print("C == 3 output image shape: ", tensor2rgb(result).shape)
return tensor2rgb(result)
else:
print("else result shape: ", result.shape)
return result
def combine(image1, image2, op, clamp_result, round_result):
image1, image2 = tensors2common(image1, image2)
if op == "union (max)":
result = torch.max(image1, image2)
elif op == "intersection (min)":
result = torch.min(image1, image2)
elif op == "difference":
result = image1 - image2
elif op == "multiply":
result = image1 * image2
elif op == "multiply_alpha":
image1 = tensor2rgba(image1)
image2 = tensor2mask(image2)
result = torch.cat((image1[:, :, :, :3], (image1[:, :, :, 3] * image2).unsqueeze(3)), dim=3)
elif op == "add":
result = image1 + image2
elif op == "greater_or_equal":
result = torch.where(image1 >= image2, 1., 0.)
elif op == "greater":
result = torch.where(image1 > image2, 1., 0.)
if clamp_result == "yes":
result = torch.min(torch.max(result, torch.tensor(0.)), torch.tensor(1.))
if round_result == "yes":
result = torch.round(result)
return result
def apply_color_correction(target_image, source_image, factor=1):
if not isinstance(source_image, (torch.Tensor, Image.Image)):
raise ValueError(f"Unexpected datatype for 'source_image' at method start. Expected a tensor or PIL Image, but received type: {type(source_image)}")
# Ensure source_image is a tensor
if isinstance(source_image, Image.Image): # Check if it's a PIL Image
source_image = pil2tensor_stacked(source_image) # Convert it to tensor
if not isinstance(source_image, (torch.Tensor, Image.Image)):
raise ValueError(f"Unexpected datatype for 'source_image'. Expected a tensor or PIL Image, but received type: {type(source_image)}")
# Get the batch size
batch_size = source_image.shape[0]
output_images = []
for i in range(batch_size):
# Convert the source and target images to NumPy arrays for the i-th image in the batch
source_numpy = source_image[i, ...].numpy()
target_numpy = target_image[i, ...].numpy()
# Convert to float32
source_numpy = source_numpy.astype(np.float32)
target_numpy = target_numpy.astype(np.float32)
# If the images have an alpha channel, remove it for the color transformations
if source_numpy.shape[-1] == 4:
source_numpy = source_numpy[..., :3]
if target_numpy.shape[-1] == 4:
target_numpy = target_numpy[..., :3]
# Compute the mean and standard deviation of the color channels for both images
target_mean, target_std = np.mean(source_numpy, axis=(0, 1)), np.std(source_numpy, axis=(0, 1))
source_mean, source_std = np.mean(target_numpy, axis=(0, 1)), np.std(target_numpy, axis=(0, 1))
adjusted_source_mean = target_mean + factor * (target_mean - source_mean)
adjusted_source_std = target_std + factor * (target_std - source_std)
# Normalize the target image (zero mean and unit variance)
target_norm = (target_numpy - target_mean) / target_std
# Scale and shift the normalized target image to match the exaggerated source image statistics
matched_rgb = target_norm * adjusted_source_std + adjusted_source_mean
# Clip values to [0, 1] and convert to PIL Image
img = Image.fromarray(np.clip(matched_rgb * 255, 0, 255).astype('uint8'), 'RGB')
# Convert the PIL Image to a tensor and append to the list
img_tensor = pil2tensor_stacked(img)
output_images.append(img_tensor)
# Stack the list of tensors to get the batch of corrected images
stacked_images = torch.stack(output_images)
return stacked_images
def PasteByMask(image_base, image_to_paste, mask, resize_behavior, mask_mapping_optional):
image_base = tensor2rgba(image_base)
image_to_paste = tensor2rgba(image_to_paste)
mask = tensor2mask(mask)
# Scale the mask to be a matching size if it isn't
B, H, W, C = image_base.shape
MB = mask.shape[0]
PB = image_to_paste.shape[0]
if mask_mapping_optional is None:
if B < PB:
assert(PB % B == 0)
image_base = image_base.repeat(PB // B, 1, 1, 1)
B, H, W, C = image_base.shape
if MB < B:
assert(B % MB == 0)
mask = mask.repeat(B // MB, 1, 1)
elif B < MB:
assert(MB % B == 0)
image_base = image_base.repeat(MB // B, 1, 1, 1)
if PB < B:
assert(B % PB == 0)
image_to_paste = image_to_paste.repeat(B // PB, 1, 1, 1)
mask = F.interpolate(mask.unsqueeze(1), size=(H, W), mode='nearest')[:,0,:,:]
MB, MH, MW = mask.shape
# masks_to_boxes errors if the tensor is all zeros, so we'll add a single pixel and zero it out at the end
is_empty = ~torch.gt(torch.max(torch.reshape(mask,[MB, MH * MW]), dim=1).values, 0.)
mask[is_empty,0,0] = 1.
boxes = masks_to_boxes(mask)
mask[is_empty,0,0] = 0.
min_x = boxes[:,0]
min_y = boxes[:,1]
max_x = boxes[:,2]
max_y = boxes[:,3]
mid_x = (min_x + max_x) / 2
mid_y = (min_y + max_y) / 2
target_width = max_x - min_x + 1
target_height = max_y - min_y + 1
result = image_base.detach().clone()
for i in range(0, MB):
if i >= len(image_to_paste):
raise ValueError(f"image_to_paste does not have an entry for mask index {i}")
if is_empty[i]:
continue
else:
image_index = i
if mask_mapping_optional is not None:
image_index = mask_mapping_optional[i].item()
source_size = image_to_paste.size()
SB, SH, SW, _ = image_to_paste.shape
# Figure out the desired size
width = int(target_width[i].item())
height = int(target_height[i].item())
if resize_behavior == "keep_ratio_fill":
target_ratio = width / height
actual_ratio = SW / SH
if actual_ratio > target_ratio:
width = int(height * actual_ratio)
elif actual_ratio < target_ratio:
height = int(width / actual_ratio)
elif resize_behavior == "keep_ratio_fit":
target_ratio = width / height
actual_ratio = SW / SH
if actual_ratio > target_ratio:
height = int(width / actual_ratio)
elif actual_ratio < target_ratio:
width = int(height * actual_ratio)
elif resize_behavior == "source_size" or resize_behavior == "source_size_unmasked":
width = SW
height = SH
# Resize the image we're pasting if needed
resized_image = image_to_paste[i].unsqueeze(0)
if SH != height or SW != width:
resized_image = F.interpolate(resized_image.permute(0, 3, 1, 2), size=(height,width), mode='bicubic').permute(0, 2, 3, 1)
pasting = torch.ones([H, W, C])
ymid = float(mid_y[i].item())
ymin = int(math.floor(ymid - height / 2)) + 1
ymax = int(math.floor(ymid + height / 2)) + 1
xmid = float(mid_x[i].item())
xmin = int(math.floor(xmid - width / 2)) + 1
xmax = int(math.floor(xmid + width / 2)) + 1
_, source_ymax, source_xmax, _ = resized_image.shape
source_ymin, source_xmin = 0, 0
if xmin < 0:
source_xmin = abs(xmin)
xmin = 0
if ymin < 0:
source_ymin = abs(ymin)
ymin = 0
if xmax > W:
source_xmax -= (xmax - W)
xmax = W
if ymax > H:
source_ymax -= (ymax - H)
ymax = H
pasting[ymin:ymax, xmin:xmax, :] = resized_image[0, source_ymin:source_ymax, source_xmin:source_xmax, :]
pasting[:, :, 3] = 1.
pasting_alpha = torch.zeros([H, W])
pasting_alpha[ymin:ymax, xmin:xmax] = resized_image[0, source_ymin:source_ymax, source_xmin:source_xmax, 3]
if resize_behavior == "keep_ratio_fill" or resize_behavior == "source_size_unmasked":
# If we explicitly want to fill the area, we are ok with extending outside
paste_mask = pasting_alpha.unsqueeze(2).repeat(1, 1, 4)
else:
paste_mask = torch.min(pasting_alpha, mask[i]).unsqueeze(2).repeat(1, 1, 4)
result[image_index] = pasting * paste_mask + result[image_index] * (1. - paste_mask)
return result
class Mask_Ops:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"text": ("STRING", {"default":"", "multiline": False}),
"separate_mask": ("INT", {"default":0, "min":0, "max":1, "step":1}),
"text_sigma": ("INT", {"default":30, "min":0, "max":150, "step":1}),
"use_text": ("INT", {"default":0, "min":0, "max":1, "step":1}),
"blend_percentage": ("FLOAT", {"default": 0, "min": 0.0, "max": 1.0, "step": 0.01}),
"black_level": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 255.0, "step": 0.1}),
"mid_level": ("FLOAT", {"default": 127.5, "min": 0.0, "max": 255.0, "step": 0.1}),
"white_level": ("FLOAT", {"default": 255, "min": 0.0, "max": 255.0, "step": 0.1}),
"channel": (["red", "green", "blue"],),
"shrink_grow": ("INT", {"default": 0, "min": -128, "max": 128, "step": 1}),
"invert": ("INT", {"default":0, "min":0, "max":1, "step":1}),
"blur_radius": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 1024, "step": 0.1}),
},
"optional": {
"mask": ("MASK",),
},
}
CATEGORY = "I2I"
RETURN_TYPES = ("IMAGE", "MASK", "MASK_MAPPING",)
RETURN_NAMES = ("mask_image", "mask", "mask mapping")
FUNCTION = "Mask_Ops"
def Mask_Ops(self, image, text, separate_mask, text_sigma, use_text, blend_percentage, black_level, mid_level, white_level, channel, shrink_grow, invert=0, blur_radius=5.0, mask=None):
channels = ["red", "green", "blue"]
# Freeze PIP modules
def packages(versions=False):
import sys
import subprocess
return [( r.decode().split('==')[0] if not versions else r.decode() ) for r in subprocess.check_output([sys.executable, '-s', '-m', 'pip', 'freeze']).split()]
# PIL to Mask
def pil2mask(image):
image_np = np.array(image.convert("L")).astype(np.float32) / 255.0
mask = torch.from_numpy(image_np)
return 1.0 - mask
def gaussian_region(image, radius=5.0):
image = ImageOps.invert(image.convert("L"))
image = image.filter(ImageFilter.GaussianBlur(radius=int(radius)))
return image.convert("RGB")
# scipy handling
if 'scipy' not in packages():
cstr("Installing `scipy` ...").msg.print()
subprocess.check_call([sys.executable, '-s', '-m', 'pip', 'install', 'scipy'])
try:
import scipy
except ImportError as e:
cstr("Unable to import tools for certain masking procedures.").msg.print()
print(e)
def smooth_region(image, tolerance):
from scipy.ndimage import gaussian_filter
image = image.convert("L")
mask_array = np.array(image)
smoothed_array = gaussian_filter(mask_array, sigma=tolerance)
threshold = np.max(smoothed_array) / 2
smoothed_mask = np.where(smoothed_array >= threshold, 255, 0).astype(np.uint8)
smoothed_image = Image.fromarray(smoothed_mask, mode="L")
return ImageOps.invert(smoothed_image.convert("RGB"))
def erode_region(image, iterations):
from scipy.ndimage import binary_erosion
image = image.convert("L")
binary_mask = np.array(image) > 0
eroded_mask = binary_erosion(binary_mask, iterations=iterations)
eroded_image = Image.fromarray(eroded_mask.astype(np.uint8) * 255, mode="L")
return ImageOps.invert(eroded_image.convert("RGB"))
def dilate_region(image, iterations):
from scipy.ndimage import binary_dilation
image = image.convert("L")
binary_mask = np.array(image) > 0
dilated_mask = binary_dilation(binary_mask, iterations=iterations)
dilated_image = Image.fromarray(dilated_mask.astype(np.uint8) * 255, mode="L")
return ImageOps.invert(dilated_image.convert("RGB"))
def erode(masks, iterations):
iterations = iterations * -1
if masks.ndim > 3:
regions = []
for mask in masks:
mask_np = np.clip(255. * mask.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
pil_image = Image.fromarray(mask_np, mode="L")
region_mask = erode_region(pil_image, iterations)
region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
regions.append(region_tensor)
regions_tensor = torch.cat(regions, dim=0)
return regions_tensor
else:
mask_np = np.clip(255. * masks.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
pil_image = Image.fromarray(mask_np, mode="L")
region_mask = erode_region(pil_image, iterations)
region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
return region_tensor
def dilate(masks, iterations):
if masks.ndim > 3:
regions = []
for mask in masks:
mask_np = np.clip(255. * mask.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
pil_image = Image.fromarray(mask_np, mode="L")
region_mask = dilate_region(pil_image, iterations)
region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
regions.append(region_tensor)
regions_tensor = torch.cat(regions, dim=0)
return regions_tensor
else:
mask_np = np.clip(255. * masks.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
pil_image = Image.fromarray(mask_np, mode="L")
region_mask = dilate_region(pil_image, iterations)
region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
return region_tensor
def separate(mask, separate_flag=1):
if separate_flag == 0:
# Create an unseparated mapping tensor of the same length as the batch dimension
mapping = torch.arange(mask.shape[0], device=mask.device, dtype=torch.int)
return mask, mapping
mask = tensor2mask(mask)
thresholded = torch.gt(mask, 0).unsqueeze(1)
B, H, W = mask.shape
components = torch.arange(B * H * W, device=mask.device, dtype=mask.dtype).reshape(B, 1, H, W) + 1
components[~thresholded] = 0
while True:
previous_components = components
components = F.max_pool2d(components, kernel_size=3, stride=1, padding=1)
components[~thresholded] = 0
if torch.equal(previous_components, components):
break
components = components.reshape(B, H, W)
segments = torch.unique(components)
result = torch.zeros([len(segments) - 1, H, W])
index = 0
mapping = torch.zeros([len(segments) - 1], device=mask.device, dtype=torch.int)
for i in range(len(segments)):
segment = segments[i].item()
if segment == 0:
continue
image_index = int((segment - 1) // (H * W))
segment_mask = (components[image_index,:,:] == segment)
result[index][segment_mask] = mask[image_index][segment_mask]
mapping[index] = image_index
index += 1
return result, mapping
image = tensor2pil(image)
use_text = int(round(use_text))
if use_text == 1:
# CLIPSeg Model Loader
from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
model = "CIDAS/clipseg-rd64-refined"
cache = os.path.join(MODELS_DIR, 'clipseg')
inputs = CLIPSegProcessor.from_pretrained(model, cache_dir=cache)
model = CLIPSegForImageSegmentation.from_pretrained(model, cache_dir=cache)
image = image.convert('RGB')
cache = os.path.join(MODELS_DIR, 'clipseg')
with torch.no_grad():
result = model(**inputs(text=text, images=image, padding=True, return_tensors="pt"))
tensor = torch.sigmoid(result[0])
mask = 1. - (tensor - tensor.min()) / tensor.max()
mask = mask.unsqueeze(0)
mask = tensor2pil(mask).convert("L")
mask = mask.resize(image.size)
sigma = text_sigma
mask = pil2mask(mask)
if mask.ndim > 3:
regions = []
for mk in mask:
mask_np = np.clip(255. * mk.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
pil_image = Image.fromarray(mask_np, mode="L")
region_mask = smooth_region(pil_image, sigma)
region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
regions.append(region_tensor)
mask = torch.cat(regions, dim=0)
else:
mask_np = np.clip(255. * mask.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
pil_image = Image.fromarray(mask_np, mode="L")
region_mask = smooth_region(pil_image, sigma)
mask = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
else:
if mask is None:
# Create a full mask for the entire image
mask_shape = (image.size[1], image.size[0]) # Assuming image is in (H, W, C) format
mask = torch.ones(mask_shape, dtype=torch.float32)
else:
mask = mask
if shrink_grow < 0:
mask = erode(mask, shrink_grow)
elif shrink_grow > 0:
mask = dilate(mask, shrink_grow)
invert = int(round(invert))
if invert == 1:
mask = 1.0 - mask
#Invert Mask
Mask_Inv = mask
#Convert Inverted Mask to Image
Inv_Mask_2_Img = Mask_Inv.reshape((-1, 1, Mask_Inv.shape[-2], Mask_Inv.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
#Convert Mask to Image
Mask_2_Img = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
#Image Blend by Mask
# Convert images to PIL
img_a = tensor2pil(Inv_Mask_2_Img)
img_b = image
mask = ImageOps.invert(tensor2pil(Mask_2_Img).convert('L'))
# Mask image
masked_img = Image.composite(img_a, img_b, mask.resize(img_a.size))
# Blend image
blend_mask = Image.new(mode="L", size=img_a.size,
color=(round(blend_percentage * 255)))
blend_mask = ImageOps.invert(blend_mask)
Blended_Image = Image.composite(img_a, masked_img, blend_mask)
Blended_Image = pil2tensor(Blended_Image)
del img_a, img_b, blend_mask, mask
#Image Levels Adjustment
# Convert image to PIL
tensor_images = []
for img in Blended_Image:
img = tensor2pil(img)
img = img.convert("RGB")
levels = self.AdjustLevels(black_level, mid_level, white_level)
tensor_images.append(pil2tensor(levels.adjust(img)))
tensor_images = torch.cat(tensor_images, dim=0)
#Convert Image to Mask
masks = tensor_images[0, :, :, channels.index(channel)]
if masks.ndim > 3:
regions = []
for mask in masks:
mask_np = np.clip(255. * mask.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
pil_image = Image.fromarray(mask_np, mode="L")
region_mask = gaussian_region(pil_image, blur_radius)
region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
regions.append(region_tensor)
result = torch.cat(regions, dim=0)
else:
mask_np = np.clip(255. * masks.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
pil_image = Image.fromarray(mask_np, mode="L")
region_mask = gaussian_region(pil_image, blur_radius)
result = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
result = result.reshape((-1, 1, result.shape[-2], result.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
if invert == 1:
result = 1.0 - result
result, mapping = separate(result, separate_mask)
if invert == 1:
result = 1.0 - result
return (result, result, mapping,)
class AdjustLevels:
def __init__(self, min_level, mid_level, max_level):
self.min_level = min_level
self.mid_level = mid_level
self.max_level = max_level
def adjust(self, im):
im_arr = np.array(im)
im_arr[im_arr < self.min_level] = self.min_level
im_arr = (im_arr - self.min_level) * \
(255 / (self.max_level - self.min_level))
im_arr[im_arr < 0] = 0
im_arr[im_arr > 255] = 255
im_arr = im_arr.astype(np.uint8)
im = Image.fromarray(im_arr)
im = ImageOps.autocontrast(im, cutoff=self.max_level)
return im
class Color_Correction:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"source_image": ("IMAGE",),
"target_image": ("IMAGE",),
"no_of_colors": ("INT", {"default": 6, "min": 0, "max": 256, "step": 1}),
"blur_radius": ("INT", {"default": 2, "min": 0, "max": 100, "step": 1}),
"blur_amount": ("INT", {"default": 2, "min": 0, "max": 100, "step": 1}),
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
"gamma": ("FLOAT", {"default": 1.0, "min": 0.10, "max": 2.0, "step": 0.1}),
"contrast": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.1}),
"brightness": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1}),
},
"optional": {
"mask": ("MASK",),
},
}
CATEGORY = "I2I"
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("image", )
FUNCTION = "ColorXfer2"
def ColorXfer2(cls, source_image, target_image, no_of_colors, blur_radius, blur_amount, strength, gamma, contrast, brightness, mask=None):
if mask is not None:
if torch.is_tensor(mask):
# Convert to grayscale if it's a 3-channel image
if mask.shape[-1] == 3:
mask = torch.mean(mask, dim=-1)
# Remove batch dimension if present
if mask.dim() == 3:
mask = mask.squeeze(0)
mask_np1 = (mask.cpu().numpy() * 255).astype(np.uint8)
else:
mask_np1 = (mask * 255).astype(np.uint8)
mask_np = mask_np1 / 255.0
mask_torch = torch.tensor(mask_np).to(device)
# If the source_image is a tensor, convert it to a numpy array
if torch.is_tensor(source_image):
Source_np = (source_image[0].cpu().numpy() * 255).astype(np.uint8)
else:
Source_np = (source_image * 255).astype(np.uint8)
# If the source_image is a tensor, convert it to a numpy array
if torch.is_tensor(target_image):
Target_np = (target_image[0].cpu().numpy() * 255).astype(np.uint8)
else:
Target_np = (target_image * 255).astype(np.uint8)
# Load the source image and convert to torch tensor
Source_np = cv2.cvtColor(Source_np, cv2.COLOR_BGR2RGB)
Source = torch.from_numpy(Source_np).float().to(device)
# Extract colors from the source image
colors1 = extract(Source, no_of_colors, mask_np=mask_torch)
# Load the target image
Target_np = cv2.cvtColor(Target_np, cv2.COLOR_BGR2RGB)
Target = torch.from_numpy(Target_np).float().to(device=device)
# Extract colors from the target image
colors2 = extract(Target, no_of_colors)
min_length = min(len(colors1), len(colors2))
colors1 = colors1[:min_length]
colors2 = colors2[:min_length]
# Segment the image
segmented_np = segment_image(Source, colors1, mask_torch=mask_torch, threshold=1)
matched_pairs = luminance_match(colors1, colors2)
result_rgb = refined_replace_and_blend_colors(Source.cpu().numpy(), segmented_np, colors1, matched_pairs, blur_radius, blur_amount, mask_torch=mask_torch)
luminance_np = retain_luminance_hsv_swap(Source.cpu().numpy(), result_rgb, strength)
gamma_contrast_np = adjust_gamma_contrast(luminance_np, gamma, contrast, brightness, mask_np=mask_np1)
final_img_np_rgb = cv2.cvtColor(gamma_contrast_np, cv2.COLOR_BGR2RGB)
# Convert the numpy array back to a PyTorch tensor
final_img_tensor = torch.tensor(final_img_np_rgb).float().to(device)
final_img_tensor = final_img_tensor.unsqueeze(0)
if final_img_tensor.max() > 1.0:
final_img_tensor /= 255.0
return (final_img_tensor, )
class MaskToRegion:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"mask": ("IMAGE",),
"force_resize_width": ("INT", {"default": 1024, "min": 0, "max": VERY_BIG_SIZE, "step": 1}),
"force_resize_height": ("INT", {"default": 1024, "min": 0, "max": VERY_BIG_SIZE, "step": 1}),
"kind": (["mask", "RGB", "RGBA"],),
"padding": ("INT", {"default": 3, "min": 0, "max": VERY_BIG_SIZE, "step": 1}),
"constraints": (["keep_ratio", "keep_ratio_divisible", "multiple_of", "ignore"],),
"constraint_x": ("INT", {"default": 64, "min": 2, "max": VERY_BIG_SIZE, "step": 1}),
"constraint_y": ("INT", {"default": 64, "min": 2, "max": VERY_BIG_SIZE, "step": 1}),
"min_width": ("INT", {"default": 0, "min": 0, "max": VERY_BIG_SIZE, "step": 1}),
"min_height": ("INT", {"default": 0, "min": 0, "max": VERY_BIG_SIZE, "step": 1}),
"batch_behavior": (["match_ratio", "match_size"],),
},
"optional": {
"mask_mapping_optional": ("MASK_MAPPING",),
},
}
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", )
RETURN_NAMES = ("cut image", "cut mask", "region")
FUNCTION = "get_region"
CATEGORY = "I2I"
def get_region(self, image, mask, force_resize_width, force_resize_height, kind, padding, constraints, constraint_x, constraint_y, min_width, min_height, batch_behavior, mask_mapping_optional = None):
mask2 = tensor2mask(mask)
mask_size = mask2.size()
mask_width = int(mask_size[2])
mask_height = int(mask_size[1])
# masks_to_boxes errors if the tensor is all zeros, so we'll add a single pixel and zero it out at the end
is_empty = ~torch.gt(torch.max(torch.reshape(mask2,[mask_size[0], mask_width * mask_height]), dim=1).values, 0.)
mask2[is_empty,0,0] = 1.
boxes = masks_to_boxes(mask2)
mask2[is_empty,0,0] = 0.
# Account for padding
min_x = torch.max(boxes[:,0] - padding, torch.tensor(0.))
min_y = torch.max(boxes[:,1] - padding, torch.tensor(0.))
max_x = torch.min(boxes[:,2] + padding, torch.tensor(mask_width))
max_y = torch.min(boxes[:,3] + padding, torch.tensor(mask_height))
width = max_x - min_x
height = max_y - min_y
# Make sure the width and height are big enough
target_width = torch.max(width, torch.tensor(min_width))
target_height = torch.max(height, torch.tensor(min_height))
if constraints == "keep_ratio":
target_width = torch.max(target_width, target_height * constraint_x // constraint_y)
target_height = torch.max(target_height, target_width * constraint_y // constraint_x)
elif constraints == "keep_ratio_divisible":
# Probably a more efficient way to do this, but given the bounds it's not too bad
max_factors = torch.min(constraint_x // target_width, constraint_y // target_height)
max_factor = int(torch.max(max_factors).item())
for i in range(1, max_factor+1):
divisible = constraint_x % i == 0 and constraint_y % i == 0
if divisible:
big_enough = ~torch.lt(target_width, constraint_x // i) * ~torch.lt(target_height, constraint_y // i)
target_width[big_enough] = constraint_x // i
target_height[big_enough] = constraint_y // i
elif constraints == "multiple_of":
target_width[torch.gt(target_width % constraint_x, 0)] = (target_width // constraint_x + 1) * constraint_x
target_height[torch.gt(target_height % constraint_y, 0)] = (target_height // constraint_y + 1) * constraint_y
if batch_behavior == "match_size":
target_width[:] = torch.max(target_width)
target_height[:] = torch.max(target_height)
elif batch_behavior == "match_ratio":
# We'll target the ratio that's closest to 1:1, but don't want to take into account empty masks
ratios = torch.abs(target_width / target_height - 1)
ratios[is_empty] = 10000
match_ratio = torch.min(ratios,dim=0).indices.item()
target_width = torch.max(target_width, target_height * target_width[match_ratio] // target_height[match_ratio])
target_height = torch.max(target_height, target_width * target_height[match_ratio] // target_width[match_ratio])
missing = target_width - width
min_x = min_x - missing // 2
max_x = max_x + (missing - missing // 2)
missing = target_height - height
min_y = min_y - missing // 2
max_y = max_y + (missing - missing // 2)
# Move the region into range if needed
bad = torch.lt(min_x,0)
max_x[bad] -= min_x[bad]
min_x[bad] = 0
bad = torch.lt(min_y,0)
max_y[bad] -= min_y[bad]
min_y[bad] = 0
bad = torch.gt(max_x, mask_width)
min_x[bad] -= (max_x[bad] - mask_width)
max_x[bad] = mask_width
bad = torch.gt(max_y, mask_height)
min_y[bad] -= (max_y[bad] - mask_height)
max_y[bad] = mask_height
region = torch.zeros((mask_size[0], mask_height, mask_width))
for i in range(0, mask_size[0]):
if not is_empty[i]:
ymin = int(min_y[i].item())
ymax = int(max_y[i].item())
xmin = int(min_x[i].item())
xmax = int(max_x[i].item())
region[i, ymin:ymax+1, xmin:xmax+1] = 1
Cut_Image = CutByMask(image, region, force_resize_width, force_resize_height, mask_mapping_optional)
#Change Channels >>>> OUTPUT TO VAE ENCODE
if kind == "mask":
Cut_Image = tensor2mask(Cut_Image)
elif kind == "RGBA":
Cut_Image = tensor2rgba(Cut_Image)
else: # RGB
Cut_Image = tensor2rgb(Cut_Image)
Cut_Mask = CutByMask(mask, region, force_resize_width, force_resize_height, mask_mapping_optional = None)
return (Cut_Image, Cut_Mask, region, )
class Combine_And_Paste_Op:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"decoded_vae": ("IMAGE",),
"Original_Image": ("IMAGE",),
"Cut_Image": ("IMAGE",),
"Cut_Mask": ("IMAGE",),
"region": ("IMAGE",),
"color_xfer_factor": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
"op": (["union (max)", "intersection (min)", "difference", "multiply", "multiply_alpha", "add", "greater_or_equal", "greater"],),
"clamp_result": (["yes", "no"],),
"round_result": (["no", "yes"],),
"resize_behavior": (["resize", "keep_ratio_fill", "keep_ratio_fit", "source_size", "source_size_unmasked"],),
},
"optional": {
"mask_mapping_optional": ("MASK_MAPPING",),
},
}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("FinalOut", )
FUNCTION = "com_paste_op"
CATEGORY = "I2I"
def com_paste_op(self, decoded_vae, Original_Image, Cut_Image, Cut_Mask, region, color_xfer_factor, op, clamp_result, round_result, resize_behavior, mask_mapping_optional = None):
Combined_Decoded = combine(decoded_vae, Cut_Mask, op, clamp_result, round_result)
Combined_Originals = combine(Cut_Image, Cut_Mask, op, clamp_result, round_result)
Cx_Decoded = apply_color_correction(Combined_Decoded, Combined_Originals, color_xfer_factor)
Cx_Decode_Mask = combine(Cx_Decoded, Cut_Mask, op, clamp_result, round_result)
FinalOut = PasteByMask(Original_Image, Cx_Decode_Mask, region, resize_behavior, mask_mapping_optional)
return (FinalOut, )
NODE_CLASS_MAPPINGS = {
"Color Transfer": Color_Correction,
"Mask Ops": Mask_Ops,
"Inpaint Segments": MaskToRegion,
"Combine and Paste": Combine_And_Paste_Op,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Color Transfer": "Color Transfer",
"Mask Ops": "Mask Ops",
"Inpaint Segments": "Inpaint Segments",
"Combine and Paste": "Combine and Paste",
}
Reference in New Issue View Git Blame Copy Permalink