Reithan
1be05ad7d2
## Summary
Collapses the experimental three-mode EPS switch
(`current`/`identity`/`true_v`) into a single, mathematically-correct
v-prediction-space path used by **all** variance-preserving
parameterizations — EPS, v-pred, and x0.
The sampler hook delivers `cond`/`uncond` as `x - x0` for every VP
parameterization, so NRS recovers the true velocity `v = (cond -
A)/factor` (with `A = x·σ²/(σ²+1)`, `factor = σ/√(σ²+1)`), runs its
geometry in v-space, and inverts exactly on return. This replaces the
prior EPS-only affine that operated on the wrong input space.
## Version: 1.1.0
MINOR, not a patch: this adds **X0 (sample) prediction support** — a new
backward-compatible capability — and **changes the default parameters**
(2/5/0.75 → 2/4/0.5). The node interface is unchanged (`eps_mode` never
shipped in 1.0.0), so it is **not** breaking/MAJOR. `pyproject.toml` and
`NRS/nodes_NRS.py __version__` are bumped to `1.1.0` in this PR; the
registry publish is triggered by the merge/release process, not by this
branch.
## Changes
- **`NRS/nodes_NRS.py`**: two-branch conversion — FLOW/CONST operated
natively (identity); all VP types (EPS/V/X0, UNKNOWN→VP) share one
true-v-space transform. Implements X0 (removes the
`NotImplementedError`). Removes `eps_mode` from every method and the
`nrs()` call. New defaults **skew 2 / stretch 4 / squash 0.5**.
- **`scripts/negative_rejection_steering_script.py`** (Forge/A1111):
removes the `eps_mode` test radio, process args, XYZ axis, and
extra-param; reverts `patch()` to 4 args; defaults 2/4/0.5.
- **`tests/test_pred_type.py`**: rewritten for the unified path — FLOW
is the sole native branch; EPS/V/X0/UNKNOWN all run the algebra.
- **`README.md`**: rewrites the V-Space step (now covers EPS/v-pred/x0
identically; flow-matching stays native) and updates the default to
2/4/0.5.
- **`CHANGELOG.md`**: `[1.1.0]` entry.
- **`pyproject.toml` / `NRS/nodes_NRS.py`**: version bumped 1.0.0 →
1.1.0.
## Rationale
Empirical A/B on an EPS checkpoint showed the true-v path is weakly
dominant: equal in normal ranges, burn-resistant at the failure edge,
and unlocks an on-style high-detail regime via squash (squash becomes a
detail/steering knob because the offset `A` is unscaled).
## Validation
- 39 passed / 9 skipped locally (torch tests skip; CI runs them).
Branch-coverage gate passed on push.
- ⚠️ **v-pred and x0 in v-space are math-validated but not yet
image-validated** — recommend a spot-check on a v-pred checkpoint before
merge/release.
350 lines
14 KiB
Python
350 lines
14 KiB
Python
# ruff: noqa: N999 -- filename predates NRS/__init__.py; mixed-case "nodes_NRS"
|
||
# only became checkable once NRS became a regular (non-namespace) package here.
|
||
# Renaming it is out of scope (would break existing imports); pyproject.toml's
|
||
# per-file-ignores already carve out N802/N804 for this same file.
|
||
import logging
|
||
import math
|
||
from enum import Enum, auto
|
||
|
||
import torch
|
||
|
||
try:
|
||
import comfy.utils as _comfy_utils
|
||
except Exception:
|
||
_comfy_utils = None
|
||
|
||
# Must be bumped together with the `version` field in pyproject.toml at release time.
|
||
__version__ = "1.1.0"
|
||
|
||
|
||
def _unpack_latents(combined, latent_shapes):
|
||
"""Split a flat packed latent [B, 1, N] back into its per-stream tensors.
|
||
|
||
Mirrors comfy.utils.unpack_latents: for each shape in latent_shapes, take
|
||
math.prod(shape[1:]) elements off the last dim and reshape that [B, 1, n]
|
||
slice back to `shape`.
|
||
"""
|
||
streams = []
|
||
offset = 0
|
||
for shape in latent_shapes:
|
||
n = math.prod(shape[1:])
|
||
chunk = combined[:, :, offset : offset + n]
|
||
streams.append(chunk.reshape(shape))
|
||
offset += n
|
||
return streams
|
||
|
||
|
||
def _pack_latents(streams):
|
||
"""Pack a list of per-stream tensors [B, C, ...] into a flat [B, 1, N] tensor.
|
||
|
||
Mirrors comfy.utils.pack_latents: each stream is reshaped to (B, 1, -1)
|
||
and concatenated on the last dim.
|
||
"""
|
||
flat = [s.reshape(s.shape[0], 1, -1) for s in streams]
|
||
return torch.cat(flat, dim=-1)
|
||
|
||
|
||
# fmt: off
|
||
class PredictionType(Enum):
|
||
EPS = auto() # ε-prediction
|
||
V = auto() # v-prediction
|
||
X0 = auto() # x₀-prediction
|
||
FLOW = auto() # flow-matching / velocity — operated natively, no VP conversion
|
||
UNKNOWN = auto() # couldn’t detect / new scheduler
|
||
|
||
|
||
_RAW_TO_ENUM = {
|
||
"eps": PredictionType.EPS,
|
||
"epsilon": PredictionType.EPS,
|
||
"flux": PredictionType.FLOW,
|
||
"chroma": PredictionType.FLOW,
|
||
"flow": PredictionType.FLOW, # FLOW models (WAN, etc.) operated natively
|
||
"wan": PredictionType.FLOW, # WAN21 is FLOW-based
|
||
"const": PredictionType.FLOW, # CONST prediction class used in FLOW models
|
||
"v": PredictionType.V,
|
||
"v_prediction": PredictionType.V,
|
||
"x0": PredictionType.X0,
|
||
"sample": PredictionType.X0,
|
||
}
|
||
# fmt: on
|
||
|
||
|
||
class NRS:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL", {"tooltip": "Input model to apply NRS to"}),
|
||
"skew": (
|
||
"FLOAT",
|
||
{
|
||
"default": 2.00,
|
||
"min": -30.0,
|
||
"max": 30.0,
|
||
"step": 0.01,
|
||
"tooltip": "Changes the 'direction' of generation, steering away from negative prompt elements. Start with CFG/2.",
|
||
},
|
||
),
|
||
"stretch": (
|
||
"FLOAT",
|
||
{
|
||
"default": 4.00,
|
||
"min": -30.0,
|
||
"max": 30.0,
|
||
"step": 0.01,
|
||
"tooltip": "Intensifies positive prompt elements. Start with your normal CFG value.",
|
||
},
|
||
),
|
||
"squash": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.50,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "Softens Skew/Stretch effects, adding micro-detailing. Keep low initially.",
|
||
},
|
||
),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL",)
|
||
FUNCTION = "patch"
|
||
|
||
CATEGORY = "advanced/model"
|
||
|
||
DESCRIPTION = "Negative Rejection Steering (NRS) replaces CFG with more nuanced guidance. IMPORTANT: Set your KSampler CFG to any value (it will be ignored). Connect your model through this node before sampling."
|
||
|
||
def _get_pred_type(self, model) -> PredictionType:
|
||
"""
|
||
In order to support Comfy, Forge, and possibly other models
|
||
and various loaders.
|
||
Walk common wrappers until we find something that looks like a
|
||
prediction-type flag, then map it to the enum.
|
||
Defaults to EPS if all else fails.
|
||
"""
|
||
|
||
def _canon(p):
|
||
if p is None:
|
||
return ""
|
||
if isinstance(p, bytes):
|
||
p = p.decode(errors="ignore")
|
||
if isinstance(p, Enum):
|
||
p = p.name
|
||
return str(p).strip().lower()
|
||
|
||
# Breadth-first search through a few well-known wrappers.
|
||
queue, seen = [model], set()
|
||
|
||
while queue:
|
||
obj = queue.pop(0)
|
||
|
||
# 1) direct hit on this object ---------------------------------
|
||
for attr in ("model_type", "prediction_type", "parameterization"):
|
||
p = _canon(getattr(obj, attr, None))
|
||
if p:
|
||
pred_type = _RAW_TO_ENUM.get(p, PredictionType.UNKNOWN)
|
||
if pred_type != PredictionType.UNKNOWN:
|
||
logging.debug(
|
||
f"NRS._get_pred_type: Found prediction type '{p}' from attribute '{attr}' -> {pred_type}"
|
||
)
|
||
return pred_type
|
||
|
||
# 2) enqueue child containers we care about -------------------
|
||
for attr in ("model", "diffusion_model", "config", "scheduler", "inner_model", "model_sampling"):
|
||
child = getattr(obj, attr, None)
|
||
if child is not None and id(child) not in seen:
|
||
seen.add(id(child))
|
||
queue.append(child)
|
||
|
||
# 3) enhanced detection for FLOW models (WAN, Flux, etc.) -------
|
||
try:
|
||
# Check model_sampling class type for FLOW models
|
||
if hasattr(model, "model_sampling") and model.model_sampling is not None:
|
||
sampling_class_name = type(model.model_sampling).__name__.lower()
|
||
logging.debug(f"NRS._get_pred_type: Found model_sampling class: {sampling_class_name}")
|
||
|
||
# CONST class is used by FLOW models (WAN21, Flux, etc.)
|
||
if "const" in sampling_class_name:
|
||
logging.debug("NRS._get_pred_type: Detected FLOW model via CONST sampling class -> FLOW")
|
||
return PredictionType.FLOW
|
||
elif "v_prediction" in sampling_class_name:
|
||
logging.debug("NRS._get_pred_type: Detected V-prediction model via sampling class -> V")
|
||
return PredictionType.V
|
||
elif "eps" in sampling_class_name:
|
||
logging.debug("NRS._get_pred_type: Detected EPS model via sampling class -> EPS")
|
||
return PredictionType.EPS
|
||
|
||
# Check model.model.model_type enum for newer models
|
||
if hasattr(model, "model") and hasattr(model.model, "model_type"):
|
||
model_type_str = _canon(str(model.model.model_type))
|
||
logging.debug(f"NRS._get_pred_type: Found model.model.model_type: {model_type_str}")
|
||
|
||
if "flow" in model_type_str or "flux" in model_type_str:
|
||
logging.debug("NRS._get_pred_type: Detected FLOW/Flux model via model_type -> FLOW")
|
||
return PredictionType.FLOW
|
||
elif "v_prediction" in model_type_str:
|
||
logging.debug("NRS._get_pred_type: Detected V-prediction model via model_type -> V")
|
||
return PredictionType.V
|
||
elif "eps" in model_type_str:
|
||
logging.debug("NRS._get_pred_type: Detected EPS model via model_type -> EPS")
|
||
return PredictionType.EPS
|
||
|
||
except Exception as e:
|
||
logging.debug(f"NRS._get_pred_type: Exception during enhanced detection: {e}")
|
||
|
||
# 4) safe default (matches docstring promise) --------------------
|
||
logging.warning("NRS._get_pred_type: Could not determine prediction type for model. Using EPS as fallback.")
|
||
logging.debug(
|
||
f"NRS._get_pred_type: Model structure: {[attr for attr in dir(model) if not attr.startswith('_')]}"
|
||
)
|
||
return PredictionType.EPS
|
||
|
||
def _is_vp(self, pred_type):
|
||
"""VP (variance-preserving) parameterizations converted to v-space: EPS, V, X0.
|
||
|
||
UNKNOWN (and any unhandled type) falls back to VP/v-space. FLOW/CONST is the only
|
||
parameterization operated natively (see _convert_to_v_space).
|
||
"""
|
||
if pred_type in (PredictionType.EPS, PredictionType.V, PredictionType.X0):
|
||
return True
|
||
if pred_type == PredictionType.FLOW:
|
||
return False
|
||
logging.warning(f"NRS: unknown prediction type {pred_type}, treating as VP (v-space)")
|
||
return True
|
||
|
||
def _convert_to_v_space(self, x_orig, sig_root, sigma, cond, uncond, pred_type):
|
||
"""Convert the (x - x0) guidance vectors into v-prediction space before the NRS geometry.
|
||
|
||
The sampler hook delivers cond/uncond as `x - x0` for every parameterization (the
|
||
model's raw output is converted to a denoised x0 before NRS sees it), so the true
|
||
velocity is recovered the same way regardless of EPS/V/X0:
|
||
v = (cond - A)/factor = (x/(sigma^2+1) - x0) * sig_root/sigma
|
||
with A = x*sigma^2/(sigma^2+1), factor = sigma/sqrt(sigma^2+1).
|
||
|
||
FLOW/CONST is operated natively: `x - x0 = sigma*out` is a pure scalar multiple of
|
||
the model's velocity (no additive offset), and the NRS geometry is scale-invariant,
|
||
so identity already runs on the native prediction. There is no VP v-space for
|
||
flow-matching (its sigma is a [0,1] flow time, not a VP karras sigma).
|
||
"""
|
||
if not self._is_vp(pred_type):
|
||
logging.debug("NRS._convert_to_v_space: flow/const operated natively (identity)")
|
||
return cond, uncond
|
||
|
||
logging.debug("NRS._convert_to_v_space: converting VP prediction to v-space")
|
||
factor = sigma / sig_root
|
||
a_off = x_orig - x_orig / (sigma**2 + 1) # A = x*sigma^2/(sigma^2+1)
|
||
return (cond - a_off) / factor, (uncond - a_off) / factor
|
||
|
||
def _finalize_from_v_space(self, x_orig, x_final, sig_root, sigma, pred_type):
|
||
"""Invert _convert_to_v_space so the hook returns `x - x0_final`. Round-trips exactly."""
|
||
if not self._is_vp(pred_type):
|
||
logging.debug("NRS._finalize_from_v_space: flow/const operated natively (identity)")
|
||
return x_final
|
||
|
||
factor = sigma / sig_root
|
||
a_off = x_orig - x_orig / (sigma**2 + 1)
|
||
return a_off + x_final * factor
|
||
|
||
def _apply_guidance(self, x_orig, cond, uncond, sigma, skew, stretch, squash, pred_type):
|
||
"""Run the NRS geometry pipeline on a single (already-unpacked, channels-first) stream."""
|
||
sigma = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1))
|
||
sig_root = (sigma**2 + 1).sqrt()
|
||
|
||
# Convert (x - x0) guidance into v-space for all VP parameterizations (EPS/V/X0);
|
||
# FLOW/CONST runs natively.
|
||
nrs_cond, nrs_uncond = self._convert_to_v_space(x_orig, sig_root, sigma, cond, uncond, pred_type)
|
||
|
||
def _dot(a, b):
|
||
return (a * b).sum(dim=1, keepdim=True) # [B,C,W,H] => [B,1,W,H]
|
||
|
||
def _nrm2(v):
|
||
return _dot(v, v)
|
||
|
||
eps = torch.finfo(nrs_cond.dtype).eps
|
||
c_dot_c = _nrm2(nrs_cond) + eps # [B,1,W,H]
|
||
u_dot_c = _dot(nrs_uncond, nrs_cond) # [B,1,W,H]
|
||
u_on_c = (u_dot_c / c_dot_c) * nrs_cond # [B,1,W,H] * [B,C,H,W]
|
||
|
||
# Amplify Cond based on length compared to projection of uncond
|
||
proj_diff = nrs_cond - u_on_c
|
||
stretched = nrs_cond + (stretch * proj_diff)
|
||
|
||
# Skew/Steer Conf based on rejection of uncond on cond
|
||
u_rej_c = nrs_uncond - u_on_c
|
||
skewed = stretched - (skew * u_rej_c)
|
||
|
||
# Squash final length back down to original length of cond
|
||
cond_len = nrs_cond.norm(dim=1, keepdim=True)
|
||
nrs_len = skewed.norm(dim=1, keepdim=True) + eps
|
||
|
||
squash_scale = (1 - squash) + (squash * (cond_len / nrs_len))
|
||
x_final = skewed * squash_scale
|
||
|
||
return self._finalize_from_v_space(x_orig, x_final, sig_root, sigma, pred_type)
|
||
|
||
def patch(self, model, skew, stretch, squash):
|
||
pred_type = self._get_pred_type(model)
|
||
logging.info(f"NRS v{__version__}: prediction type detected -> {pred_type.name}")
|
||
warned = {"done": False}
|
||
|
||
def nrs(args):
|
||
logging.debug(f"NRS.nrs: Skew: {skew}, Stretch: {stretch}, Squash: {squash}")
|
||
cond = args["cond"]
|
||
uncond = args["uncond"]
|
||
x_orig = args["input"]
|
||
sigma = args["sigma"]
|
||
|
||
shapes = getattr(args["model"], "latent_shapes", None)
|
||
if shapes and len(shapes) > 1:
|
||
if _comfy_utils is not None and hasattr(_comfy_utils, "unpack_latents"):
|
||
cond_streams = _comfy_utils.unpack_latents(cond, shapes)
|
||
uncond_streams = _comfy_utils.unpack_latents(uncond, shapes)
|
||
x_streams = _comfy_utils.unpack_latents(x_orig, shapes)
|
||
else:
|
||
cond_streams = _unpack_latents(cond, shapes)
|
||
uncond_streams = _unpack_latents(uncond, shapes)
|
||
x_streams = _unpack_latents(x_orig, shapes)
|
||
else:
|
||
cond_streams, uncond_streams, x_streams = [cond], [uncond], [x_orig]
|
||
|
||
if not warned["done"]:
|
||
for stream in x_streams:
|
||
if stream.shape[1] == 1:
|
||
logging.warning(
|
||
f"NRS.nrs: routed stream has a singleton reduction axis {tuple(stream.shape)}; "
|
||
"NRS geometry (dot/proj/skew) will degenerate to a no-op on this stream."
|
||
)
|
||
warned["done"] = True
|
||
break
|
||
|
||
results = [
|
||
self._apply_guidance(
|
||
x_streams[i],
|
||
cond_streams[i],
|
||
uncond_streams[i],
|
||
sigma,
|
||
skew,
|
||
stretch,
|
||
squash,
|
||
pred_type,
|
||
)
|
||
for i in range(len(cond_streams))
|
||
]
|
||
|
||
if len(results) == 1:
|
||
return results[0]
|
||
|
||
if _comfy_utils is not None and hasattr(_comfy_utils, "pack_latents"):
|
||
return _comfy_utils.pack_latents(results)[0]
|
||
return _pack_latents(results)
|
||
|
||
m = model.clone()
|
||
m.set_model_sampler_cfg_function(nrs, True)
|
||
return (m,)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"NRS": NRS,
|
||
}
|