From da497db11049cd82c247dd9e2be7b1a643964aa2 Mon Sep 17 00:00:00 2001 From: mcDandy Date: Mon, 20 Apr 2026 19:16:50 +0200 Subject: [PATCH] AI: add Color math node --- README.md | 2 +- more_math/ColorMathNode.py | 168 +++++++++++++++++++++++++++++++++++++ more_math/nodes.py | 4 +- 3 files changed, 172 insertions(+), 2 deletions(-) create mode 100644 more_math/ColorMathNode.py diff --git a/README.md b/README.md index 2b75584..d89ec46 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ You can also get the node from comfy manager under the name of More math. - Functions and variables in math expressions. - Conversion between `INT` and `FLOAT`, `INT` and `BOOLEAN`, and `AUDIO` and `IMAGE` - image/audio conversion encodes frequency content into RGB channels. -- Node support for `FLOAT`, `STRING`, `CONDITIONING`, `LATENT`, `IMAGE`, `MASK`, `NOISE`, `AUDIO`, `VIDEO`, `MODEL`, `CLIP`, `VAE`, `SIGMAS`, and `GUIDER`. +- Node support for `FLOAT`, `STRING`, `CONDITIONING`, `LATENT`, `IMAGE`, `MASK`, `COLOR`, `NOISE`, `AUDIO`, `VIDEO`, `MODEL`, `CLIP`, `VAE`, `SIGMAS`, and `GUIDER`. - Vector math with list literals `[v1, v2, ...]` and mixed operations between lists, scalars, and tensors. - Custom functions `funcname(variable, variable, ...)->expression;` - can be used later in expressions and in other custom functions, diff --git a/more_math/ColorMathNode.py b/more_math/ColorMathNode.py new file mode 100644 index 0000000..b0b0370 --- /dev/null +++ b/more_math/ColorMathNode.py @@ -0,0 +1,168 @@ +"""Color math node using integer color representation.""" + +from __future__ import annotations + +import copy +from typing import Any + +import torch +from comfy_api.latest import io + +from .ParseTree import MrmthParseTree +from .Parser.UnifiedMathVisitor import UnifiedMathVisitor +from .Stack import MrmthStack +from .helper_functions import checkLazyNew, get_f_variable, get_v_variable, parse_expr + + +class ColorMathNode(io.ComfyNode): + """Apply expression math to Color values converted to packed RGB integers.""" + + @classmethod + def define_schema(cls) -> io.Schema: + """Define node schema.""" + return io.Schema( + node_id="mrmth_ag_ColorMathNode", + category="More math", + display_name="Color math", + inputs=[ + io.Autogrow.Input( + id="V", + template=io.Autogrow.TemplatePrefix( + io.Color.Input("values", optional=True), + prefix="V", + min=1, + max=50, + ), + ), + io.Autogrow.Input( + id="F", + template=io.Autogrow.TemplatePrefix( + io.Float.Input( + "float", + default=0.0, + optional=True, + lazy=True, + force_input=True, + ), + prefix="F", + min=1, + max=50, + ), + ), + io.MultiType.Input( + io.String.Input("Expression", default="V0", multiline=False), + types=[io.String, MrmthParseTree], + tooltip="Expression evaluated over packed color integers.", + ), + io.Boolean.Input( + id="remember_stack", + default=False, + display_name="Remember stack across batch", + ), + MrmthStack.Input( + id="stack", + tooltip="Access stack between nodes", + optional=True, + ), + ], + outputs=[ + io.Color.Output(id="color"), + MrmthStack.Output(), + ], + ) + + @staticmethod + def _hex_to_int(color_hex: str) -> int: + """Convert #RRGGBB/#RGB into packed integer 0xRRGGBB.""" + s = str(color_hex).strip() + if s.startswith("#"): + s = s[1:] + if len(s) == 3: + s = "".join(ch * 2 for ch in s) + if len(s) != 6: + raise ValueError(f"Invalid color format: {color_hex}") + return int(s, 16) & 0xFFFFFF + + @staticmethod + def _int_to_hex(value: int) -> str: + """Convert numeric value to #RRGGBB.""" + s = f"{int(value) & 0xFFFFFF:06x}" + return f"#{s}" + + @classmethod + def check_lazy_status( + cls, + Expression: Any, + V: dict[str, str], + F: dict[str, float], + remember_stack: bool = False, + stack: dict | None = None, + ) -> list[str]: + """Return unresolved lazy inputs for expression.""" + return checkLazyNew(Expression, V, F) + + @classmethod + def execute( + cls, + V: dict[str, str], + F: dict[str, float], + Expression: Any, + remember_stack: bool = False, + stack: dict | None = None, + ) -> tuple[str, dict]: + """Execute color expression and return resulting color.""" + work_stack = stack if remember_stack else ( + copy.deepcopy(stack) if stack is not None else {} + ) + + if not V: + raise ValueError("At least one color input is required.") + + variables: dict[str, Any] = {} + + # Legacy aliases. + variables["a"] = cls._hex_to_int(V.get("V0", "#000000")) + variables["b"] = cls._hex_to_int(V.get("V1", "#000000")) + variables["c"] = cls._hex_to_int(V.get("V2", "#000000")) + variables["d"] = cls._hex_to_int(V.get("V3", "#000000")) + + # All color inputs as packed ints. + color_int_inputs: dict[str, int] = {} + for k, val in V.items(): + color_int_inputs[k] = cls._hex_to_int(val if val is not None else "#000000") + variables[k] = color_int_inputs[k] + + v_stacked, v_cnt = get_v_variable(color_int_inputs, length_mismatch="do nothing") + if v_stacked is not None: + variables["V"] = v_stacked + variables["Vcnt"] = float(v_cnt) + variables["V_count"] = float(v_cnt) + + # Float aliases and F array. + variables["w"] = F.get("F0", 0.0) if F.get("F0") is not None else 0.0 + variables["x"] = F.get("F1", 0.0) if F.get("F1") is not None else 0.0 + variables["y"] = F.get("F2", 0.0) if F.get("F2") is not None else 0.0 + variables["z"] = F.get("F3", 0.0) if F.get("F3") is not None else 0.0 + + f_stacked, f_cnt = get_f_variable(F) + if f_stacked is not None: + variables["F"] = f_stacked + variables["Fcnt"] = float(f_cnt) + variables["F_count"] = float(f_cnt) + + for k, val in F.items(): + variables[k] = val if val is not None else 0.0 + + tree = parse_expr(Expression) if isinstance(Expression, str) else Expression + + visitor = UnifiedMathVisitor( + variables=variables, + shape=(1,), + device=torch.device("cpu"), + state_storage=work_stack, + ) + result = visitor.visit(tree) + out_color = cls._int_to_hex(result) + + returned_stack = work_stack if remember_stack else copy.deepcopy(work_stack) + return out_color, returned_stack diff --git a/more_math/nodes.py b/more_math/nodes.py index 6686912..ff5b320 100644 --- a/more_math/nodes.py +++ b/more_math/nodes.py @@ -23,6 +23,7 @@ from .VaeMathNode import VAEMathNode from .ClipMathNode import CLIPMathNode from .SpectrogramToAudioNode import SpectrogramToAudio from .AudioToSpectrogramNode import AudioToSpectrogram +from .ColorMathNode import ColorMathNode from .NoiseMathNode import NoiseMathNode from .AudioMathNode import AudioMathNode @@ -172,7 +173,8 @@ class MoreMathExtension(ComfyExtension): AudioMathNode, StringMathNode, ScriptTextInput, - SelectiveGuiderMathNode + SelectiveGuiderMathNode, + ColorMathNode ]