AI: add Color math node

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,

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

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
         ]