356 lines
12 KiB
Python
356 lines
12 KiB
Python
"""
|
||
BoyoNodes — Experimental Pack
|
||
Node: BoyoExperimentalSigmaPlot
|
||
|
||
Generates a sigma schedule for flow models (Wan, LTX, Krea2).
|
||
Sigmas live in [0, 1] — this is non-negotiable for flow models.
|
||
sigma_max=1.0, sigma_min is a small positive value near zero.
|
||
|
||
Outputs:
|
||
- SIGMAS → wire into SamplerCustomAdvanced
|
||
- IMAGE → wire into PreviewImage
|
||
|
||
Strategies
|
||
----------
|
||
simple : Linear [1.0 → sigma_min]. Matches BasicScheduler/simple exactly.
|
||
linear : Same as simple — explicit alias.
|
||
beta : Beta distribution spacing (alpha=0.6, beta=0.6). Middle-weighted.
|
||
beta57 : Beta (alpha=5, beta=7). Skewed toward low-noise / detail end.
|
||
karras : Karras curve in [0,1] space.
|
||
exponential : Exponential decay in [0,1] space.
|
||
chaotic_logistic: Logistic map drives spacing. Seeded, deterministic chaos.
|
||
chaotic_henon : Hénon map drives spacing. Different chaotic texture.
|
||
non_monotonic : Simple base with seeded upward reinjection events.
|
||
|
||
Author : Boyo / Dragon Diffusion UK Ltd
|
||
Pack : BoyoNodes (experimental prefix)
|
||
"""
|
||
|
||
import torch
|
||
import numpy as np
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
|
||
from comfy.k_diffusion import sampling as k_diff_sampling
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Sigma generation — all return float32 tensor length (steps+1), last = 0
|
||
# All operate in [0, 1] space — flow models only.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _sigmas_simple(steps, sigma_min):
|
||
"""Linear descent 1.0 → sigma_min. Matches BasicScheduler simple exactly."""
|
||
sigmas = torch.linspace(1.0, sigma_min, steps)
|
||
return torch.cat([sigmas, sigmas.new_zeros(1)])
|
||
|
||
|
||
def _sigmas_linear(steps, sigma_min):
|
||
return _sigmas_simple(steps, sigma_min)
|
||
|
||
|
||
def _beta_quantiles(steps, alpha, beta_v):
|
||
from scipy.special import betaincinv
|
||
p = np.linspace(0.01, 0.99, steps)
|
||
q = betaincinv(alpha, beta_v, p)
|
||
return torch.tensor(q, dtype=torch.float32)
|
||
|
||
|
||
def _sigmas_beta(steps, sigma_min):
|
||
quantiles = _beta_quantiles(steps, 0.6, 0.6)
|
||
quantiles = (quantiles - quantiles.min()) / (quantiles.max() - quantiles.min())
|
||
sigmas = sigma_min + (1.0 - quantiles) * (1.0 - sigma_min)
|
||
sigmas, _ = sigmas.sort(descending=True)
|
||
return torch.cat([sigmas, sigmas.new_zeros(1)])
|
||
|
||
|
||
def _sigmas_beta57(steps, sigma_min):
|
||
quantiles = _beta_quantiles(steps, 5.0, 7.0)
|
||
quantiles = (quantiles - quantiles.min()) / (quantiles.max() - quantiles.min())
|
||
sigmas = sigma_min + (1.0 - quantiles) * (1.0 - sigma_min)
|
||
sigmas, _ = sigmas.sort(descending=True)
|
||
return torch.cat([sigmas, sigmas.new_zeros(1)])
|
||
|
||
|
||
def _sigmas_karras(steps, sigma_min):
|
||
return k_diff_sampling.get_sigmas_karras(
|
||
steps, sigma_min, 1.0, rho=7.0
|
||
)
|
||
|
||
|
||
def _sigmas_exponential(steps, sigma_min):
|
||
sigmas = torch.linspace(np.log(1.0), np.log(sigma_min), steps).exp()
|
||
return torch.cat([sigmas, sigmas.new_zeros(1)])
|
||
|
||
|
||
def _sigmas_chaotic_logistic(steps, sigma_min, seed):
|
||
rng = np.random.default_rng(seed)
|
||
x = rng.uniform(0.01, 0.99)
|
||
r = 3.99
|
||
raw = []
|
||
for _ in range(steps):
|
||
x = r * x * (1 - x)
|
||
raw.append(x)
|
||
raw = np.array(raw)
|
||
rescaled = raw * (1.0 - sigma_min) + sigma_min
|
||
rescaled = np.sort(rescaled)[::-1].copy()
|
||
sigmas = torch.tensor(rescaled, dtype=torch.float32)
|
||
return torch.cat([sigmas, sigmas.new_zeros(1)])
|
||
|
||
|
||
def _sigmas_chaotic_henon(steps, sigma_min, seed):
|
||
rng = np.random.default_rng(seed)
|
||
x = rng.uniform(-0.5, 0.5)
|
||
y = rng.uniform(-0.5, 0.5)
|
||
a, b = 1.4, 0.3
|
||
raw = []
|
||
for _ in range(steps):
|
||
x, y = 1 - a * x**2 + y, b * x
|
||
raw.append(x)
|
||
raw = np.array(raw)
|
||
r_min, r_max = raw.min(), raw.max()
|
||
if r_max - r_min < 1e-8:
|
||
raw = np.linspace(0.99, 0.01, steps)
|
||
else:
|
||
raw = (raw - r_min) / (r_max - r_min)
|
||
rescaled = raw * (1.0 - sigma_min) + sigma_min
|
||
rescaled = np.sort(rescaled)[::-1].copy()
|
||
sigmas = torch.tensor(rescaled, dtype=torch.float32)
|
||
return torch.cat([sigmas, sigmas.new_zeros(1)])
|
||
|
||
|
||
def _sigmas_non_monotonic(steps, sigma_min, seed):
|
||
base = _sigmas_simple(steps, sigma_min)
|
||
sigmas = base.clone()
|
||
rng = np.random.default_rng(seed)
|
||
max_events = max(1, steps // 3)
|
||
n_events = int(rng.integers(1, max_events + 1))
|
||
eligible = list(range(2, steps - 1))
|
||
if not eligible:
|
||
return sigmas
|
||
event_steps = rng.choice(eligible, size=min(n_events, len(eligible)), replace=False)
|
||
for idx in event_steps:
|
||
bump = float(sigmas[idx]) * rng.uniform(0.05, 0.20)
|
||
sigmas[idx] = min(float(sigmas[idx]) + bump, 1.0)
|
||
return sigmas
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Strategy dispatcher
|
||
# ---------------------------------------------------------------------------
|
||
|
||
STRATEGIES = [
|
||
"simple",
|
||
"beta",
|
||
"beta57",
|
||
"linear",
|
||
"karras",
|
||
"exponential",
|
||
"chaotic_logistic",
|
||
"chaotic_henon",
|
||
"non_monotonic",
|
||
]
|
||
|
||
|
||
def _build_sigmas(steps, sigma_min, strategy, seed):
|
||
if strategy in ("simple", "linear"):
|
||
return _sigmas_simple(steps, sigma_min)
|
||
elif strategy == "beta":
|
||
return _sigmas_beta(steps, sigma_min)
|
||
elif strategy == "beta57":
|
||
return _sigmas_beta57(steps, sigma_min)
|
||
elif strategy == "karras":
|
||
return _sigmas_karras(steps, sigma_min)
|
||
elif strategy == "exponential":
|
||
return _sigmas_exponential(steps, sigma_min)
|
||
elif strategy == "chaotic_logistic":
|
||
return _sigmas_chaotic_logistic(steps, sigma_min, seed)
|
||
elif strategy == "chaotic_henon":
|
||
return _sigmas_chaotic_henon(steps, sigma_min, seed)
|
||
elif strategy == "non_monotonic":
|
||
return _sigmas_non_monotonic(steps, sigma_min, seed)
|
||
else:
|
||
raise ValueError(f"Unknown strategy: {strategy}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Plot generation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
W, H = 512, 320
|
||
PAD_L, PAD_R, PAD_T, PAD_B = 60, 20, 20, 50
|
||
BG_COLOUR = (18, 18, 22)
|
||
GRID_COLOUR = (45, 45, 55)
|
||
LINE_COLOUR = (120, 220, 180)
|
||
DOT_COLOUR = (220, 160, 80)
|
||
TEXT_COLOUR = (180, 180, 190)
|
||
LABEL_COLOUR = (120, 120, 140)
|
||
|
||
|
||
def _sigma_to_px(sigma, sigma_max, sigma_min):
|
||
frac = (sigma - sigma_min) / max(sigma_max - sigma_min, 1e-8)
|
||
return int(PAD_T + (1.0 - frac) * (H - PAD_T - PAD_B))
|
||
|
||
|
||
def _step_to_px(step_idx, n_points):
|
||
plot_w = W - PAD_L - PAD_R
|
||
frac = step_idx / max(n_points - 1, 1)
|
||
return int(PAD_L + frac * plot_w)
|
||
|
||
|
||
def _render_plot(sigmas_tensor, strategy, steps, sigma_min):
|
||
sigma_max = 1.0
|
||
img = Image.new("RGB", (W, H), BG_COLOUR)
|
||
draw = ImageDraw.Draw(img)
|
||
|
||
try:
|
||
font_sm = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 11)
|
||
font_md = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 13)
|
||
except Exception:
|
||
font_sm = ImageFont.load_default()
|
||
font_md = font_sm
|
||
|
||
for frac in [0.0, 0.33, 0.66, 1.0]:
|
||
sigma_val = sigma_min + frac * (sigma_max - sigma_min)
|
||
y = _sigma_to_px(sigma_val, sigma_max, sigma_min)
|
||
draw.line([(PAD_L, y), (W - PAD_R, y)], fill=GRID_COLOUR, width=1)
|
||
draw.text((2, y - 7), f"{sigma_val:.3f}", fill=LABEL_COLOUR, font=font_sm)
|
||
|
||
draw.line([(PAD_L, PAD_T), (PAD_L, H - PAD_B)], fill=TEXT_COLOUR, width=1)
|
||
draw.line([(PAD_L, H - PAD_B), (W - PAD_R, H - PAD_B)], fill=TEXT_COLOUR, width=1)
|
||
|
||
vals = sigmas_tensor[:-1].tolist()
|
||
n = len(vals)
|
||
points = [(_step_to_px(i, n), _sigma_to_px(v, sigma_max, sigma_min)) for i, v in enumerate(vals)]
|
||
|
||
if len(points) > 1:
|
||
draw.line(points, fill=LINE_COLOUR, width=2)
|
||
|
||
for px, py in points:
|
||
r = 4
|
||
draw.ellipse([(px - r, py - r), (px + r, py + r)], fill=DOT_COLOUR)
|
||
|
||
for i in range(n):
|
||
px = _step_to_px(i, n)
|
||
draw.text((px - 3, H - PAD_B + 6), str(i + 1), fill=LABEL_COLOUR, font=font_sm)
|
||
|
||
title = f"{strategy} | {steps} steps | σ 1.000→{sigma_min:.4f} [flow]"
|
||
draw.text((PAD_L, 4), title, fill=TEXT_COLOUR, font=font_md)
|
||
|
||
for i, (px, py) in enumerate(points):
|
||
label = f"{vals[i]:.3f}"
|
||
offset_y = -18 if py > PAD_T + 20 else 10
|
||
draw.text((px - 14, py + offset_y), label, fill=LINE_COLOUR, font=font_sm)
|
||
|
||
return img
|
||
|
||
|
||
def _pil_to_comfy_image(pil_img):
|
||
arr = np.array(pil_img).astype(np.float32) / 255.0
|
||
tensor = torch.from_numpy(arr).unsqueeze(0)
|
||
return tensor
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ComfyUI node class
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class BoyoExperimentalSigmaPlot:
|
||
"""
|
||
Sigma schedule generator for flow models (Wan, LTX, Krea2).
|
||
Sigmas always in [0, 1] — the correct range for flow/rectified models.
|
||
|
||
Outputs SIGMAS for SamplerCustomAdvanced and IMAGE for PreviewImage.
|
||
"""
|
||
|
||
CATEGORY = "BoyoNodes/Experimental"
|
||
RETURN_TYPES = ("SIGMAS", "IMAGE")
|
||
RETURN_NAMES = ("sigmas", "plot")
|
||
FUNCTION = "generate"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"steps": (
|
||
"INT",
|
||
{
|
||
"default": 8,
|
||
"min": 4,
|
||
"max": 12,
|
||
"step": 1,
|
||
"tooltip": "Sampling steps. Flow model turbo range: 4–12.",
|
||
},
|
||
),
|
||
"sigma_min": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.03,
|
||
"min": 0.001,
|
||
"max": 0.5,
|
||
"step": 0.001,
|
||
"round": False,
|
||
"tooltip": (
|
||
"Lower sigma bound. Flow models: typically 0.03–0.125.\n"
|
||
"BasicScheduler/simple at 8 steps ends at 0.125.\n"
|
||
"Lower values spend more time in fine detail."
|
||
),
|
||
},
|
||
),
|
||
"strategy": (
|
||
STRATEGIES,
|
||
{
|
||
"default": "simple",
|
||
"tooltip": (
|
||
"simple/linear : matches BasicScheduler — use as baseline.\n"
|
||
"beta : middle-weighted spacing.\n"
|
||
"beta57 : skewed toward detail end.\n"
|
||
"karras : karras curve in [0,1] space.\n"
|
||
"exponential : exponential decay.\n"
|
||
"chaotic_* : chaos-map driven spacing (seeded).\n"
|
||
"non_monotonic : linear base with reinjection events."
|
||
),
|
||
},
|
||
),
|
||
"seed": (
|
||
"INT",
|
||
{
|
||
"default": 42,
|
||
"min": 0,
|
||
"max": 0xFFFFFFFFFFFFFFFF,
|
||
"tooltip": "Seed for chaotic and non_monotonic strategies.",
|
||
},
|
||
),
|
||
},
|
||
}
|
||
|
||
def generate(self, steps, sigma_min, strategy, seed):
|
||
if sigma_min >= 1.0:
|
||
raise ValueError(
|
||
f"BoyoExperimentalSigmaPlot: sigma_min ({sigma_min}) must be less than 1.0"
|
||
)
|
||
|
||
sigmas = _build_sigmas(steps, sigma_min, strategy, seed)
|
||
|
||
print(
|
||
f"[BoyoExperimentalSigmaPlot] {strategy} | {steps} steps\n"
|
||
f" sigmas: {[round(float(s), 4) for s in sigmas.tolist()]}"
|
||
)
|
||
|
||
plot_pil = _render_plot(sigmas, strategy, steps, sigma_min)
|
||
plot_tensor = _pil_to_comfy_image(plot_pil)
|
||
|
||
return (sigmas, plot_tensor)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node registration
|
||
# ---------------------------------------------------------------------------
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"BoyoExperimentalSigmaPlot": BoyoExperimentalSigmaPlot,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"BoyoExperimentalSigmaPlot": "🧪 Sigma Plot (Experimental)",
|
||
}
|