33 lines
1015 B
Python
33 lines
1015 B
Python
|
import torch
|
||
|
|
import torchvision.transforms.functional as tf
|
||
|
|
import torchvision.transforms.v2 as v2
|
||
|
|
from .mmaker_color_enhance_core import color_enhance
|
||
|
|
|
||
|
|
class ColorEnhanceComfyNode:
|
||
|
|
def __init__(self):
|
||
|
|
pass
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def INPUT_TYPES(s):
|
||
|
|
return {
|
||
|
|
"required": {
|
||
|
|
"image": ("IMAGE",),
|
||
|
|
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
RETURN_TYPES = ("IMAGE",)
|
||
|
|
FUNCTION = "apply_color_enhance"
|
||
|
|
CATEGORY = "postprocessing/Effects"
|
||
|
|
|
||
|
|
def apply_color_enhance(self, image: torch.Tensor, strength: float):
|
||
|
|
images = []
|
||
|
|
|
||
|
|
for img in image:
|
||
|
|
edited_image = v2.ToDtype(dtype=torch.uint8, scale=True)(img).squeeze()
|
||
|
|
edited_image = color_enhance(edited_image.detach().cpu().numpy(), strength)
|
||
|
|
edited_image = tf.to_tensor(edited_image)
|
||
|
|
images.append(edited_image)
|
||
|
|
|
||
|
|
return (torch.stack(images).permute(0, 2, 3, 1),)
|