feat(nodes): three V3 nodes + ComfyExtension entrypoint + folder_paths + pytest config
This commit is contained in:
akashzeno
@@ -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"]
|
||||
@@ -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.
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -0,0 +1,3 @@
|
||||
[pytest]
|
||||
# rootdir resolves here (tests/), above which the package's V3 __init__.py is never imported.
|
||||
addopts = --import-mode=importlib
|
||||
@@ -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)
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import math
|
||||
import torch
|
||||
|
||||
from src.engine import spectral
|
||||
|
||||
Reference in New Issue
Block a user