898 lines
28 KiB
Python
898 lines
28 KiB
Python
import os
|
|
import sys
|
|
import copy
|
|
import torch
|
|
import torchvision.transforms
|
|
import numpy as np
|
|
import cv2
|
|
from pymatting import *
|
|
|
|
try:
|
|
from cv2.ximgproc import guidedFilter
|
|
except ImportError:
|
|
print("\033[33mUnable to import guidedFilter, make sure you have only opencv-contrib-python or run the import_error_install.bat script\033[m")
|
|
|
|
import comfy.model_management
|
|
|
|
MAX_RESOLUTION=8192
|
|
|
|
class AlphaClean:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"radius": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1
|
|
}),
|
|
"fill_holes": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 16,
|
|
"step": 1
|
|
}),
|
|
"white_threshold": ("FLOAT", {
|
|
"default": 0.9,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"extra_clip": ("FLOAT", {
|
|
"default": 0.98,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "alpha_clean"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def alpha_clean(self, images: torch.Tensor, radius: int, fill_holes: int, white_threshold: float, extra_clip: float):
|
|
|
|
d = radius * 2 + 1
|
|
i_dup = copy.deepcopy(images.cpu().numpy())
|
|
|
|
for index, image in enumerate(i_dup):
|
|
|
|
cleaned = cv2.bilateralFilter(image, 9, 0.05, 8)
|
|
|
|
alpha = np.clip((image - white_threshold) / (1 - white_threshold), 0, 1)
|
|
rgb = image * alpha
|
|
|
|
alpha = cv2.GaussianBlur(alpha, (d,d), 0) * 0.99 + np.average(alpha) * 0.01
|
|
rgb = cv2.GaussianBlur(rgb, (d,d), 0) * 0.99 + np.average(rgb) * 0.01
|
|
|
|
rgb = rgb / np.clip(alpha, 0.00001, 1)
|
|
rgb = rgb * extra_clip
|
|
|
|
cleaned = np.clip(cleaned / rgb, 0, 1)
|
|
|
|
if fill_holes > 0:
|
|
fD = fill_holes * 2 + 1
|
|
gamma = cleaned * cleaned
|
|
kD = np.ones((fD, fD), np.uint8)
|
|
kE = np.ones((fD + 2, fD + 2), np.uint8)
|
|
gamma = cv2.dilate(gamma, kD, iterations=1)
|
|
gamma = cv2.erode(gamma, kE, iterations=1)
|
|
gamma = cv2.GaussianBlur(gamma, (fD, fD), 0)
|
|
cleaned = np.maximum(cleaned, gamma)
|
|
|
|
i_dup[index] = cleaned
|
|
|
|
return (torch.from_numpy(i_dup),)
|
|
|
|
class AlphaMatte:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"alpha_trimap": ("IMAGE",),
|
|
"preblur": ("INT", {
|
|
"default": 8,
|
|
"min": 0,
|
|
"max": 256,
|
|
"step": 1
|
|
}),
|
|
"blackpoint": ("FLOAT", {
|
|
"default": 0.01,
|
|
"min": 0.0,
|
|
"max": 0.99,
|
|
"step": 0.01
|
|
}),
|
|
"whitepoint": ("FLOAT", {
|
|
"default": 0.99,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"max_iterations": ("INT", {
|
|
"default": 1000,
|
|
"min": 100,
|
|
"max": 10000,
|
|
"step": 100
|
|
}),
|
|
"estimate_fg": (["true", "false"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE",)
|
|
RETURN_NAMES = ("alpha", "fg", "bg",)
|
|
FUNCTION = "alpha_matte"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def alpha_matte(self, images, alpha_trimap, preblur, blackpoint, whitepoint, max_iterations, estimate_fg):
|
|
|
|
d = preblur * 2 + 1
|
|
|
|
i_dup = copy.deepcopy(images.cpu().numpy().astype(np.float64))
|
|
a_dup = copy.deepcopy(alpha_trimap.cpu().numpy().astype(np.float64))
|
|
fg = copy.deepcopy(images.cpu().numpy().astype(np.float64))
|
|
bg = copy.deepcopy(images.cpu().numpy().astype(np.float64))
|
|
|
|
|
|
for index, image in enumerate(i_dup):
|
|
trimap = a_dup[index][:,:,0] # convert to single channel
|
|
if preblur > 0:
|
|
trimap = cv2.GaussianBlur(trimap, (d, d), 0)
|
|
trimap = fix_trimap(trimap, blackpoint, whitepoint)
|
|
|
|
alpha = estimate_alpha_cf(image, trimap, laplacian_kwargs={"epsilon": 1e-6}, cg_kwargs={"maxiter":max_iterations})
|
|
|
|
if estimate_fg == "true":
|
|
fg[index], bg[index] = estimate_foreground_ml(image, alpha, return_background=True)
|
|
|
|
a_dup[index] = np.stack([alpha, alpha, alpha], axis = -1) # convert back to rgb
|
|
|
|
return (
|
|
torch.from_numpy(a_dup.astype(np.float32)), # alpha
|
|
torch.from_numpy(fg.astype(np.float32)), # fg
|
|
torch.from_numpy(bg.astype(np.float32)), # bg
|
|
)
|
|
|
|
class BlurImageFast:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"radius_x": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
"radius_y": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "blur_image"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def blur_image(self, images, radius_x, radius_y):
|
|
|
|
if radius_x + radius_y == 0:
|
|
return (images,)
|
|
|
|
dx = radius_x * 2 + 1
|
|
dy = radius_y * 2 + 1
|
|
|
|
dup = copy.deepcopy(images.cpu().numpy())
|
|
|
|
for index, image in enumerate(dup):
|
|
dup[index] = cv2.GaussianBlur(image, (dx, dy), 0)
|
|
|
|
return (torch.from_numpy(dup),)
|
|
|
|
class BlurMaskFast:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"masks": ("MASK",),
|
|
"radius_x": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
"radius_y": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "blur_mask"
|
|
|
|
CATEGORY = "mask/filters"
|
|
|
|
def blur_mask(self, masks, radius_x, radius_y):
|
|
|
|
if radius_x + radius_y == 0:
|
|
return (masks,)
|
|
|
|
dx = radius_x * 2 + 1
|
|
dy = radius_y * 2 + 1
|
|
|
|
dup = copy.deepcopy(masks.cpu().numpy())
|
|
|
|
for index, mask in enumerate(dup):
|
|
dup[index] = cv2.GaussianBlur(mask, (dx, dy), 0)
|
|
|
|
return (torch.from_numpy(dup),)
|
|
|
|
class DilateErodeMask:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"masks": ("MASK",),
|
|
"radius": ("INT", {
|
|
"default": 0,
|
|
"min": -1023,
|
|
"max": 1023,
|
|
"step": 1
|
|
}),
|
|
"shape": (["box", "circle"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "dilate_mask"
|
|
|
|
CATEGORY = "mask/filters"
|
|
|
|
def dilate_mask(self, masks, radius, shape):
|
|
|
|
if radius == 0:
|
|
return (masks,)
|
|
|
|
s = abs(radius)
|
|
d = s * 2 + 1
|
|
k = np.zeros((d, d), np.uint8)
|
|
if shape == "circle":
|
|
k = cv2.circle(k, (s,s), s, 1, -1)
|
|
else:
|
|
k += 1
|
|
|
|
dup = copy.deepcopy(masks.cpu().numpy())
|
|
|
|
for index, mask in enumerate(dup):
|
|
if radius > 0:
|
|
dup[index] = cv2.dilate(mask, k, iterations=1)
|
|
else:
|
|
dup[index] = cv2.erode(mask, k, iterations=1)
|
|
|
|
return (torch.from_numpy(dup),)
|
|
|
|
class EnhanceDetail:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"filter_radius": ("INT", {
|
|
"default": 2,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1
|
|
}),
|
|
"sigma": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.01,
|
|
"max": 100.0,
|
|
"step": 0.01
|
|
}),
|
|
"denoise": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.01
|
|
}),
|
|
"detail_mult": ("FLOAT", {
|
|
"default": 2.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.1
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "enhance"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def enhance(self, images: torch.Tensor, filter_radius: int, sigma: float, denoise: float, detail_mult: float):
|
|
|
|
if filter_radius == 0:
|
|
return (images,)
|
|
|
|
d = filter_radius * 2 + 1
|
|
s = sigma / 10
|
|
n = denoise / 10
|
|
|
|
dup = copy.deepcopy(images.cpu().numpy())
|
|
|
|
for index, image in enumerate(dup):
|
|
imgB = image
|
|
if denoise>0.0:
|
|
imgB = cv2.bilateralFilter(image, d, n, d)
|
|
|
|
imgG = np.clip(guidedFilter(image, image, d, s), 0.001, 1)
|
|
|
|
details = (imgB/imgG - 1) * detail_mult + 1
|
|
dup[index] = np.clip(details*imgG - imgB + image, 0, 1)
|
|
|
|
return (torch.from_numpy(dup),)
|
|
|
|
class GuidedFilterAlpha:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"alpha": ("IMAGE",),
|
|
"filter_radius": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1
|
|
}),
|
|
"sigma": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "guided_filter_alpha"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def guided_filter_alpha(self, images: torch.Tensor, alpha: torch.Tensor, filter_radius: int, sigma: float):
|
|
|
|
d = filter_radius * 2 + 1
|
|
s = sigma / 10
|
|
|
|
i_dup = copy.deepcopy(images.cpu().numpy())
|
|
a_dup = copy.deepcopy(alpha.cpu().numpy())
|
|
|
|
for index, image in enumerate(i_dup):
|
|
alpha_work = a_dup[index]
|
|
i_dup[index] = guidedFilter(image, alpha_work, d, s)
|
|
|
|
return (torch.from_numpy(i_dup),)
|
|
|
|
class RemapRange:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"blackpoint": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"whitepoint": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "remap"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def remap(self, image: torch.Tensor, blackpoint: float, whitepoint: float):
|
|
|
|
bp = min(blackpoint, whitepoint - 0.001)
|
|
scale = 1 / (whitepoint - bp)
|
|
|
|
i_dup = copy.deepcopy(image.cpu().numpy())
|
|
i_dup = np.clip((i_dup - bp) * scale, 0.0, 1.0)
|
|
|
|
return (torch.from_numpy(i_dup),)
|
|
|
|
class ClampOutliers:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"latents": ("LATENT", ),
|
|
"std_dev": ("FLOAT", {"default": 3.0, "min": 0.1, "max": 100.0, "step": 0.1, "round": 0.1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "clamp_outliers"
|
|
|
|
CATEGORY = "latent/filters"
|
|
|
|
def clamp_outliers(self, latents, std_dev):
|
|
latents_copy = copy.deepcopy(latents)
|
|
t = latents_copy["samples"]
|
|
|
|
for i, latent in enumerate(t):
|
|
for j, channel in enumerate(latent):
|
|
sd, mean = torch.std_mean(channel, dim=None)
|
|
t[i,j] = torch.clamp(channel, min = -sd * std_dev + mean, max = sd * std_dev + mean)
|
|
|
|
latents_copy["samples"] = t
|
|
return (latents_copy,)
|
|
|
|
class AdainLatent:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"latents": ("LATENT", ),
|
|
"reference": ("LATENT", ),
|
|
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "batch_normalize"
|
|
|
|
CATEGORY = "latent/filters"
|
|
|
|
def batch_normalize(self, latents, reference, factor):
|
|
latents_copy = copy.deepcopy(latents)
|
|
t = latents_copy["samples"] # [B x C x H x W]
|
|
|
|
t = t.movedim(0,1) # [C x B x H x W]
|
|
for c in range(t.size(0)):
|
|
for i in range(t.size(1)):
|
|
r_sd, r_mean = torch.std_mean(reference["samples"][i, c], dim=None) # index by original dim order
|
|
i_sd, i_mean = torch.std_mean(t[c, i], dim=None)
|
|
|
|
t[c, i] = ((t[c, i] - i_mean) / i_sd) * r_sd + r_mean
|
|
|
|
latents_copy["samples"] = torch.lerp(latents["samples"], t.movedim(1,0), factor) # [B x C x H x W]
|
|
return (latents_copy,)
|
|
|
|
class AdainImage:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", ),
|
|
"reference": ("IMAGE", ),
|
|
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "batch_normalize"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def batch_normalize(self, images, reference, factor):
|
|
t = copy.deepcopy(images) # [B x H x W x C]
|
|
|
|
t = t.movedim(-1,0) # [C x B x H x W]
|
|
for c in range(t.size(0)):
|
|
for i in range(t.size(1)):
|
|
r_sd, r_mean = torch.std_mean(reference[i, :, :, c], dim=None) # index by original dim order
|
|
i_sd, i_mean = torch.std_mean(t[c, i], dim=None)
|
|
|
|
t[c, i] = ((t[c, i] - i_mean) / i_sd) * r_sd + r_mean
|
|
|
|
t = torch.lerp(images, t.movedim(0,-1), factor) # [B x H x W x C]
|
|
return (t,)
|
|
|
|
class BatchNormalizeLatent:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"latents": ("LATENT", ),
|
|
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "batch_normalize"
|
|
|
|
CATEGORY = "latent/filters"
|
|
|
|
def batch_normalize(self, latents, factor):
|
|
latents_copy = copy.deepcopy(latents)
|
|
t = latents_copy["samples"] # [B x C x H x W]
|
|
|
|
t = t.movedim(0,1) # [C x B x H x W]
|
|
for c in range(t.size(0)):
|
|
c_sd, c_mean = torch.std_mean(t[c], dim=None)
|
|
|
|
for i in range(t.size(1)):
|
|
i_sd, i_mean = torch.std_mean(t[c, i], dim=None)
|
|
|
|
t[c, i] = (t[c, i] - i_mean) / i_sd
|
|
|
|
t[c] = t[c] * c_sd + c_mean
|
|
|
|
latents_copy["samples"] = torch.lerp(latents["samples"], t.movedim(1,0), factor) # [B x C x H x W]
|
|
return (latents_copy,)
|
|
|
|
class BatchNormalizeImage:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", ),
|
|
"factor": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01, "round": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "batch_normalize"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def batch_normalize(self, images, factor):
|
|
t = copy.deepcopy(images) # [B x H x W x C]
|
|
|
|
t = t.movedim(-1,0) # [C x B x H x W]
|
|
for c in range(t.size(0)):
|
|
c_sd, c_mean = torch.std_mean(t[c], dim=None)
|
|
|
|
for i in range(t.size(1)):
|
|
i_sd, i_mean = torch.std_mean(t[c, i], dim=None)
|
|
|
|
t[c, i] = (t[c, i] - i_mean) / i_sd
|
|
|
|
t[c] = t[c] * c_sd + c_mean
|
|
|
|
t = torch.lerp(images, t.movedim(0,-1), factor) # [B x H x W x C]
|
|
return (t,)
|
|
|
|
class DifferenceChecker:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images1": ("IMAGE", ),
|
|
"images2": ("IMAGE", ),
|
|
"multiplier": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 1000.0, "step": 0.01, "round": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "difference_checker"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def difference_checker(self, images1, images2, multiplier):
|
|
t = copy.deepcopy(images1)
|
|
t = torch.abs(images1 - images2) * multiplier
|
|
return (torch.clamp(t, min=0, max=1),)
|
|
|
|
class ImageConstant:
|
|
def __init__(self, device="cpu"):
|
|
self.device = device
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
|
|
"height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
|
"color": ("COLOR",),
|
|
}}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "generate"
|
|
|
|
CATEGORY = "image"
|
|
|
|
def generate(self, width, height, batch_size, color):
|
|
# rgb = color.lstrip('#')
|
|
rgb = tuple(int(color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))
|
|
|
|
r = torch.full([batch_size, height, width, 1], rgb[0]/255)
|
|
g = torch.full([batch_size, height, width, 1], rgb[1]/255)
|
|
b = torch.full([batch_size, height, width, 1], rgb[2]/255)
|
|
return (torch.cat((r, g, b), dim=-1), )
|
|
|
|
class OffsetLatentImage:
|
|
def __init__(self):
|
|
self.device = comfy.model_management.intermediate_device()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
|
|
"height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
|
"offset_0": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.1, "round": 0.1}),
|
|
"offset_1": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.1, "round": 0.1}),
|
|
"offset_2": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.1, "round": 0.1}),
|
|
"offset_3": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.1, "round": 0.1}),
|
|
}}
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "generate"
|
|
|
|
CATEGORY = "latent"
|
|
|
|
def generate(self, width, height, batch_size, offset_0, offset_1, offset_2, offset_3):
|
|
latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
|
|
latent[:,0,:,:] = offset_0
|
|
latent[:,1,:,:] = offset_1
|
|
latent[:,2,:,:] = offset_2
|
|
latent[:,3,:,:] = offset_3
|
|
return ({"samples":latent}, )
|
|
|
|
class LatentStats:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"latent": ("LATENT", ),}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("stats",)
|
|
FUNCTION = "notify"
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "utils"
|
|
|
|
def notify(self, latent):
|
|
latents = latent["samples"]
|
|
width, height = latents.size(3), latents.size(2)
|
|
|
|
text = ["",]
|
|
text[0] = f"batch size: {latents.size(0)}"
|
|
text.append(f"width: {width} ({width * 8})")
|
|
text.append(f"height: {height} ({height * 8})")
|
|
|
|
for i in range(4):
|
|
minimum = torch.min(latents[:,i,:,:]).item()
|
|
maximum = torch.max(latents[:,i,:,:]).item()
|
|
std_dev, mean = torch.std_mean(latents[:,i,:,:], dim=None)
|
|
|
|
text.append(f"c{i} mean: {mean:.1f} std_dev: {std_dev:.1f} min: {minimum:.1f} max: {maximum:.1f}")
|
|
|
|
|
|
printtext = "\033[36mLatent Stats:\033[m"
|
|
for t in text:
|
|
printtext += "\n " + t
|
|
|
|
returntext = ""
|
|
for i in range(len(text)):
|
|
if i > 0:
|
|
returntext += "\n"
|
|
returntext += text[i]
|
|
|
|
print(printtext)
|
|
return (returntext,)
|
|
|
|
def sRGBtoLinear(npArray):
|
|
less = npArray <= 0.0404482362771082
|
|
npArray[less] = npArray[less] / 12.92
|
|
npArray[~less] = np.power((npArray[~less] + 0.055) / 1.055, 2.4)
|
|
|
|
def linearToSRGB(npArray):
|
|
less = npArray <= 0.0031308
|
|
npArray[less] = npArray[less] * 12.92
|
|
npArray[~less] = np.power(npArray[~less], 1/2.4) * 1.055 - 0.055
|
|
|
|
def linearToTonemap(npArray):
|
|
more = npArray > 0.06
|
|
SLog3 = np.clip((np.log10((npArray + 0.01)/0.19) * 261.5 + 420) / 1023, 0, 1)
|
|
npArray[more] = np.power(1 / (1 + (1 / np.power(SLog3[more] / (1 - SLog3[more]), 1.7))), 1.7)
|
|
|
|
def tonemapToLinear(npArray):
|
|
more = npArray > 0.06
|
|
x = np.power(np.clip(npArray, 0, 1), 1/1.7)
|
|
ut = 1 / (1 + np.power((-1 / x) * (x - 1), 1/1.7))
|
|
npArray[more] = np.power(10, (ut[more] * 1023 - 420)/261.5) * 0.19 - 0.01
|
|
|
|
class Tonemap:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"input_mode": (["linear", "sRGB"],),
|
|
"output_mode": (["sRGB", "linear"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def apply(self, images, input_mode, output_mode):
|
|
t = images.detach().clone().cpu().numpy().astype(np.float32)
|
|
|
|
if input_mode == "sRGB":
|
|
sRGBtoLinear(t[:,:,:,:3])
|
|
|
|
linearToTonemap(t[:,:,:,:3])
|
|
|
|
if output_mode == "sRGB":
|
|
linearToSRGB(t[:,:,:,:3])
|
|
t = np.clip(t, 0, 1)
|
|
|
|
t = torch.from_numpy(t)
|
|
return (t,)
|
|
|
|
class UnTonemap:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"input_mode": (["sRGB", "linear"],),
|
|
"output_mode": (["linear", "sRGB"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def apply(self, images, input_mode, output_mode):
|
|
t = images.detach().clone().cpu().numpy().astype(np.float32)
|
|
|
|
if input_mode == "sRGB":
|
|
sRGBtoLinear(t[:,:,:,:3])
|
|
|
|
tonemapToLinear(t[:,:,:,:3])
|
|
|
|
if output_mode == "sRGB":
|
|
linearToSRGB(t[:,:,:,:3])
|
|
t = np.clip(t, 0, 1)
|
|
|
|
t = torch.from_numpy(t)
|
|
return (t,)
|
|
|
|
def exposure(npArray, stops):
|
|
more = npArray > 0
|
|
npArray[more] *= pow(2, stops)
|
|
|
|
class ExposureAdjust:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"stops": ("FLOAT", {"default": 0.0, "min": -100, "max": 100, "step": 0.01}),
|
|
"input_mode": (["sRGB", "linear"],),
|
|
"output_mode": (["sRGB", "linear"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply"
|
|
|
|
CATEGORY = "image/filters"
|
|
|
|
def apply(self, images, stops, input_mode, output_mode):
|
|
t = images.detach().clone().cpu().numpy().astype(np.float32)
|
|
|
|
if input_mode == "sRGB":
|
|
sRGBtoLinear(t[:,:,:,:3])
|
|
|
|
exposure(t[:,:,:,:3], stops)
|
|
|
|
if output_mode == "sRGB":
|
|
linearToSRGB(t[:,:,:,:3])
|
|
t = np.clip(t, 0, 1)
|
|
|
|
t = torch.from_numpy(t)
|
|
return (t,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AlphaClean": AlphaClean,
|
|
"AlphaMatte": AlphaMatte,
|
|
"BlurImageFast": BlurImageFast,
|
|
"BlurMaskFast": BlurMaskFast,
|
|
"DilateErodeMask": DilateErodeMask,
|
|
"EnhanceDetail": EnhanceDetail,
|
|
"GuidedFilterAlpha": GuidedFilterAlpha,
|
|
"RemapRange": RemapRange,
|
|
"ClampOutliers": ClampOutliers,
|
|
"AdainLatent": AdainLatent,
|
|
"AdainImage": AdainImage,
|
|
"BatchNormalizeLatent": BatchNormalizeLatent,
|
|
"BatchNormalizeImage": BatchNormalizeImage,
|
|
"DifferenceChecker": DifferenceChecker,
|
|
"ImageConstant": ImageConstant,
|
|
"OffsetLatentImage": OffsetLatentImage,
|
|
"LatentStats": LatentStats,
|
|
"Tonemap": Tonemap,
|
|
"UnTonemap": UnTonemap,
|
|
"ExposureAdjust": ExposureAdjust,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AlphaClean": "Alpha Clean",
|
|
"AlphaMatte": "Alpha Matte",
|
|
"BlurImageFast": "Blur Image (Fast)",
|
|
"BlurMaskFast": "Blur Mask (Fast)",
|
|
"DilateErodeMask": "Dilate/Erode Mask",
|
|
"EnhanceDetail": "Enhance Detail",
|
|
"GuidedFilterAlpha": "Guided Filter Alpha",
|
|
"RemapRange": "Remap Range",
|
|
"ClampOutliers": "Clamp Outliers",
|
|
"AdainLatent": "AdaIN (Latent)",
|
|
"AdainImage": "AdaIN (Image)",
|
|
"BatchNormalizeLatent": "Batch Normalize (Latent)",
|
|
"BatchNormalizeImage": "Batch Normalize (Image)",
|
|
"DifferenceChecker": "Difference Checker",
|
|
"ImageConstant": "Image Constant Color (RGB)",
|
|
"OffsetLatentImage": "Offset Latent Image",
|
|
"LatentStats": "Latent Stats",
|
|
"Tonemap": "Tonemap",
|
|
"UnTonemap": "UnTonemap",
|
|
"ExposureAdjust": "Exposure Adjust",
|
|
} |