nodes/wmpmiles-comfyui-some-image-processing-stuff
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit.

Browse Source
This commit is contained in:
William Miles
2024-11-23 12:10:51 +11:00
commit 3f0cd9bfad
6 changed files with 1949 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -0,0 +1 @@
__pycache__
__init__.py
+9
View File
@@ -0,0 +1,9 @@
from .resampling import NODE_CLASS_MAPPINGS as resampling
from .inpainting import NODE_CLASS_MAPPINGS as inpainting
from .color_grading import NODE_CLASS_MAPPINGS as color_grading
NODE_CLASS_MAPPINGS = {
**resampling,
**inpainting,
**color_grading,
}
color_grading.py
+91
View File
@@ -0,0 +1,91 @@
import torch
from torch import Tensor
from math import log2
from . import gtfu
def _color_grade(
image: Tensor,
red: float,
green: float,
blue: float,
exposure: float,
saturation: float,
contrast: float,
tonemap: str,
) -> tuple[Tensor]:
linear = gtfu.colorspace_srgb_linear_from_gamma(image)
exposure_mult = 2 ** exposure
color_mult = torch.tensor([red, green, blue]).reshape(1, 1, 1, 3)
ec_mult = exposure_mult * color_mult
ec_graded = ec_mult * linear
ec_luminosity = gtfu.convert_luminance_from_linear_srgb(ec_graded, 3)
saturation_graded = (ec_luminosity + saturation * (ec_graded - ec_luminosity)).clamp(min=0)
sg_log2 = gtfu.colorspace_log2_from_linear(saturation_graded)
middle_log2 = log2(0.18)
contrast_graded_log2 = middle_log2 + (sg_log2 - middle_log2) * contrast
contrast_graded = (gtfu.colorspace_linear_from_log2(contrast_graded_log2)).clamp(min=0)
cg_luminance = gtfu.convert_luminance_from_linear_srgb(contrast_graded, 3)
cg_whitepoint = gtfu.util_tensor_max(cg_luminance, (1, 2, 3))
match tonemap:
case "Reinhard":
tonemapped = gtfu.tonemap_reinhard_extended_luminance(contrast_graded, cg_luminance, cg_whitepoint)
case "Reinhard-Jodie":
tonemapped = gtfu.tonemap_reinhard_jodie_extended(contrast_graded, cg_luminance, cg_whitepoint)
case "Uncharted 2":
tonemapped = gtfu.tonemap_uncharted_2(contrast_graded)
case "ACES":
tonemapped = gtfu.tonemap_aces(contrast_graded, 3)
tonemap_gamma = gtfu.colorspace_srgb_gamma_from_linear(tonemapped)
quantized = gtfu.filter_quantize(tonemap_gamma, 256, gtfu.RoundingMode.ROUND)
return quantized
# NODE
class ColorGrading:
@staticmethod
def INPUT_TYPES():
return {"required": {
"image": ("IMAGE", {}),
"red": ("FLOAT", {"default": 1, "step": 0.01}),
"green": ("FLOAT", {"default": 1, "step": 0.01}),
"blue": ("FLOAT", {"default": 1, "step": 0.01}),
"exposure_bias": ("FLOAT", {"default": 0, "step": 0.01, "min": -1000}),
"saturation": ("FLOAT", {"default": 1, "step": 0.01}),
"contrast": ("FLOAT", {"default": 1, "step": 0.01}),
"tonemap": (["Reinhard", "Reinhard-Jodie", "Uncharted 2", "ACES"], {})
}}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("image", )
CATEGORY = "inpaint"
FUNCTION = "f"
@staticmethod
def f(
image: Tensor,
red: float,
green: float,
blue: float,
exposure_bias: float,
saturation: float,
contrast: float,
tonemap: str,
) -> tuple[Tensor]:
graded = _color_grade(image, red, green, blue, exposure_bias, saturation, contrast, tonemap)
return (graded, )
NODE_CLASS_MAPPINGS = {
"Color Grading": ColorGrading,
}
__all__ = ["NODE_CLASS_MAPPINGS"]
gtfu.py
+1291
View File
File diff suppressed because it is too large Load Diff
inpainting.py
+254
View File
@@ -0,0 +1,254 @@
import torch
from torch import Tensor
from . import gtfu
from .resampling import Scaler, Resampler
class MaskCropInpaintPre:
@staticmethod
def INPUT_TYPES():
return {"required": {
"image": ("IMAGE", {}),
"mask": ("MASK", {}),
"resampler": ("RESAMPLER", {}),
"scaler": ("SCALER", {}),
"square": ("BOOLEAN", {"default": True}),
"area_mult": ("FLOAT", {"default": 1, "min": 1, "step": 0.1}),
}}
RETURN_TYPES = ("IMAGE", "MASK", "MASK", "CONTEXT")
RETURN_NAMES = ("image", "mask", "latent_mask", "context")
CATEGORY = "inpaint"
FUNCTION = "f"
@staticmethod
def f(image: Tensor, mask: Tensor, resampler: Resampler, scaler: Scaler, square: bool, area_mult: float) -> tuple[Tensor, Tensor, Tensor, tuple]:
if mask.shape[0] != 1:
raise ValueError("Mask must have a batch size of 1.")
if image.shape[1:3] != mask.shape[1:3]:
raise ValueError("Mask and image are not the same resolution.")
# Create bbox from mask
binary_mask = mask.to(torch.bool)
bbox = gtfu.bbox_from_2d_binary(binary_mask[0])
if square:
bbox = gtfu.bbox_outer_square(bbox)
bbox = gtfu.bbox_expand_area(bbox, area_mult)
# Crop image and mask to bbox
cropped_image = gtfu.transform_crop_to_bbox(image, bbox, 1, 2)
cropped_mask = gtfu.transform_crop_to_bbox(mask, bbox, 1, 2)
cropped_res = tuple(cropped_image.shape[1:3])
# Rescale image to chosen size
inpaint_res = scaler(cropped_res)
cropped_image_linear = gtfu.colorspace_srgb_linear_from_gamma(cropped_image)
rescaled_image_linear = resampler(cropped_image_linear, inpaint_res, (1, 2))
rescaled_image = gtfu.colorspace_srgb_gamma_from_linear(rescaled_image_linear)
rescaled_mask = gtfu.resample_nearest_neighbor_2d(cropped_mask, inpaint_res, (1, 2))
# Pad image to be a multiple of 8 in height and width
padded_res = tuple((gtfu.util_round_up_to_mult_of(x, 8) for x in inpaint_res))
padding = tuple(((0, y - x) for x, y in zip(inpaint_res, padded_res)))
padded_image = gtfu.transform_pad_dim2_reflect(rescaled_image, (1, 2), padding)
padded_mask = gtfu.transform_pad_dim2_zero(rescaled_mask, (1, 2), padding)
# Rescale mask to the correct size for latent masking
latent_mask_res = tuple((x // 8 for x in padded_res))
latent_mask = gtfu.resample_nearest_neighbor_2d(cropped_mask, latent_mask_res, (1, 2))
# Store context needed for post inpainting compositing
context = ("mask_crop_inpaint", image, bbox, inpaint_res, cropped_res)
return (padded_image, padded_mask, latent_mask, context)
class MaskCropInpaintPost:
@staticmethod
def INPUT_TYPES():
return {"required": {
"image": ("IMAGE", {}),
"mask": ("MASK", {}),
"context": ("CONTEXT", {}),
"resampler": ("RESAMPLER", {}),
}}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("image", )
CATEGORY = "inpaint"
FUNCTION = "f"
@staticmethod
def f(image: Tensor, mask: Tensor, context: tuple, resampler: Resampler) -> tuple[Tensor]:
if context[0] != "mask_crop_inpaint":
raise ValueError("Invalid context.")
original_image, bbox, inpaint_res, cropped_res = context[1:]
if original_image.shape[0] != image.shape[0]:
raise ValueError("Image must have same batch size as source.")
if mask.shape[1:3] != original_image.shape[1:3]:
raise ValueError("Mask must have same resolution as original image.")
# Prepare inpainted, cropped section
unpadded = image[:, :inpaint_res[0], :inpaint_res[1], :]
unpadded_linear = gtfu.colorspace_srgb_linear_from_gamma(unpadded)
resampled_linear = resampler(unpadded_linear, cropped_res, (1, 2))
uncropped_linear = gtfu.transform_uncrop_from_bbox(resampled_linear, bbox, 1, 2)
# Composite
original_linear = gtfu.colorspace_srgb_linear_from_gamma(original_image)
composite_mask = mask.unsqueeze(3).expand(*original_image.shape)
masked_original = (1 - composite_mask) * original_linear
masked_inpainted = composite_mask * uncropped_linear
composited_linear = masked_original + masked_inpainted
composited = gtfu.colorspace_srgb_gamma_from_linear(composited_linear)
# We quantize to prevent colorspace conversion artifacts
composited_quantized = gtfu.filter_quantize(composited, 256, gtfu.RoundingMode.ROUND)
return (composited_quantized, )
class LatentZeroMask:
@staticmethod
def INPUT_TYPES():
return {"required": {
"latent": ("LATENT", ), "mask": ("MASK", )
}}
RETURN_NAMES = ("latent", )
RETURN_TYPES = ("LATENT", )
CATEGORY = "inpaint"
FUNCTION = "f"
@staticmethod
def f(latent: dict[str, Tensor], mask: Tensor) -> tuple[dict[str, Tensor]]:
if latent["samples"].shape[0] != mask.shape[0]:
raise ValueError("Latent and mask batch size must match.")
if latent["samples"].shape[2:4] != mask.shape[1:3]:
raise ValueError("Latent and mask must have same resolution.")
new_latent = latent.copy()
samples = latent["samples"].clone()
mask_reshaped = mask.unsqueeze(1).expand(*samples.shape)
samples *= (1.0 - mask)
new_latent["samples"] = samples
return (new_latent, )
class MaskCropPre:
@staticmethod
def INPUT_TYPES():
return {"required": {
"image": ("IMAGE", {}),
"mask": ("MASK", {}),
"square": ("BOOLEAN", {"default": True}),
"area_mult": ("FLOAT", {"default": 1, "min": 1, "step": 0.1}),
}}
RETURN_TYPES = ("IMAGE", "MASK", "CONTEXT")
RETURN_NAMES = ("image", "mask", "context")
CATEGORY = "inpaint"
FUNCTION = "f"
@staticmethod
def f(image: Tensor, mask: Tensor, square: bool, area_mult: float) -> tuple[Tensor, Tensor, tuple]:
if mask.shape[0] != 1:
raise ValueError("Mask must have a batch size of 1.")
if image.shape[1:3] != mask.shape[1:3]:
raise ValueError("Mask and image are not the same resolution.")
# Create bbox from mask
binary_mask = mask.to(torch.bool)
bbox = gtfu.bbox_from_2d_binary(binary_mask[0])
if square:
bbox = gtfu.bbox_outer_square(bbox)
bbox = gtfu.bbox_expand_area(bbox, area_mult)
# Crop image and mask to bbox
cropped_image = gtfu.transform_crop_to_bbox(image, bbox, 1, 2)
cropped_mask = gtfu.transform_crop_to_bbox(mask, bbox, 1, 2)
cropped_res = tuple(cropped_image.shape[1:3])
context = ("mask_crop", image, bbox)
return (cropped_image, cropped_mask, context)
class MaskCropPost:
@staticmethod
def INPUT_TYPES():
return {"required": {
"image": ("IMAGE", {}),
"context": ("CONTEXT", {}),
}}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("image", )
CATEGORY = "inpaint"
FUNCTION = "f"
@staticmethod
def f(image: Tensor, context: tuple) -> tuple[Tensor]:
if context[0] != "mask_crop":
raise ValueError("Invalid context.")
original_image, bbox = context[1:]
if (bbox.height, bbox.width) != image.shape[1:3]:
raise ValueError("Image must be same resolution as the original crop.")
if image.shape[0] != original_image.shape[0]:
raise ValueError("Image must have same batch size as the original image.")
if image.shape[3] != original_image.shape[3]:
raise ValueError("Image must have same channel count as the original image.")
pasted = original_image.clone()
pasted[:, bbox.up:bbox.down_offset, bbox.left:bbox.right_offset, :] = image
return (pasted, )
class BlurMask:
@staticmethod
def INPUT_TYPES():
return {"required": {
"mask": ("MASK", ),
"blur": ("FLOAT", {"default": 1, "min": 0, "step": 0.01}),
"expand": ("INT", {"default": 0, "min": 0}),
}}
RETURN_TYPES = ("MASK", "IMAGE")
RETURN_NAMES = ("mask", "mask_image")
CATEGORY = "inpaint"
FUNCTION = "f"
@staticmethod
def f(mask: Tensor, blur: float, expand: int) -> tuple[Tensor]:
blurred_mask = blur_expand_mask(mask, blur, expand)
mask_image = blurred_mask.unsqueeze(3).expand(*blurred_mask.shape, 3)
return (blurred_mask, mask_image)
NODE_CLASS_MAPPINGS = {
"Mask-Crop Inpaint | Pre": MaskCropInpaintPre,
"Mask-Crop Inpaint | Post": MaskCropInpaintPost,
"Mask-Crop | Pre": MaskCropPre,
"Mask-Crop | Post": MaskCropPost,
"LatentZeroMask": LatentZeroMask,
"Blur Mask": BlurMask,
}
__all__ = ["NODE_CLASS_MAPPINGS"]
def blur_expand_mask(mask: Tensor, sigma: float, dilation_radius: int) -> Tensor:
if dilation_radius > 0:
dilated_mask = gtfu.filter_morpho_dilate(mask, dilation_radius, (1, 2))
else:
dilated_mask = mask
if sigma > 0:
gaussian_radius = int(gtfu.special_gaussian_area_radius(0.99, sigma))
kernel_xs = torch.arange(-gaussian_radius, gaussian_radius + 1).to(torch.float)
gaussian_kernel_1d = gtfu.special_gaussian(kernel_xs, sigma)
blurred_mask = gtfu.filter_convolve_2d_separable(dilated_mask, (gaussian_kernel_1d, gaussian_kernel_1d), (1, 2))
else:
blurred_mask = dilated_mask
return blurred_mask
resampling.py
+303
View File
@@ -0,0 +1,303 @@
import torch
from . import gtfu
from torch import Tensor
from typing import TypeAlias, Callable
from math import sqrt
from copy import copy
Resampler: TypeAlias = Callable[[Tensor, tuple[int, int], tuple[int, int]], Tensor]
Scaler: TypeAlias = Callable[[tuple[int, int]], tuple[int, int]]
class ResampleImage:
@staticmethod
def INPUT_TYPES():
return {"required": {
"image": ("IMAGE", {}),
"scaler": ("SCALER", {}),
"resampler": ("RESAMPLER", {}),
}}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("image", )
CATEGORY = "inpaint/resampling"
FUNCTION = "f"
@staticmethod
def f(image: Tensor, scaler: Scaler, resampler: Resampler) -> tuple[Tensor]:
_, h, w, _ = image.shape
new_res = scaler((h, w))
linear = gtfu.colorspace_srgb_linear_from_gamma(image)
resampled_linear = resampler(linear, new_res, (1, 2))
resampled = gtfu.colorspace_srgb_gamma_from_linear(resampled_linear)
return (resampled, )
class ResampleMask:
@staticmethod
def INPUT_TYPES():
return {"required": {
"mask": ("MASK", {}),
"scaler": ("SCALER", {}),
"resampler": ("RESAMPLER", {}),
}}
RETURN_TYPES = ("MASK", )
RETURN_NAMES = ("mask", )
CATEGORY = "inpaint/resampling"
FUNCTION = "f"
@staticmethod
def f(mask: Tensor, scaler: Scaler, resampler: Resampler) -> tuple[Tensor]:
_, h, w = mask.shape
new_res = scaler((h, w))
resampled = resampler(mask, new_res, (1, 2))
return (resampled, )
class ResampleLatent:
@staticmethod
def INPUT_TYPES():
return {"required": {
"latent": ("LATENT", {}),
"scaler": ("SCALER", {}),
"resampler": ("RESAMPLER", {}),
}}
RETURN_TYPES = ("LATENT", )
RETURN_NAMES = ("latent", )
CATEGORY = "inpaint/resampling"
FUNCTION = "f"
@staticmethod
def f(latent: dict[str, Tensor], scaler: Scaler, resampler: Resampler) -> tuple[Tensor]:
samples = latent["samples"]
_, _, h, w = samples.shape
new_res = scaler((h, w))
resampled = resampler(samples, new_res, (2, 3))
new_latent = copy(latent)
new_latent["samples"] = resampled
return (resampled, )
class ResamplerBase:
RETURN_TYPES = ("RESAMPLER", )
RETURN_NAMES = ("resampler", )
CATEGORY = "inpaint/resampling"
FUNCTION = "f"
class ResamplerNearestNeighbor(ResamplerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {}}
@staticmethod
def f() -> tuple[Resampler]:
def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
resampled = gtfu.resample_nearest_neighbor_2d(tensor, resolution, dims)
return resampled
return (resampler, )
class ResamplerTriangle(ResamplerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {
"radius": ("INT", {"min": 1, "step": 1, "default": 1})
}}
@staticmethod
def f(radius: int) -> tuple[Resampler]:
def filter(x: Tensor) -> Tensor:
return gtfu.window_triangle(x, radius)
def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
resampled = gtfu.resample_filter_2d_separable(tensor, resolution, radius, (filter, filter), dims)
return resampled
return (resampler, )
class ResamplerLanczos(ResamplerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {
"radius": ("INT", {"min": 1, "step": 1, "default": 3})
}}
@staticmethod
def f(radius: int) -> tuple[Resampler]:
def filter(x: Tensor) -> Tensor:
return torch.sinc(x) * gtfu.window_lanczos(x, radius)
def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
resampled = gtfu.resample_filter_2d_separable(tensor, resolution, radius, (filter, filter), dims)
return resampled
return (resampler, )
class ResamplerMitchellNetravali(ResamplerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {
"b": ("FLOAT", {"step": 0.01, "default": 0.33}),
"c": ("FLOAT", {"step": 0.01, "default": 0.33}),
}}
@staticmethod
def f(b: float, c: float) -> tuple[Resampler]:
radius = gtfu.window_mitchell_netravali_radius()
def filter(x: Tensor) -> Tensor:
return gtfu.window_mitchell_netravali(x, b, c)
def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
resampled = gtfu.resample_filter_2d_separable(tensor, resolution, radius, (filter, filter), dims)
return resampled
return (resampler, )
class ResamplerArea(ResamplerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {}}
@staticmethod
def f() -> tuple[Resampler]:
radius = gtfu.window_area_radius()
def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
_, h, w, _ = tensor.shape
H, W = resolution
def filter_h(x: Tensor) -> Tensor:
return gtfu.window_area(x, h, H)
def filter_w(x: Tensor) -> Tensor:
return gtfu.window_area(x, w, W)
resampled = gtfu.resample_filter_2d_separable(tensor, resolution, radius, (filter_h, filter_w), dims)
return resampled
return (resampler, )
class ResamplerJincLanczos(ResamplerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {
"radius": ("INT", {"min": 1, "step": 1, "default": 3})
}}
@staticmethod
def f(radius: int) -> tuple[Resampler]:
def filter(x: Tensor) -> Tensor:
return gtfu.special_jinc(x) * gtfu.window_lanczos(x, radius)
def resampler(tensor: Tensor, resolution: tuple[int, int], dims: tuple[int, int]) -> Tensor:
resampled = gtfu.resample_filter_2d(tensor, resolution, radius, filter, dims)
return resampled
return (resampler, )
class ScalerBase:
RETURN_TYPES = ("SCALER", )
RETURN_NAMES = ("scaler", )
CATEGORY = "inpaint/resampling"
FUNCTION = "f"
class ScalerSide(ScalerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {
"side_multiplier": ("FLOAT", {"default": 1, "min": 0.001, "step": 0.001}),
}}
@staticmethod
def f(side_multiplier: float) -> tuple[Scaler]:
def scaler(res: tuple[int, int]) -> tuple[int, int]:
scaled = tuple((int(round(x * side_multiplier)) for x in res))
return scaled
return (scaler, )
class ScalerArea(ScalerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {
"area_multiplier": ("FLOAT", {"default": 1, "min": 0.001, "step": 0.001}),
}}
@staticmethod
def f(area_multiplier: float) -> tuple[Scaler]:
side_multiplier = sqrt(area_multiplier)
def scaler(res: tuple[int, int]) -> tuple[int, int]:
scaled = tuple((int(round(x * side_multiplier)) for x in res))
return scaled
return (scaler, )
class ScalerUnlinked(ScalerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {
"width_multiplier": ("FLOAT", {"default": 1, "min": 0.001, "step": 0.001}),
"height_multiplier": ("FLOAT", {"default": 1, "min": 0.001, "step": 0.001}),
}}
@staticmethod
def f(width_multiplier: float, height_multiplier: float) -> tuple[Scaler]:
def scaler(res: tuple[int, int]) -> tuple[int, int]:
h, w = res
scaled = (height_multiplier * h, width_multiplier * w)
return scaled
return (scaler, )
class ScalerPixelDeltas(ScalerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {
"width_delta": ("INT", {"default": 0}),
"height_delta": ("INT", {"default": 0}),
}}
@staticmethod
def f(width_multiplier: float, height_multiplier: float) -> tuple[Scaler]:
def scaler(res: tuple[int, int]) -> tuple[int, int]:
h, w = res
scaled = (h + height_delta, w + width_delta)
return scaled
return (scaler, )
class ScalerMegapixels(ScalerBase):
@staticmethod
def INPUT_TYPES():
return {"required": {
"megapixels": ("FLOAT", {"default": 1.0, "min": 0.001, "step": 0.001}),
}}
@staticmethod
def f(megapixels: float) -> tuple[Scaler]:
def scaler(res: tuple[int, int]) -> tuple[int, int]:
h, w = res
cur_megapixels = h * w / 1_000_000
area_mult = megapixels / cur_megapixels
side_mult = sqrt(area_mult)
scaled = tuple((int(round(side_mult * x)) for x in res))
return scaled
return (scaler, )
NODE_CLASS_MAPPINGS = {
"Resample Image" : ResampleImage,
"Resample Mask" : ResampleMask,
"Resample Latent" : ResampleLatent,
"Resampler | Nearest-Neighbor" : ResamplerNearestNeighbor,
"Resampler | Triangle" : ResamplerTriangle,
"Resampler | Lanczos" : ResamplerLanczos,
"Resampler | Jinc-Lanczos" : ResamplerJincLanczos,
"Resampler | Mitchell-Netravali" : ResamplerMitchellNetravali,
"Resampler | Area" : ResamplerArea,
"Scaler | Side" : ScalerSide,
"Scaler | Area" : ScalerArea,
"Scaler | Sides Unlinked" : ScalerUnlinked,
"Scaler | Pixel Deltas" : ScalerPixelDeltas,
"Scaler | Megapixels" : ScalerMegapixels,
}
__all__ = ["NODE_CLASS_MAPPINGS"]