commit 3e39e9bb09b80f0cd142c63ecb9cebc9f39ad140 Author: Immac Date: Sun Mar 1 16:36:11 2026 -0600 Copied from ImmacTools, diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..d15dd99 --- /dev/null +++ b/LICENSE @@ -0,0 +1,22 @@ +MIT License + +Copyright (c) 2025, Immac + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..6514960 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,9 @@ +include LICENSE +include README.md + +recursive-exclude * __pycache__ +recursive-exclude * *.py[co] + +recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif + + diff --git a/README.md b/README.md new file mode 100644 index 0000000..5f3a64d --- /dev/null +++ b/README.md @@ -0,0 +1,99 @@ + +# Immac Tools + +A small collection of ComfyUI extension nodes created for personal use. + +## Quick summary + +- Package: `ComfyUI-ImmacTools` +- Purpose: ComfyUI extension providing utility nodes. + +## Installation + +This repository is intended to be used as a ComfyUI custom node. Install it by one of the following methods: + +A) ComfyUI Manager (When/If published) + +- Use ComfyUI's built-in extension/manager (if available) to add this repo directly by URL. The manager will clone the repository into the appropriate `custom_nodes` folder and keep it updated. + +B) Clone into your ComfyUI `custom_nodes` folder + +1. From your ComfyUI repository root (or wherever you keep custom nodes), clone this repo into the `custom_nodes` directory: + +```bash +# Example (from ComfyUI repo root) +git clone https://github.com/Immac/ComfyUI-ImmacTools custom_nodes/ComfyUI-ImmacTools +``` + +2. Restart ComfyUI. The extension's `comfy_entrypoint` should be discovered and the nodes will appear in the node list. + +# Immac Tools + +A small collection of ComfyUI custom nodes created for personal use and convenience. + +## Quick summary + +- Package: `immac_tools` +- Purpose: Utility nodes for ComfyUI (sigma schedule helpers and small helpers). + +## Installation + +This repository is intended to be installed as a ComfyUI custom node. Two common installation methods: + +1) Clone into your ComfyUI `custom_nodes` folder + +From your ComfyUI repository root (or wherever you keep custom nodes), clone this repo into the `custom_nodes` directory: + +```bash +# Example (from ComfyUI repo root) +git clone https://github.com/Immac/ComfyUI-ImmacTools custom_nodes/immac_tools +``` + +Then restart ComfyUI. The extension entrypoint should be discovered and the nodes will appear in the node list. + +2) ComfyUI Manager (if available) + +If you use a community manager/extension manager that supports installing extensions by Git URL, add this repo's URL and let the manager clone/update it for you. + +## Where the code lives + +- Node implementations: `src/immac_tools/nodes.py` and `src/immac_tools/forwarding_nodes.py`. +- Python package metadata: `pyproject.toml` and top-level package under `src/immac_tools`. + +## Nodes included + +The main utility nodes included today are: + +- Concatenate Sigmas Node + - Class: `ConcatenateSigmasImmacTools` + - Display name: `Concatenate Sigmas Node` + - Inputs: `sigmas_1`, `sigmas_2` (sigma schedules as tensors or lists) + - Outputs: concatenated sigma schedule + - Behavior: concatenates two sigma schedules, tolerating None inputs and non-torch objects by converting them to tensors. + +- Splice Sigmas At % + - Class: `SpliceSigmasAtImmacTools` + - Display name: `Splice Sigmas At %` + - Inputs: `sigmas_a`, `sigmas_b`, `splice` (float 0.0 - 1.0) + - Outputs: `spliced_sigmas` (full), `first_part`, `second_part` + - Behavior: picks a boundary on `sigmas_a` based on `splice` and combines `sigmas_a` and `sigmas_b` around that boundary. Accepts tensor or list inputs. + +- Skip Every Nth Image + - Class: `SkipEveryNthImagesImmacTools` + - Display name: `Skip Every Nth Image` + - Inputs: `images` (batch as list/tuple or torch tensor), `n` (int; n<=0 yields empty output) + - Outputs: filtered images batch, indices removed + - Behavior: drops every nth image while preserving the order of remaining items; returns the kept batch plus the zero-based indices that were removed. + +For full details see the source files mentioned above. + + + +## Contributing + +- Bug reports and feature requests: open an issue on the GitHub repository. + +## License & contact + +MIT — see `LICENSE`. +For questions or issues open an issue at: https://github.com/Immac/ComfyUI-ImmacTools diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..aad8d55 --- /dev/null +++ b/__init__.py @@ -0,0 +1,19 @@ +"""Top-level package for immac_tools.""" + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "WEB_DIRECTORY", +] + +__author__ = """Immac""" +__email__ = "immac.gm@gmail.com" +__version__ = "0.0.4" +__title__ = "Immac Tools" +__description__ = "A collection of Nodes for personal use." +__icon_ = "🧰" +from .src.immac_tools import NODE_CLASS_MAPPINGS +from .src.immac_tools import NODE_DISPLAY_NAME_MAPPINGS +from .src.immac_tools import WEB_DIRECTORY + + diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..3be24e0 Binary files /dev/null and b/__pycache__/__init__.cpython-312.pyc differ diff --git a/immac_tools.json b/immac_tools.json new file mode 100644 index 0000000..bf86dbf --- /dev/null +++ b/immac_tools.json @@ -0,0 +1,7 @@ +{ + "name": "Immac Tools", + "icon": "🧰", + "description": "A collection of Nodes for personal use.", + "author": "Immac", + "version": "0.0.2" +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4d46380 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,91 @@ +[build-system] +requires = ["setuptools>=70.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "ComfyUI-ImmacStyleMixer" +version = "0.0.1" +description = "A collection of Nodes for personal use." +authors = [ + {name = "Immac", email = "immac.gm@gmail.com"} +] +readme = "README.md" +license = {text = "MIT license"} +requires-python = ">=3.10" +classifiers = [] +dependencies = [ + +] + +[project.optional-dependencies] +dev = [ + "bump-my-version", + "coverage", # testing + "mypy", # linting + "pre-commit", # runs linting on commit + "pytest", # testing + "ruff", # linting +] + +[project.urls] +Repository = "https://github.com/Immac/ComfyUI-ImmacTools" +BugTracker = "https://github.com/Immac/ComfyUI-ImmacTools/issues" +Documentation = "https://github.com/Immac/ComfyUI-ImmacTools/wiki" + + +[tool.comfy] +PublisherId = "immac" +DisplayName = "Immac's Tools" +Icon = "" +Tags = ["sigmas"] +Repository = "https://github.com/Immac/ComfyUI-ImmacTools" + +includes = [] + +[tool.setuptools] +package-dir = {"" = "src"} + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-data] +"*" = ["*.*"] + +[tool.pytest.ini_options] +minversion = "8.0" +testpaths = [ + "tests", +] + +[tool.mypy] +files = "." + +# Use strict defaults +strict = true +warn_unreachable = true +warn_no_return = true + +[[tool.mypy.overrides]] +# Don't require test functions to include types +module = "tests.*" +allow_untyped_defs = true +disable_error_code = "attr-defined" + +[tool.ruff] +# extend-exclude = ["static", "ci/templates"] +line-length = 140 +src = ["src", "tests"] +target-version = "py310" + +# Add rules to ban exec/eval +[tool.ruff.lint] +select = [ + "S102", # exec-builtin + "S307", # eval-used + "W293", + "F", # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names. + # See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f +] + +[tool.ruff.lint.flake8-quotes] +inline-quotes = "double" diff --git a/src/immac_tools/__init__.py b/src/immac_tools/__init__.py new file mode 100644 index 0000000..f57bff5 --- /dev/null +++ b/src/immac_tools/__init__.py @@ -0,0 +1,40 @@ +"""ComfyUI-ImmacTools package exports used by ComfyUI. + +This module exposes the node registration mappings and the `WEB_DIRECTORY` +constant expected by ComfyUI web extensions. +""" +from .nodes import ConcatenateSigmasNode, SpliceSigmasAtNode, ResampleSigmas, SkipEveryNthImages, MatchContrastNode +from .forwarding_nodes import ForwardAnyNode, ForwardConditioningNode, ForwardModelNode + +# Mapping from node_id to node class +NODE_CLASS_MAPPINGS: dict[str, type] = { + "ConcatenateSigmasImmacTools": ConcatenateSigmasNode, + "SpliceSigmasAtImmacTools": SpliceSigmasAtNode, + "ResampleSigmasImmacTools": ResampleSigmas, + "SkipEveryNthImagesImmacTools": SkipEveryNthImages, + "MatchContrastImmacTools": MatchContrastNode, + "ForwardAnyImmacTools": ForwardAnyNode, + "ForwardConditioningImmacTools": ForwardConditioningNode, + "ForwardModelImmacTools": ForwardModelNode, +} + +# Mapping from node_id to human readable display name +NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = { + "ConcatenateSigmasImmacTools": "Concatenate Sigmas Node", + "SpliceSigmasAtImmacTools": "Splice Sigmas At Node", + "ResampleSigmasImmacTools": "Resample Sigmas", + "SkipEveryNthImagesImmacTools": "Skip Every Nth Image", + "MatchContrastImmacTools": "Match Contrast", + "ForwardAnyImmacTools": "Forward Any", + "ForwardConditioningImmacTools": "Forward Conditioning", + "ForwardModelImmacTools": "Forward Model", +} + +# ComfyUI expects this for web extensions +WEB_DIRECTORY = "./web" + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "WEB_DIRECTORY", +] diff --git a/src/immac_tools/__pycache__/__init__.cpython-312.pyc b/src/immac_tools/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..bf4adad Binary files /dev/null and b/src/immac_tools/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/immac_tools/__pycache__/forwarding_nodes.cpython-312.pyc b/src/immac_tools/__pycache__/forwarding_nodes.cpython-312.pyc new file mode 100644 index 0000000..1ee0b9a Binary files /dev/null and b/src/immac_tools/__pycache__/forwarding_nodes.cpython-312.pyc differ diff --git a/src/immac_tools/__pycache__/nodes.cpython-312.pyc b/src/immac_tools/__pycache__/nodes.cpython-312.pyc new file mode 100644 index 0000000..139d119 Binary files /dev/null and b/src/immac_tools/__pycache__/nodes.cpython-312.pyc differ diff --git a/src/immac_tools/forwarding_nodes.py b/src/immac_tools/forwarding_nodes.py new file mode 100644 index 0000000..55bf83c --- /dev/null +++ b/src/immac_tools/forwarding_nodes.py @@ -0,0 +1,119 @@ +from typing_extensions import override + +from comfy_api.latest import ComfyExtension, io +import torch + + +class ForwardAnyNode(io.ComfyNode): + """ + ForwardingNode + A simple pass-through node that returns its single input unchanged. + + Inputs: + - value: any type accepted by ComfyUI inputs (Image, Latent, Float, Int, Sigmas, etc.) + + Outputs: + - value: same as input + """ + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="ForwardAnyImmacTools", + display_name="Forward Any", + category="Example", + inputs=[ + io.AnyType.Input("value",display_name="->"), + ], + outputs=[ + io.AnyType.Output(display_name="->"), + ], + ) + + @classmethod + def check_lazy_status(cls, value): + return [] + + @classmethod + def execute(cls, value) -> io.NodeOutput: + # Pass-through: simply return the input as output + return io.NodeOutput(value) + + +class ForwardConditioningNode(io.ComfyNode): + """ + ConditioningForwardingNode + A pass-through node specifically for Conditioning types. + + Inputs: + - conditioning: Conditioning input + + Outputs: + - conditioning: same Conditioning output + """ + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="ForwardConditioningImmacTools", + display_name="Forward Conditioning", + category="Example", + inputs=[ + io.Conditioning.Input("conditioning", display_name="->"), + ], + outputs=[ + io.Conditioning.Output(display_name="->"), + ], + ) + + @classmethod + def check_lazy_status(cls, conditioning): + return [] + + @classmethod + def execute(cls, conditioning) -> io.NodeOutput: + # Pass-through: simply return the conditioning as output + return io.NodeOutput(conditioning) + + +class ForwardModelNode(io.ComfyNode): + """ + ModelForwardingNode + A pass-through node specifically for Model types. + + Inputs: + - model: Model input + + Outputs: + - model: same Model output + """ + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="ForwardModelImmacTools", + display_name="Forward Model", + category="Example", + inputs=[ + io.Model.Input("model", display_name="->"), + ], + outputs=[ + io.Model.Output(display_name="->"), + ], + ) + + @classmethod + def check_lazy_status(cls, model): + return [] + + @classmethod + def execute(cls, model) -> io.NodeOutput: + # Pass-through: simply return the model as output + return io.NodeOutput(model) + + +class ExampleForwardingExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ForwardAnyNode, ForwardConditioningNode, ForwardModelNode] + + + + diff --git a/src/immac_tools/nodes.py b/src/immac_tools/nodes.py new file mode 100644 index 0000000..19dcb7f --- /dev/null +++ b/src/immac_tools/nodes.py @@ -0,0 +1,684 @@ +from typing_extensions import override +from collections.abc import Sequence + +from comfy_api.latest import ComfyExtension, io +from .forwarding_nodes import ForwardAnyNode, ForwardConditioningNode, ForwardModelNode +import torch +import numpy as np + +class ConcatenateSigmasNode(io.ComfyNode): + """ + ConcatenateSigmasNode + A ComfyUI-compatible node that concatenates two "sigmas" inputs into a single 1D sigma tensor. + + Inputs: + - sigmas_1: torch.Tensor or tensor-like (1D or scalar). May be None. + - sigmas_2: torch.Tensor or tensor-like (1D or scalar). May be None. + + Output: + - io.NodeOutput containing the concatenated torch.Tensor, or the single non-None input, or None if both inputs are None. + + Notes and edge cases: + - If inputs have more than 1 dimension, concatenation occurs along the first dimension; the remaining dimensions must match. + - If inputs are on different devices or have incompatible dtypes/shapes, torch.cat will raise an error. Convert or move tensors to a common device/dtype before feeding them to this node if necessary. + - The node's check_lazy_status does not request any additional lazy evaluation; all provided inputs are processed directly. + + Edge cases and recommendations: + - When inputs are on different devices (CPU vs GPU) or have different dtypes, callers should coerce them to a common device/dtype before connecting to this node; otherwise torch.cat will raise. + - If inputs have trailing dimensions (i.e. dim() > 1), those trailing shapes must match exactly; the node will concatenate along dim=0. + + Examples: + - Given sigmas_1 = torch.tensor([0.1, 0.2]) and sigmas_2 = torch.tensor([0.3]), the output is tensor([0.1, 0.2, 0.3]). + - Given sigmas_1 = 0.5 (scalar) and sigmas_2 = [0.6, 0.7], the output is tensor([0.5, 0.6, 0.7]) after conversion and unsqueezing. + """ + @classmethod + def define_schema(cls) -> io.Schema: + """ + Return a schema which contains all information about the node. + Some types: "Model", "Vae", "Clip", "Conditioning", "Latent", "Image", "Int", "String", "Float", "Combo". + For outputs the "io.Model.Output" should be used, for inputs the "io.Model.Input" can be used. + The type can be a "Combo" - this will be a list for selection. + """ + return io.Schema( + node_id="ConcatenateSigmasImmacTools", + display_name="Concatenate Sigmas Node", + category="Example", + inputs=[ + io.Sigmas.Input("sigmas_1"), + io.Sigmas.Input("sigmas_2"), + ], + outputs=[ + io.Sigmas.Output(), + ], + ) + + @classmethod + def check_lazy_status(cls, sigmas_1, sigmas_2): + """ + Return a list of input names that need to be evaluated. + + This function will be called if there are any lazy inputs which have not yet been + evaluated. As long as you return at least one field which has not yet been evaluated + (and more exist), this function will be called again once the value of the requested + field is available. + + Any evaluated inputs will be passed as arguments to this function. Any unevaluated + inputs will have the value None. + """ + return [] + + @classmethod + def execute(cls, sigmas_1, sigmas_2) -> io.NodeOutput: + if sigmas_1 is None: + return io.NodeOutput(sigmas_2) + if sigmas_2 is None: + return io.NodeOutput(sigmas_1) + + if not isinstance(sigmas_1, torch.Tensor): + sigmas_1 = torch.as_tensor(sigmas_1) + if not isinstance(sigmas_2, torch.Tensor): + sigmas_2 = torch.as_tensor(sigmas_2) + + if sigmas_1.dim() == 0: + sigmas_1 = sigmas_1.unsqueeze(0) + if sigmas_2.dim() == 0: + sigmas_2 = sigmas_2.unsqueeze(0) + + # If the last value of sigmas_1 is very close to the first value of sigmas_2, + # avoid duplicating that boundary sigma so the concatenated schedule has + # exactly one sample at the splice point. + # Ensure both tensors are on the same device/dtype for comparison + try: + sig1_last = sigmas_1[-1] + sig2_first = sigmas_2[0] + except Exception: + # If indexing fails (shouldn't for non-empty 1D tensors), fall back to simple concat + out = torch.cat((sigmas_1, sigmas_2), dim=0) + return io.NodeOutput(out) + + # Move comparison tensors to a common dtype/device for isclose + comp_dtype = sigmas_1.dtype + comp_device = sigmas_1.device + if sigmas_2.device != comp_device: + sigmas_2 = sigmas_2.to(comp_device) + sig2_first = sig2_first.to(comp_device) + if sigmas_2.dtype != comp_dtype: + sigmas_2 = sigmas_2.to(dtype=comp_dtype) + sig2_first = sig2_first.to(dtype=comp_dtype) + + # Use a relative tolerance based on magnitude plus a small absolute tolerance + atol = 1e-6 + rtol = 1e-4 + if torch.isclose(sig1_last, sig2_first, atol=atol, rtol=rtol): + # Concatenate but skip the duplicated first value of sigmas_2 + if sigmas_2.numel() > 1: + out = torch.cat((sigmas_1, sigmas_2[1:]), dim=0) + else: + # sigmas_2 only contained the duplicated boundary, so result is just sigmas_1 + out = sigmas_1.clone() + else: + out = torch.cat((sigmas_1, sigmas_2), dim=0) + + return io.NodeOutput(out) + + """ + The node will always be re executed if any of the inputs change but + this method can be used to force the node to execute again even when the inputs don't change. + You can make this node return a number or a string. This value will be compared to the one returned the last time the node was + executed, if it is different the node will be executed again. + This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash + changes between executions the LoadImage node is executed again. + """ + + #@classmethod + #def fingerprint_inputs(s, image, string_field, int_field, float_field, print_to_screen): + # return "" + +class SpliceSigmasAtNode(io.ComfyNode): + """ + Splice two sigma schedules at a relative point of the first schedule. + + Inputs: + - sigmas_a: first sigma schedule (descending) + - sigmas_b: second sigma schedule + - boundary: float in [0.0, 1.0] representing where to splice relative to the first sigma (0 -> start, 1 -> end) + + Outputs: + - full concatenated sigma (first_part + second_part without duplicated boundary) + - first_part (prefix of sigmas_a up to the boundary, boundary appended if needed) + - second_part (boundary prepended then suffix of sigmas_b <= boundary) + + Edge cases and behavior: + - If `boundary` is outside [0.0, 1.0], it is clamped to the range [0.0, 1.0]. + - The `boundary` value is interpreted as a fraction of the starting sigma value (sigmas_a[0] * (1.0 - boundary)). + For example, `boundary=0.0` results in a boundary equal to `sigmas_a[0]` (i.e. start), and `boundary=1.0` results in a boundary of 0. + - If no values in `sigmas_a` are >= boundary value, `first_part` will be empty (unless `include_boundary` is True, in which case the boundary value will be appended). + - If no values in `sigmas_b` are <= boundary value, `second_part` will be empty (unless `include_boundary` is True, in which case the boundary value will be prepended). + - When `include_boundary` is True but both schedules do not contain the boundary, the boundary will be inserted into both `first_part` and `second_part` (so `full` will contain the boundary once). + - If the inputs are on different devices or have different dtypes, callers should coerce them to a common device/dtype before connecting to this node; the node assumes inputs are compatible for comparisons and concatenation. + """ + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="SpliceSigmasAtImmacTools", + display_name="Splice Sigmas At %", + category="Example", + inputs=[ + io.Sigmas.Input("sigmas_a"), + io.Sigmas.Input("sigmas_b"), + io.Float.Input("boundary", default=0.5, min=0.0, max=1.0, step=0.005), + io.Boolean.Input( + "include_boundary", + display_name="include boundary?", + tooltip="If enabled, the boundary sigma is included at the splice point, which can smooth transitions but adds an extra step." + ) + ], + outputs=[ + io.Sigmas.Output("full","spliced_sigmas"), + io.Sigmas.Output("first_part","first_part"), + io.Sigmas.Output("second_part","second_part"), + ], + ) + + @classmethod + def check_lazy_status(cls, sigmas_a, sigmas_b, boundary, include_boundary): + # Signature matches the defined inputs: (sigmas_a, sigmas_b, boundary, include_boundary) + return [] + + @classmethod + def execute(cls, sigmas_a, sigmas_b, boundary, include_boundary) -> io.NodeOutput: + # If either schedule missing, pass through sensible defaults + if sigmas_a is None and sigmas_b is None: + return io.NodeOutput(None, None, None) + if sigmas_a is None: + return io.NodeOutput(sigmas_b, None, sigmas_b) + if sigmas_b is None: + return io.NodeOutput(sigmas_a, sigmas_a, None) + + # Ensure tensors and 1D + if not isinstance(sigmas_a, torch.Tensor): + sigmas_a = torch.as_tensor(sigmas_a) + if not isinstance(sigmas_b, torch.Tensor): + sigmas_b = torch.as_tensor(sigmas_b) + if sigmas_a.dim() == 0: + sigmas_a = sigmas_a.unsqueeze(0) + if sigmas_b.dim() == 0: + sigmas_b = sigmas_b.unsqueeze(0) + + # Clamp boundary to [0,1] + try: + p = float(boundary) + except Exception: + p = 0.5 + p = max(0.0, min(1.0, p)) + + # Boundary value: interpreted as the remaining fraction of the initial sigma + # (example: boundary=0.5 => boundary = sigmas_a[0] * 0.5) + start_val = float(sigmas_a[0].item()) + boundary_val = start_val * (1.0 - p) + + dtype = sigmas_a.dtype + device = sigmas_a.device + boundary_tensor = torch.as_tensor([boundary_val], dtype=dtype, device=device) + + # Build first part: all values from sigmas_a >= boundary, then ensure boundary appended if needed + mask_a = sigmas_a >= boundary_val + if mask_a.any(): + last_idx = torch.nonzero(mask_a, as_tuple=True)[0][-1].item() + first_part = sigmas_a[: last_idx + 1].clone() + else: + first_part = torch.empty((0,), dtype=dtype, device=device) + + # Build second part: include values from sigmas_b <= boundary + mask_b = sigmas_b <= boundary_val + if mask_b.any(): + suffix_b = sigmas_b[mask_b].clone() + else: + suffix_b = torch.empty((0,), dtype=dtype, device=device) + + if include_boundary: + # Append boundary to first_part if not present + if first_part.numel() == 0 or not torch.isclose(first_part[-1], boundary_tensor[0], atol=1e-6): + first_part = torch.cat((first_part, boundary_tensor)) + # Prepend boundary to second_part + second_part = torch.cat((boundary_tensor, suffix_b)) + # Concatenate without duplicating the boundary value + if second_part.numel() > 1: + full = torch.cat((first_part, second_part[1:])) + else: + full = torch.cat((first_part, second_part[0:0])) # just first_part if second_part has only boundary + else: + # No boundary is added, so just use first_part and suffix_b directly + second_part = suffix_b + # Concatenate first_part and second_part directly + full = torch.cat((first_part, second_part)) + + return io.NodeOutput(full, first_part, second_part) + +class ResampleSigmas(io.ComfyNode): + """ + ResampleSigmas + + Resample / re-interpret a sigma schedule to a requested number of steps. + + Inputs: + - sigmas: torch.Tensor or tensor-like (1D or scalar). May be None. + - steps: int - number of steps to produce (clamped to >= 1) + + Output: + - io.NodeOutput containing the resampled torch.Tensor (shape: [steps + 1, ...]) or None if input sigmas is None + + Behavior: + - The node interprets `steps` as the number of intervals; the returned schedule includes both endpoints and + therefore has length `steps + 1`. + - If `sigmas` is scalar, the output is a length-`steps + 1` tensor filled with that value. + - If `sigmas` is 1D (T,), output is produced by linear interpolation of the curve defined by the T points at + positions [0..T-1] sampled at `steps + 1` evenly spaced positions in that same range. + - If `sigmas` has trailing dimensions (T, ...), interpolation is applied across the first (time) dimension and + trailing dimensions are preserved. + - Dtype and device of the output match the input. + """ + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="ResampleSigmasImmacTools", + display_name="Resample Sigmas", + category="Example", + inputs=[ + io.Sigmas.Input("sigmas"), + io.Int.Input("steps", default=10, min=1), + ], + outputs=[ + io.Sigmas.Output(), + ], + ) + + @classmethod + def check_lazy_status(cls, sigmas, steps): + return [] + + @classmethod + def execute(cls, sigmas, steps) -> io.NodeOutput: + # Pass-throughs + if sigmas is None: + return io.NodeOutput(None) + + # Ensure steps is an int and at least 1 + try: + s = int(steps) + except Exception: + s = 1 + s = max(1, s) + # We interpret `steps` as the number of intervals; the returned sigma schedule + # should include both endpoints, so it has length `s + 1`. + out_len = s + 1 + + # Ensure tensor + if not isinstance(sigmas, torch.Tensor): + sigmas = torch.as_tensor(sigmas) + + # Defensive: ensure we have a floating dtype for linspace and interpolation. + # If sigmas is integer or non-floating, cast to the default floating dtype + # while keeping device. This prevents unexpected integer linspace behavior. + if not torch.is_floating_point(sigmas): + sigmas = sigmas.to(dtype=torch.get_default_dtype()) + + # If scalar -> repeat the single value for each timestep including endpoints + if sigmas.dim() == 0: + out = sigmas.unsqueeze(0).repeat(out_len) + return io.NodeOutput(out) + + # Work along first dimension (time). Preserve trailing dims. + T = sigmas.shape[0] + dtype = sigmas.dtype + device = sigmas.device + + # If only one input time point, just repeat that point for each output timestep + if T == 1: + out = sigmas.unsqueeze(0).repeat(out_len, *([1] * (sigmas.dim() - 1))) + return io.NodeOutput(out) + + # Flatten trailing dims so we can vectorize interpolation + trailing_shape = sigmas.shape[1:] + D = 1 + for d in trailing_shape: + D *= d + sig_flat = sigmas.contiguous().view(T, D) + + # Source x positions and target x positions + src_x = torch.linspace(0.0, float(T - 1), steps=T, dtype=dtype, device=device) + # Target positions: produce `out_len` samples evenly spaced across [0, T-1] + tgt_x = torch.linspace(0.0, float(T - 1), steps=out_len, dtype=dtype, device=device) + + # For each target position find right-side index + idx = torch.searchsorted(src_x, tgt_x) + # idx in [0..T]; left = clamp(idx-1,0), right = clamp(idx, T-1) + right = torch.clamp(idx, 0, T - 1) + left = torch.clamp(right - 1, 0, T - 1) + + xL = src_x[left] + xR = src_x[right] + + # Gather yL/yR -> shape (out_len, D) + yL = sig_flat[left] # indexing with (out_len,) yields (out_len, D) + yR = sig_flat[right] + + # Avoid division by zero where xR == xL (shouldn't happen except maybe at boundaries) + denom = (xR - xL) + # Where denom==0, set weight to 0 to pick yL + denom_safe = denom.clone() + denom_safe[denom_safe == 0] = 1.0 + weight = ((tgt_x - xL) / denom_safe).unsqueeze(1) + weight = weight.clamp(0.0, 1.0) + + out_flat = yL * (1.0 - weight) + yR * weight + + out = out_flat.view((out_len,) + trailing_shape) + + return io.NodeOutput(out) + + +class SkipEveryNthImages(io.ComfyNode): + """ + Skip every nth image in a batch. + + Accepts a batch of images (list/tuple of frames or torch tensor with batch dim) and drops the nth item while preserving order of the rest. Outputs the filtered batch and the indices that were removed. + """ + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="SkipEveryNthImagesImmacTools", + display_name="Skip Every Nth Image", + category="Example", + inputs=[ + io.AnyType.Input("images", display_name="images", tooltip="Batch of images/frames (list or tensor)."), + io.Int.Input("n", default=2, min=1, tooltip="Skip every nth image; n<=0 yields empty output."), + ], + outputs=[ + io.AnyType.Output("filtered", display_name="images"), + io.AnyType.Output("removed_indices", display_name="removed indices"), + ], + ) + + @classmethod + def check_lazy_status(cls, images, n): + return [] + + @classmethod + def execute(cls, images, n) -> io.NodeOutput: + # Normalize n; if invalid, treat as <=0 and drop everything + try: + step = int(n) + except Exception: + step = 0 + + # Torch tensor batch: slice along batch dim + if isinstance(images, torch.Tensor): + total = images.shape[0] if images.dim() > 0 else 0 + if total == 0 or step <= 0: + empty = images.new_empty((0,) + tuple(images.shape[1:])) + removed = list(range(total)) + return io.NodeOutput(empty, removed) + + keep_mask = torch.ones(total, dtype=torch.bool, device=images.device) + keep_mask[step - 1 :: step] = False # zero-based index for the nth elements + filtered = images[keep_mask] if total > 0 else images.new_empty((0,) + tuple(images.shape[1:])) + removed = [i for i in range(total) if (i + 1) % step == 0] + return io.NodeOutput(filtered, removed) + + # Sequence (list/tuple) batch: iterate and filter + if isinstance(images, Sequence) and not isinstance(images, (str, bytes)): + total = len(images) + if step <= 0: + removed = list(range(total)) + return io.NodeOutput([], removed) + + kept = [] + removed = [] + for idx, img in enumerate(images): + if (idx + 1) % step == 0: + removed.append(idx) + else: + kept.append(img) + return io.NodeOutput(kept, removed) + + # Fallback: unsupported type, pass through unchanged + return io.NodeOutput(images, []) + + +# --------------------------------------------------------------------------- +# Histogram matching helpers +# --------------------------------------------------------------------------- + +def _match_histograms_channel(src: np.ndarray, ref: np.ndarray) -> np.ndarray: + """Match the histogram of *src* to that of *ref*. + + Both arrays must be float32 in [0, 1] and can have any shape. + The mapping is built by aligning the cumulative distribution functions (CDFs) + of the two arrays. We quantise to uint16 (65 536 levels) which gives + sub-0.002 % precision while keeping the LUT small and fast. + + Returns a float32 array with the same shape as *src*. + """ + src_u = np.clip(np.round(src * 65535.0), 0, 65535).astype(np.uint16) + ref_u = np.clip(np.round(ref * 65535.0), 0, 65535).astype(np.uint16) + + src_hist = np.bincount(src_u.ravel(), minlength=65536) + ref_hist = np.bincount(ref_u.ravel(), minlength=65536) + + src_cdf = np.cumsum(src_hist, dtype=np.float64) / float(src_u.size) + ref_cdf = np.cumsum(ref_hist, dtype=np.float64) / float(ref_u.size) + + # For each source quantisation level, find the reference level whose CDF + # value is closest (i.e. has the same mass to the left). + lut = np.searchsorted(ref_cdf, src_cdf, side="left") + lut = np.clip(lut, 0, 65535).astype(np.uint16) + + matched_u = lut[src_u] + return matched_u.astype(np.float32) / 65535.0 + + +# ---- pure-numpy sRGB <-> linear <-> XYZ <-> LAB helpers (D65) ---- + +def _srgb_to_linear(c: np.ndarray) -> np.ndarray: + return np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4) + +def _linear_to_srgb(c: np.ndarray) -> np.ndarray: + c = np.clip(c, 0.0, 1.0) + return np.where(c <= 0.0031308, c * 12.92, 1.055 * c ** (1.0 / 2.4) - 0.055) + +_M_RGB_TO_XYZ = np.array([ + [0.4124564, 0.3575761, 0.1804375], + [0.2126729, 0.7151522, 0.0721750], + [0.0193339, 0.1191920, 0.9503041], +], dtype=np.float64) + +_M_XYZ_TO_RGB = np.array([ + [ 3.2404542, -1.5371385, -0.4985314], + [-0.9692660, 1.8760108, 0.0415560], + [ 0.0556434, -0.2040259, 1.0572252], +], dtype=np.float64) + +_D65 = np.array([0.95047, 1.00000, 1.08883], dtype=np.float64) + + +def _f_lab(t: np.ndarray) -> np.ndarray: + delta = 6.0 / 29.0 + return np.where(t > delta ** 3, np.cbrt(t), t / (3.0 * delta ** 2) + 4.0 / 29.0) + + +def _f_lab_inv(t: np.ndarray) -> np.ndarray: + delta = 6.0 / 29.0 + return np.where(t > delta, t ** 3, 3.0 * delta ** 2 * (t - 4.0 / 29.0)) + + +def _rgb_to_lab(img: np.ndarray) -> np.ndarray: + """Convert [H, W, 3] float32 sRGB image to CIE LAB.""" + lin = _srgb_to_linear(img.astype(np.float64)) + xyz = lin @ _M_RGB_TO_XYZ.T + xyz_n = xyz / _D65 + f = _f_lab(xyz_n) + L = 116.0 * f[..., 1] - 16.0 + a = 500.0 * (f[..., 0] - f[..., 1]) + b = 200.0 * (f[..., 1] - f[..., 2]) + return np.stack([L, a, b], axis=-1).astype(np.float32) + + +def _lab_to_rgb(lab: np.ndarray) -> np.ndarray: + """Convert [H, W, 3] float32 CIE LAB image to sRGB [0, 1] float32.""" + lab64 = lab.astype(np.float64) + L, a, b = lab64[..., 0], lab64[..., 1], lab64[..., 2] + fy = (L + 16.0) / 116.0 + fx = a / 500.0 + fy + fz = fy - b / 200.0 + xyz = np.stack([_f_lab_inv(fx), _f_lab_inv(fy), _f_lab_inv(fz)], axis=-1) * _D65 + lin = xyz @ _M_XYZ_TO_RGB.T + return np.clip(_linear_to_srgb(lin), 0.0, 1.0).astype(np.float32) + + +class MatchContrastNode(io.ComfyNode): + """ + MatchContrastNode + + Adjusts the contrast (tonal distribution) of *image* to match that of + *reference* using CDF-based histogram matching. + + Inputs: + image – the image whose contrast will be adjusted [B, H, W, C] + reference – the image whose contrast is the target [B, H, W, C] + channel_mode: + "rgb" – match each R/G/B channel independently + "luminance" – match only the L channel in CIE LAB space + (preserves the hue and saturation of *image*, + changes only brightness/contrast) + + Output: + Adjusted image with the same spatial dimensions as *image*. + + Notes: + • If *reference* contains fewer frames than *image*, the first frame + of *reference* is used for all source frames. + • The matching is done per-frame (not across the batch). + • No external dependencies beyond NumPy and PyTorch. + """ + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="MatchContrastImmacTools", + display_name="Match Contrast", + category="image/adjust", + inputs=[ + io.Image.Input("image", display_name="image"), + io.Image.Input("reference", display_name="reference"), + io.Combo.Input( + "channel_mode", + options=["rgb", "luminance"], + default="luminance", + display_name="channel mode", + tooltip=( + "rgb: match each R/G/B channel independently. " + "luminance: match only the L channel in CIE LAB " + "(keeps source colours, adjusts brightness/contrast)." + ), + ), + io.Float.Input( + "strength", + default=1.0, + min=0.0, + max=1.0, + step=0.01, + display_name="strength", + tooltip="Blend between the original image (0.0) and the fully matched image (1.0).", + ), + ], + outputs=[ + io.Image.Output(display_name="image"), + ], + ) + + @classmethod + def check_lazy_status(cls, image, reference, channel_mode, strength): + return [] + + @classmethod + def execute(cls, image: torch.Tensor, reference: torch.Tensor, channel_mode: str, strength: float) -> io.NodeOutput: + # image / reference: [B, H, W, C] float32 in [0, 1] + B = image.shape[0] + ref_B = reference.shape[0] + s = float(np.clip(strength, 0.0, 1.0)) + + out_frames: list[np.ndarray] = [] + + for i in range(B): + src_np = image[i].cpu().numpy() # [H, W, C] + # Broadcast reference: use frame i if available, else frame 0 + ref_np = reference[min(i, ref_B - 1)].cpu().numpy() # [H, W, C] + + if channel_mode == "luminance": + matched = cls._match_luminance(src_np, ref_np) + else: # "rgb" + matched = cls._match_rgb(src_np, ref_np) + + # Blend: 0.0 = original, 1.0 = fully matched + blended = src_np * (1.0 - s) + matched * s + out_frames.append(blended) + + result = torch.from_numpy(np.stack(out_frames, axis=0)).to(image.device) + return io.NodeOutput(result) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _match_rgb(src: np.ndarray, ref: np.ndarray) -> np.ndarray: + """Match each RGB channel of *src* to *ref* independently.""" + C = src.shape[-1] + channels = [ + _match_histograms_channel(src[..., c], ref[..., c]) + for c in range(C) + ] + return np.stack(channels, axis=-1) + + @staticmethod + def _match_luminance(src: np.ndarray, ref: np.ndarray) -> np.ndarray: + """Match only the LAB L-channel of *src* to *ref*; keep src a/b.""" + src_lab = _rgb_to_lab(src) # [H, W, 3] + ref_lab = _rgb_to_lab(ref) + + # Normalise L from [0, 100] to [0, 1] for the histogram helper + L_src = src_lab[..., 0] / 100.0 + L_ref = ref_lab[..., 0] / 100.0 + + L_matched = _match_histograms_channel(L_src, L_ref) * 100.0 # back to [0,100] + + out_lab = src_lab.copy() + out_lab[..., 0] = L_matched + return _lab_to_rgb(out_lab) + + +# Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension +# WEB_DIRECTORY = "./somejs" + + +# Add custom API routes, using router + +class ExampleExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + ConcatenateSigmasNode, + SpliceSigmasAtNode, + ResampleSigmas, + SkipEveryNthImages, + MatchContrastNode, + ForwardAnyNode, + ForwardConditioningNode, + ForwardModelNode, + ] + +async def comfy_entrypoint() -> ExampleExtension: # ComfyUI calls this to load your extension and its nodes. + return ExampleExtension() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..2ad204b --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Unit test package for immac_tools.""" diff --git a/tests/__pycache__/__init__.cpython-312.pyc b/tests/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..48d704c Binary files /dev/null and b/tests/__pycache__/__init__.cpython-312.pyc differ diff --git a/tests/__pycache__/conftest.cpython-312-pytest-8.4.2.pyc b/tests/__pycache__/conftest.cpython-312-pytest-8.4.2.pyc new file mode 100644 index 0000000..381cb4b Binary files /dev/null and b/tests/__pycache__/conftest.cpython-312-pytest-8.4.2.pyc differ diff --git a/tests/__pycache__/test_immac_tools.cpython-312-pytest-8.4.2.pyc b/tests/__pycache__/test_immac_tools.cpython-312-pytest-8.4.2.pyc new file mode 100644 index 0000000..ea2c7fa Binary files /dev/null and b/tests/__pycache__/test_immac_tools.cpython-312-pytest-8.4.2.pyc differ diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..310609c --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,6 @@ +import os +import sys + +# Add the project root directory to Python path +# This allows the tests to import the project +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 0000000..95c76f1 --- /dev/null +++ b/tests/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +testpaths = . # Run tests in the current directory +python_files = test_*.py # Run tests in files that start with "test_" +norecursedirs = .. # Don't run tests in the parent directory diff --git a/tests/test_immac_tools.py b/tests/test_immac_tools.py new file mode 100644 index 0000000..3c049b5 --- /dev/null +++ b/tests/test_immac_tools.py @@ -0,0 +1,21 @@ +#!/usr/bin/env python + +"""Tests for `immac_tools` package.""" + +import pytest +from src.immac_tools.nodes import ConcatenateSigmasNode + +@pytest.fixture +def example_node(): + """Fixture to create an Example node instance.""" + return ConcatenateSigmasNode() + +def test_example_node_initialization(example_node): + """Test that the node can be instantiated.""" + assert isinstance(example_node, ConcatenateSigmasNode) + +def test_return_types(): + """Test the node's metadata.""" + assert ConcatenateSigmasNode.RETURN_TYPES == ("IMAGE",) + assert ConcatenateSigmasNode.FUNCTION == "test" + assert ConcatenateSigmasNode.CATEGORY == "Example"