diff --git a/README.md b/README.md index 7e12a2e..2af0f42 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # Color Enhance -Script for [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) to enhance colors. +Script for [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) and node for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) to enhance colors. This is the same algorithm GIMP/GEGL uses for color enhancement. The gist of this implementation is that it converts the color space to [CIELCh(ab)](https://en.wikipedia.org/wiki/CIELUV#Cylindrical_representation_(CIELCh)) and normalizes the chroma (or ["colorfulness"](https://en.wikipedia.org/wiki/Colorfulness)) component. Original source can be found in the link below. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6553283 --- /dev/null +++ b/__init__.py @@ -0,0 +1,9 @@ +from . import mmaker_color_enhance_comfyui + +NODE_CLASS_MAPPINGS = { + "MMakerColorEnhance": mmaker_color_enhance_comfyui.ColorEnhanceComfyNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MMakerColorEnhance": "Color Enhance", +} \ No newline at end of file diff --git a/mmaker_color_enhance_comfyui.py b/mmaker_color_enhance_comfyui.py new file mode 100644 index 0000000..6954740 --- /dev/null +++ b/mmaker_color_enhance_comfyui.py @@ -0,0 +1,32 @@ +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),) diff --git a/mmaker_color_enhance_core.py b/mmaker_color_enhance_core.py new file mode 100644 index 0000000..87b0522 --- /dev/null +++ b/mmaker_color_enhance_core.py @@ -0,0 +1,14 @@ +import numpy as np +import skimage.color +from PIL import Image +import imageio.core.util + +imageio.core.util._precision_warn = lambda *args, **kwargs: None + +def color_enhance(arr, strength: float = 1) -> Image.Image: + lch = skimage.color.lab2lch(lab=skimage.color.rgb2lab(rgb=np.array(arr, dtype=np.uint8))) + lch[:, :, 1] *= 100/(lerp(100, lch[:, :, 1].max(), strength)) # Normalize chroma component + return Image.fromarray(np.array(skimage.color.lab2rgb(lab=skimage.color.lch2lab(lch=lch)) * 255, dtype=np.uint8)) + +def lerp(a: float, b: float, t: float) -> float: + return (1 - t) * a + t * b \ No newline at end of file diff --git a/scripts/color_enhance.py b/scripts/color_enhance.py index 339567e..9d920b6 100644 --- a/scripts/color_enhance.py +++ b/scripts/color_enhance.py @@ -1,14 +1,8 @@ import gradio as gr -import imageio.core.util -import numpy as np -import skimage.color -from PIL import Image from modules import scripts_postprocessing from modules.ui_components import FormRow - - -imageio.core.util._precision_warn = lambda *args, **kwargs: None +import mmaker_color_enhance_core as lib class ScriptPostprocessingColorEnhance(scripts_postprocessing.ScriptPostprocessing): @@ -25,14 +19,8 @@ class ScriptPostprocessingColorEnhance(scripts_postprocessing.ScriptPostprocessi return info_bak = {} if not hasattr(pp.image, "info") else pp.image.info - pp.image = self._color_enhance(pp.image, strength) + pp.image = lib.color_enhance(pp.image, strength) pp.image.info = info_bak pp.info["Color Enhance"] = strength - def _lerp(self, a: float, b: float, t: float) -> float: - return (1 - t) * a + t * b - - def _color_enhance(self, arr, strength: float = 1) -> Image.Image: - lch = skimage.color.lab2lch(lab=skimage.color.rgb2lab(rgb=np.array(arr, dtype=np.uint8))) - lch[:, :, 1] *= 100/(self._lerp(100, lch[:, :, 1].max(), strength)) # Normalize chroma component - return Image.fromarray(np.array(skimage.color.lab2rgb(lab=skimage.color.lch2lab(lch=lch)) * 255, dtype=np.uint8)) \ No newline at end of file + \ No newline at end of file