diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..219c8b1 --- /dev/null +++ b/__init__.py @@ -0,0 +1,9 @@ +"""ComfyUI-ColoredNoiseSampling — V3 custom-node pack. + +This package registers exclusively via the V3 ``comfy_entrypoint``. It deliberately does NOT +define ``NODE_CLASS_MAPPINGS`` — doing so would make ComfyUI's loader take the V1 branch and +never call ``comfy_entrypoint`` (silent failure). +""" +from .src.extension import comfy_entrypoint + +__all__ = ["comfy_entrypoint"] diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..3409c9b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,18 @@ +[project] +name = "comfyui-colorednoisesampling" +version = "0.1.0" +description = "Colored (frequency-shaped) noise sampling for ComfyUI: parametric + faithful gamma-matrix modes." +license = { file = "LICENSE" } +requires-python = ">=3.9" +dependencies = [] + +[project.urls] +Repository = "https://github.com/akashzeno/ComfyUI-ColoredNoiseSampling" + +[tool.comfy] +PublisherId = "akashzeno" +DisplayName = "ComfyUI-ColoredNoiseSampling" + +# NOTE: pytest config lives in tests/pytest.ini (NOT here). Keeping it out of pyproject.toml +# means pytest's rootdir is tests/, so the package's V3 __init__.py (which uses ComfyUI's +# package loader for its relative import) is never imported as a bare module during testing. diff --git a/src/extension.py b/src/extension.py new file mode 100644 index 0000000..506e487 --- /dev/null +++ b/src/extension.py @@ -0,0 +1,43 @@ +"""V3 extension entrypoint + gamma-matrix folder registration.""" +import os + +import folder_paths +from typing_extensions import override +from comfy_api.latest import ComfyExtension, io + +GAMMA_FOLDER = "colored_noise_gamma" + + +def _register_gamma_folder(): + """Register `colored_noise_gamma` so the bundled and user gamma matrices are listed.""" + package_root = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) + bundled = os.path.join(package_root, "gamma_matrices") + user = os.path.join(folder_paths.models_dir, GAMMA_FOLDER) + paths = [p for p in (bundled, user) if p] + exts = {".pt", ".pth"} + if GAMMA_FOLDER in folder_paths.folder_names_and_paths: + existing_paths, existing_exts = folder_paths.folder_names_and_paths[GAMMA_FOLDER] + for p in paths: + if p not in existing_paths: + existing_paths.append(p) + existing_exts.update(exts) + else: + folder_paths.folder_names_and_paths[GAMMA_FOLDER] = (paths, exts) + + +_register_gamma_folder() + +# Import nodes AFTER folder registration so their Combo option lists populate correctly. +from .nodes.sampler_select import ColoredNoise_SamplerSelect # noqa: E402 +from .nodes.noise import ColoredNoise_Noise # noqa: E402 +from .nodes.ksampler import ColoredNoise_KSampler # noqa: E402 + + +class ColoredNoiseExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ColoredNoise_SamplerSelect, ColoredNoise_Noise, ColoredNoise_KSampler] + + +async def comfy_entrypoint() -> ColoredNoiseExtension: + return ColoredNoiseExtension() diff --git a/src/nodes/ksampler.py b/src/nodes/ksampler.py new file mode 100644 index 0000000..0cc97bd --- /dev/null +++ b/src/nodes/ksampler.py @@ -0,0 +1,107 @@ +"""ColoredNoise_KSampler — all-in-one colored-noise sampler (model+conds -> latent). + +Mirrors comfy_extras.nodes_custom_sampler.SamplerCustom: computes sigmas internally, builds +the colored sampler, optionally swaps in colored initial noise, runs comfy.sample.sample_custom. +""" +import comfy.samplers +import comfy.sample +import comfy.utils +import latent_preview +from comfy_api.latest import io + +from ..sampling.registry import stochastic_samplers +from ..sampling.colored_ksampler import build_colored_sampler +from ..sampling.initial_noise import Noise_ColoredInitial +from ..engine.colored_noise import ColorParams +from .sampler_select import resolve_color_params, gamma_choices, default_sampler, NONE + + +class ColoredNoise_KSampler(io.ComfyNode): + @classmethod + def define_schema(cls): + samplers = list(stochastic_samplers().keys()) + return io.Schema( + node_id="ColoredNoise_KSampler", + display_name="Colored Noise KSampler", + category="sampling/colored_noise", + inputs=[ + io.Model.Input("model"), + io.Conditioning.Input("positive"), + io.Conditioning.Input("negative"), + io.Latent.Input("latent_image"), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True), + io.Int.Input("steps", default=20, min=1, max=10000), + io.Float.Input("cfg", default=8.0, min=0.0, max=100.0, step=0.1, round=0.01), + io.Combo.Input("base_sampler", options=samplers, default=default_sampler(samplers)), + io.Combo.Input("scheduler", options=comfy.samplers.SCHEDULER_NAMES), + io.Float.Input("denoise", default=1.0, min=0.0, max=1.0, step=0.01), + io.Boolean.Input("color_initial_noise", default=False, + tooltip="Also color the initial latent noise (uses alpha_start as a constant color)."), + io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True), + io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True), + io.Combo.Input("mode", options=["parametric", "gamma_matrix"], default="parametric"), + io.Float.Input("alpha_start", default=0.0, min=-8.0, max=8.0, step=0.05, round=False), + io.Float.Input("alpha_end", default=-1.0, min=-8.0, max=8.0, step=0.05, round=False), + io.Combo.Input("interpolation", options=["linear", "exponential"], default="linear", advanced=True), + io.Float.Input("exp_sharpness", default=4.0, min=0.1, max=16.0, step=0.1, round=False, advanced=True), + io.Combo.Input("gamma_matrix", options=gamma_choices(), default=NONE, advanced=True), + io.Float.Input("gamma_divider", default=1.0, min=0.1, max=10.0, step=0.01, round=False, advanced=True), + io.Combo.Input("gamma_shaping", options=["none", "sqrt", "power"], default="none", advanced=True), + io.Float.Input("power_gamma", default=1.0, min=0.1, max=8.0, step=0.05, round=False, advanced=True), + io.Float.Input("alpha_tilting", default=0.0, min=-8.0, max=8.0, step=0.05, round=False, advanced=True), + io.Float.Input("energy_scale", default=1.0, min=0.0, max=4.0, step=0.01, round=False, advanced=True), + ], + outputs=[ + io.Latent.Output(display_name="output"), + io.Latent.Output(display_name="denoised_output"), + ], + ) + + @classmethod + def execute(cls, model, positive, negative, latent_image, seed, steps, cfg, base_sampler, + scheduler, denoise, color_initial_noise, eta, s_noise, mode, alpha_start, alpha_end, + interpolation, exp_sharpness, gamma_matrix, gamma_divider, gamma_shaping, + power_gamma, alpha_tilting, energy_scale) -> io.NodeOutput: + params = resolve_color_params(mode, alpha_start, alpha_end, interpolation, exp_sharpness, + gamma_matrix, gamma_divider, gamma_shaping, power_gamma, + alpha_tilting, energy_scale) + base_fn = stochastic_samplers()[base_sampler] + sampler = build_colored_sampler( + base_fn, {"eta": eta, "s_noise": s_noise, "r": 0.5, "solver_type": "midpoint"}, params) + + # sigmas (mirror BasicScheduler denoise handling) + total_steps = steps if denoise >= 1.0 else max(1, int(steps / max(denoise, 1e-6))) + model_sampling = model.get_model_object("model_sampling") + sigmas = comfy.samplers.calculate_sigmas(model_sampling, scheduler, total_steps).cpu() + sigmas = sigmas[-(steps + 1):] + + latent = latent_image.copy() + latent_samples = comfy.sample.fix_empty_latent_channels(model, latent["samples"]) + latent["samples"] = latent_samples + + if color_initial_noise: + init_params = ColorParams(mode="parametric", alpha_start=alpha_start, + alpha_end=alpha_start, energy_scale=energy_scale) + noise = Noise_ColoredInitial(seed, init_params).generate_noise(latent) + else: + noise = comfy.sample.prepare_noise(latent_samples, seed, latent.get("batch_index", None)) + + noise_mask = latent.get("noise_mask", None) + x0_output = {} + callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output) + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED + + samples = comfy.sample.sample_custom( + model, noise, cfg, sampler, sigmas, positive, negative, latent_samples, + noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) + + out = latent.copy() + out["samples"] = samples + if "x0" in x0_output: + out_denoised = latent.copy() + out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu()) + else: + out_denoised = out + return io.NodeOutput(out, out_denoised) + + sample = execute diff --git a/src/nodes/noise.py b/src/nodes/noise.py new file mode 100644 index 0000000..9d8b6a4 --- /dev/null +++ b/src/nodes/noise.py @@ -0,0 +1,34 @@ +"""ColoredNoise_Noise — outputs a NOISE object that produces colored initial noise.""" +from comfy_api.latest import io + +from ..sampling.initial_noise import Noise_ColoredInitial +from ..engine.colored_noise import ColorParams + +COLOR_ALPHA = {"white": 0.0, "pink": 1.0, "brown": 2.0, "blue": -1.0, "violet": -2.0} + + +class ColoredNoise_Noise(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="ColoredNoise_Noise", + display_name="Colored Noise (Initial)", + category="sampling/colored_noise", + inputs=[ + io.Int.Input("noise_seed", default=0, min=0, max=0xffffffffffffffff, + control_after_generate=True), + io.Combo.Input("color", options=list(COLOR_ALPHA.keys()) + ["custom"], default="white"), + io.Float.Input("alpha", default=0.0, min=-8.0, max=8.0, step=0.05, round=False, + tooltip="Used when color = custom. PSD ~ 1/f^alpha; +red/pink, -blue/violet."), + io.Float.Input("energy_scale", default=1.0, min=0.0, max=4.0, step=0.01, round=False, advanced=True), + ], + outputs=[io.Noise.Output()], + ) + + @classmethod + def execute(cls, noise_seed, color, alpha, energy_scale) -> io.NodeOutput: + a = alpha if color == "custom" else COLOR_ALPHA[color] + params = ColorParams(mode="parametric", alpha_start=a, alpha_end=a, energy_scale=energy_scale) + return io.NodeOutput(Noise_ColoredInitial(noise_seed, params)) + + get_noise = execute diff --git a/src/nodes/sampler_select.py b/src/nodes/sampler_select.py new file mode 100644 index 0000000..a014042 --- /dev/null +++ b/src/nodes/sampler_select.py @@ -0,0 +1,90 @@ +"""ColoredNoise_SamplerSelect — outputs a SAMPLER that injects per-step colored noise.""" +import folder_paths +from comfy_api.latest import io + +from ..sampling.registry import stochastic_samplers +from ..sampling.colored_ksampler import build_colored_sampler +from ..engine.colored_noise import ColorParams +from ..engine.schedule import load_gamma_matrix + +GAMMA_FOLDER = "colored_noise_gamma" +NONE = "none (parametric)" + + +def gamma_choices(): + try: + return [NONE] + folder_paths.get_filename_list(GAMMA_FOLDER) + except Exception: + return [NONE] + + +def resolve_color_params(mode, alpha_start, alpha_end, interpolation, exp_sharpness, + gamma_matrix, gamma_divider, gamma_shaping, power_gamma, + alpha_tilting, energy_scale): + """Build a ColorParams, loading the gamma matrix only when gamma mode is active.""" + gm = None + if mode == "gamma_matrix" and gamma_matrix != NONE: + path = folder_paths.get_full_path(GAMMA_FOLDER, gamma_matrix) + if path is None: + raise FileNotFoundError(f"gamma matrix '{gamma_matrix}' not found in {GAMMA_FOLDER}") + gm = load_gamma_matrix(path) + return ColorParams( + mode=("gamma_matrix" if (mode == "gamma_matrix" and gm is not None) else "parametric"), + alpha_start=alpha_start, alpha_end=alpha_end, interpolation=interpolation, + exp_sharpness=exp_sharpness, gamma_matrix=gm, gamma_divider=gamma_divider, + gamma_shaping=gamma_shaping, power_gamma=power_gamma, alpha_tilting=alpha_tilting, + energy_scale=energy_scale, + ) + + +def default_sampler(samplers): + return "dpmpp_2m_sde" if "dpmpp_2m_sde" in samplers else samplers[0] + + +class ColoredNoise_SamplerSelect(io.ComfyNode): + @classmethod + def define_schema(cls): + samplers = list(stochastic_samplers().keys()) + return io.Schema( + node_id="ColoredNoise_SamplerSelect", + display_name="Colored Noise Sampler", + category="sampling/colored_noise", + inputs=[ + io.Combo.Input("base_sampler", options=samplers, default=default_sampler(samplers), + tooltip="Stochastic base sampler whose per-step noise is colored."), + io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False), + io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False), + io.Combo.Input("mode", options=["parametric", "gamma_matrix"], default="parametric"), + io.Float.Input("alpha_start", default=0.0, min=-8.0, max=8.0, step=0.05, round=False, + tooltip="Spectral exponent at the start (high sigma). 0=white, +red/pink, -blue/violet."), + io.Float.Input("alpha_end", default=-1.0, min=-8.0, max=8.0, step=0.05, round=False, + tooltip="Spectral exponent at the end (low sigma). Interpolated over the trajectory."), + io.Combo.Input("interpolation", options=["linear", "exponential"], default="linear"), + io.Float.Input("exp_sharpness", default=4.0, min=0.1, max=16.0, step=0.1, round=False, advanced=True), + io.Combo.Input("gamma_matrix", options=gamma_choices(), default=NONE, advanced=True, + tooltip="gamma_matrix mode only. Drop .pt matrices into models/colored_noise_gamma."), + io.Float.Input("gamma_divider", default=1.0, min=0.1, max=10.0, step=0.01, round=False, advanced=True), + io.Combo.Input("gamma_shaping", options=["none", "sqrt", "power"], default="none", advanced=True), + io.Float.Input("power_gamma", default=1.0, min=0.1, max=8.0, step=0.05, round=False, advanced=True), + io.Float.Input("alpha_tilting", default=0.0, min=-8.0, max=8.0, step=0.05, round=False, advanced=True), + io.Float.Input("energy_scale", default=1.0, min=0.0, max=4.0, step=0.01, round=False, advanced=True, + tooltip="Scales noise std after unit-variance renorm. 1.0 = neutral; >1 injects more 'heat'."), + io.Combo.Input("noise_device", options=["gpu", "cpu"], default="gpu", advanced=True), + ], + outputs=[io.Sampler.Output()], + ) + + @classmethod + def execute(cls, base_sampler, eta, s_noise, mode, alpha_start, alpha_end, interpolation, + exp_sharpness, gamma_matrix, gamma_divider, gamma_shaping, power_gamma, + alpha_tilting, energy_scale, noise_device) -> io.NodeOutput: + base_fn = stochastic_samplers()[base_sampler] + params = resolve_color_params(mode, alpha_start, alpha_end, interpolation, exp_sharpness, + gamma_matrix, gamma_divider, gamma_shaping, power_gamma, + alpha_tilting, energy_scale) + # eta/s_noise/r/solver_type are filtered to what the chosen base actually accepts. + base_opts = {"eta": eta, "s_noise": s_noise, "r": 0.5, "solver_type": "midpoint"} + sampler = build_colored_sampler(base_fn, base_opts, params) + return io.NodeOutput(sampler) + + get_sampler = execute diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 0000000..9e3237c --- /dev/null +++ b/tests/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +# rootdir resolves here (tests/), above which the package's V3 __init__.py is never imported. +addopts = --import-mode=importlib diff --git a/tests/test_schedule.py b/tests/test_schedule.py index 7f3e371..ccc29a6 100644 --- a/tests/test_schedule.py +++ b/tests/test_schedule.py @@ -37,8 +37,6 @@ def test_gamma_scaling_residual(): def test_load_real_gamma_matrix(): - pkg = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - # pkg above resolves to .../ComfyUI/custom_nodes ; rebuild to the package dir: pkg_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) p = os.path.join(pkg_dir, "gamma_matrices", "gamma_matrix_scaled.pt") g = schedule.load_gamma_matrix(p) diff --git a/tests/test_spectral.py b/tests/test_spectral.py index ab64cc6..f088b76 100644 --- a/tests/test_spectral.py +++ b/tests/test_spectral.py @@ -1,4 +1,3 @@ -import math import torch from src.engine import spectral