67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
import torch
|
|
|
|
class ColorTint:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.1,
|
|
"max": 1.0,
|
|
"step": 0.1
|
|
}),
|
|
"mode": (["sepia", "red", "green", "blue", "cyan", "magenta", "yellow", "purple", "orange", "warm", "cool", "lime", "navy", "vintage", "rose", "teal", "maroon", "peach", "lavender", "olive"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "color_tint"
|
|
|
|
CATEGORY = "postprocessing/Color Adjustments"
|
|
|
|
def color_tint(self, image: torch.Tensor, strength: float, mode: str = "sepia"):
|
|
if strength == 0:
|
|
return (image,)
|
|
|
|
sepia_weights = torch.tensor([0.2989, 0.5870, 0.1140]).view(1, 1, 1, 3).to(image.device)
|
|
|
|
mode_filters = {
|
|
"sepia": torch.tensor([1.0, 0.8, 0.6]),
|
|
"red": torch.tensor([1.0, 0.6, 0.6]),
|
|
"green": torch.tensor([0.6, 1.0, 0.6]),
|
|
"blue": torch.tensor([0.6, 0.8, 1.0]),
|
|
"cyan": torch.tensor([0.6, 1.0, 1.0]),
|
|
"magenta": torch.tensor([1.0, 0.6, 1.0]),
|
|
"yellow": torch.tensor([1.0, 1.0, 0.6]),
|
|
"purple": torch.tensor([0.8, 0.6, 1.0]),
|
|
"orange": torch.tensor([1.0, 0.7, 0.3]),
|
|
"warm": torch.tensor([1.0, 0.9, 0.7]),
|
|
"cool": torch.tensor([0.7, 0.9, 1.0]),
|
|
"lime": torch.tensor([0.7, 1.0, 0.3]),
|
|
"navy": torch.tensor([0.3, 0.4, 0.7]),
|
|
"vintage": torch.tensor([0.9, 0.85, 0.7]),
|
|
"rose": torch.tensor([1.0, 0.8, 0.9]),
|
|
"teal": torch.tensor([0.3, 0.8, 0.8]),
|
|
"maroon": torch.tensor([0.7, 0.3, 0.5]),
|
|
"peach": torch.tensor([1.0, 0.8, 0.6]),
|
|
"lavender": torch.tensor([0.8, 0.6, 1.0]),
|
|
"olive": torch.tensor([0.6, 0.7, 0.4]),
|
|
}
|
|
|
|
scale_filter = mode_filters[mode].view(1, 1, 1, 3).to(image.device)
|
|
|
|
grayscale = torch.sum(image * sepia_weights, dim=-1, keepdim=True)
|
|
tinted = grayscale * scale_filter
|
|
|
|
result = tinted * strength + image * (1 - strength)
|
|
return (result,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ColorTint": ColorTint
|
|
}
|