Add ComfyUI support

my bad y'all

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.
 

__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",
+}

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),)

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

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
-
+import mmaker_color_enhance_core as lib
-
-imageio.core.util._precision_warn = lambda *args, **kwargs: None
 
 
 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))