Add ComfyUI support

my bad y'all
This commit is contained in:
MMaker 2023-11-15 19:58:17 -05:00
parent b9e4a2cfa1
commit 78a01e432e
Signed by: mmaker
GPG Key ID: CCE79B8FEDA40FB2
5 changed files with 59 additions and 16 deletions

View File

@ -1,6 +1,6 @@
# Color Enhance # 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. 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.

9
__init__.py Normal file
View File

@ -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",
}

View File

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

View File

@ -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

View File

@ -1,14 +1,8 @@
import gradio as gr 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 import scripts_postprocessing
from modules.ui_components import FormRow 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): class ScriptPostprocessingColorEnhance(scripts_postprocessing.ScriptPostprocessing):
@ -25,14 +19,8 @@ class ScriptPostprocessingColorEnhance(scripts_postprocessing.ScriptPostprocessi
return return
info_bak = {} if not hasattr(pp.image, "info") else pp.image.info 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.image.info = info_bak
pp.info["Color Enhance"] = strength 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))