Compare commits

..
2 Commits
Author SHA1 Message Date
Reithan 1be05ad7d2 v1.1.0: Unify NRS on a single v-space path (EPS/V/X0) + X0 support; defaults 2/4/0.5 (#44)
## 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.
2026-09-03 05:30:03 -07:00
Reithan 244f8c656f Fix NODE_DISPLAY_NAME_MAPPINGS typo and NRS skew log label (#43)
## Summary
Applies `.claude/patches/nrs-display-name-and-log-fix.patch`, fixing two
small defects:

- **`__init__.py`** — corrects the `NODE_DISPLAY_NAME_MAPPINS` typo to
`NODE_DISPLAY_NAME_MAPPINGS`. The misspelled name meant ComfyUI never
picked up the human-readable display name, and `__all__` exported a name
that didn't exist.
- **`scripts/negative_rejection_steering_script.py`** — the debug log
labeled `self.skew` as "Squash"; now correctly labeled "Skew".

## Test plan
- Pre-commit + pre-push hooks (ruff, pytest branch-coverage gate) passed
on push.
2026-09-03 01:49:16 -07:00
8 changed files with 112 additions and 104 deletions
+16
View File
@@ -4,6 +4,22 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
## [1.1.0] - Unreleased
### Added
- **X0 (sample) prediction support** (#44): x0-prediction models are now handled through the shared v-prediction-space path, alongside EPS and v-pred.
### Fixed
- **Correct v-space conversion for all variance-preserving parameterizations** (#44). The sampler hook delivers `cond`/`uncond` as `x - x0` for EPS, v-pred, and x0 alike, so NRS now recovers the true velocity `v = (cond - A)/factor` and runs its geometry in v-prediction space, then inverts exactly on return. This replaces the prior EPS-only affine, which operated on an incorrect input-space assumption. Flow-matching (FLOW/CONST) models remain operated natively — their prediction is already a pure scalar multiple of the velocity, so no conversion is applied.
### Changed / Upgrade notes
- **Default parameters changed** 2/5/0.75 → **2/4/0.5** (Skew/Stretch/Squash) in both the ComfyUI node and the A1111-family (Forge/reForge/Forge Neo) script.
- **v-prediction models now run the v-space conversion** instead of operating on the raw guidance. For typical config ranges the output change is expected to be minimal (verified on EPS; v-pred/x0 are math-validated but **not yet image-validated** — spot-check and retune if needed).
- **Reproducibility note:** the same seed + config may produce a slightly different image than 1.0.0 because of the corrected v-space handling and the new defaults.
## [1.0.0] - 2026-08-14
### Fixed
+56 -49
View File
@@ -14,7 +14,7 @@ except Exception:
_comfy_utils = None
# Must be bumped together with the `version` field in pyproject.toml at release time.
__version__ = "1.0.0"
__version__ = "1.1.0"
def _unpack_latents(combined, latent_shapes):
@@ -88,7 +88,7 @@ class NRS:
"stretch": (
"FLOAT",
{
"default": 5.00,
"default": 4.00,
"min": -30.0,
"max": 30.0,
"step": 0.01,
@@ -98,7 +98,7 @@ class NRS:
"squash": (
"FLOAT",
{
"default": 0.75,
"default": 0.50,
"min": 0.0,
"max": 1.0,
"step": 0.01,
@@ -200,60 +200,60 @@ class NRS:
)
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):
x_div = None
v_cond = cond
v_uncond = uncond
if pred_type in (PredictionType.V, PredictionType.FLOW):
logging.debug("NRS._convert_to_v_space: already in v/flow, no pre-scale needed")
pass # already in v space / flow-matching operated natively
elif pred_type == PredictionType.EPS:
# ε → v conversion
logging.debug("NRS._convert_to_v_space: generating x_div, v_cond, and v_uncond for eps")
x_div = x_orig / (sigma**2 + 1)
factor = sigma / sig_root
"""Convert the (x - x0) guidance vectors into v-prediction space before the NRS geometry.
v_cond = x_orig - (x_div - cond * factor)
v_uncond = x_orig - (x_div - uncond * factor)
elif pred_type == PredictionType.X0:
raise NotImplementedError("NRS._convert_to_v_space: x0-prediction not supported yet.")
else:
# Fallback: treat UNKNOWN as EPS and convert to V-space
logging.warning(f"NRS._convert_to_v_space: Unknown prediction type {pred_type}, treating as EPS")
logging.debug("NRS._convert_to_v_space: generating x_div, v_cond, and v_uncond for eps (fallback)")
x_div = x_orig / (sigma**2 + 1)
factor = sigma / sig_root
v_cond = x_orig - (x_div - cond * factor)
v_uncond = x_orig - (x_div - uncond * factor)
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).
return x_div, v_cond, v_uncond
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
def _finalize_from_v_space(self, x_orig, x_div, x_final, sig_root, sigma, pred_type):
nrs_result = x_final
if pred_type in (PredictionType.V, PredictionType.FLOW):
# already in v space / flow-matching operated natively
logging.debug("NRS._finalize_from_v_space: already in v/flow, no post-scale needed")
pass
elif pred_type == PredictionType.EPS:
# v → ε conversion
logging.debug("NRS._finalize_from_v_space: generating cfg_result for eps")
nrs_result = (x_div - (x_orig - x_final)) * (sig_root / sigma)
elif pred_type == PredictionType.X0:
raise NotImplementedError("NRS._finalize_from_v_space: x0-prediction not supported yet.")
else:
# Fallback: treat UNKNOWN as EPS and convert from V-space
logging.warning(f"NRS._finalize_from_v_space: Unknown prediction type {pred_type}, treating as EPS")
logging.debug("NRS._finalize_from_v_space: generating cfg_result for eps (fallback)")
nrs_result = (x_div - (x_orig - x_final)) * (sig_root / sigma)
return nrs_result
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()
# V and FLOW models are operated natively (identity); EPS models are converted to v-space.
x_div, nrs_cond, nrs_uncond = self._convert_to_v_space(x_orig, sig_root, sigma, cond, uncond, pred_type)
# 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]
@@ -281,7 +281,7 @@ class NRS:
squash_scale = (1 - squash) + (squash * (cond_len / nrs_len))
x_final = skewed * squash_scale
return self._finalize_from_v_space(x_orig, x_div, x_final, sig_root, sigma, pred_type)
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)
@@ -320,7 +320,14 @@ class NRS:
results = [
self._apply_guidance(
x_streams[i], cond_streams[i], uncond_streams[i], sigma, skew, stretch, squash, pred_type
x_streams[i],
cond_streams[i],
uncond_streams[i],
sigma,
skew,
stretch,
squash,
pred_type,
)
for i in range(len(cond_streams))
]
+2 -2
View File
@@ -20,7 +20,7 @@ NRS seeks to replace the 'naive' linear interpolation of Classifier Free Guidanc
<img align="right" src="Examples/NRS_graph.png" alt="Graph of NRS vs CFG" style="width: 40%; float: right;">
### NRS is Applied in Three Steps:
0. ***V-Space**: Optional pre-NRS step* If the model is not using v-prediction, we transform the EPS `cond` and `uncond` into v-prediction space before continuing, then revert to eps-space before return.
0. ***V-Space**: pre-NRS step* The sampler hands NRS its `cond`/`uncond` as `x - x0` for every variance-preserving parameterization (EPS, v-prediction, and x0-prediction alike), so NRS recovers the true velocity `v` from them and runs its geometry in v-prediction space, then inverts the transform before returning. This one v-space path handles EPS, v-pred, and x0 models identically. Flow-matching models (flux, chroma, wan, and other flow/CONST families) are operated natively — their prediction is already a pure scalar multiple of the velocity, so no v-space conversion is applied.
1. **Skewing**: The conditioned output tensor is skewed away from the direction of the rejection of the unconditioned tensor on the conditioned tensor. This lengthens the tensor in a direction perpendicular to its direction without affecting the positive guidance. The tensor is displaced by the rejection multiplied by the Skew parameter.[^1]
2. **Stretching**: The skewed tensor is stretched towards the direction of the original conditioned tensor based on its difference from the projection of uncond on cond. The stretch is multiplied by the Stretch parameter.[^1]
3. **Squashing**: The skewed and stretched tensor is rescaled towards the original length of the conditioned tensor. 100% squashing outputs the original length of the conditioned tensor simply 'steered' towards the skewed & squashed version's direction.[^1]
@@ -63,7 +63,7 @@ NRS seeks to replace the 'naive' linear interpolation of Classifier Free Guidanc
## Beginner How-To
1. Set Skew to 1/2 of your normal CFG Scale setting and Stretch to your full normal CFG Scale. Set Squash to 0.0.<br>
*Alternatively, try starting with the default of 2/5/0.75, or at 1/1/1 to get a baseline.*
*Alternatively, try starting with the default of 2/4/0.5, or at 1/1/1 to get a baseline.*
2. Test some outputs. Results should be similar in quality to CFG.
3. Adjust Skew to change the intensity of your outputs adherence to your positive and negative prompts. This primarily effects composition of the output.
4. Adjust Stretch to intensify your positive prompt's aspects and colors where they differ from the negative prompt. This primarily effects color and texture.
+1 -1
View File
@@ -1,5 +1,5 @@
from .NRS.nodes_NRS import *
NODE_CLASS_MAPPINGS = {"NRS": NRS}
NODE_DISPLAY_NAME_MAPPINS = {"NRS": "Negative Rejection Steering"}
NODE_DISPLAY_NAME_MAPPINGS = {"NRS": "Negative Rejection Steering"}
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+1 -1
View File
@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
name = "negative_rejection_steering"
description = "NRS seeks to replace the 'naive' linear interpolation of Classifier Free Guidance with a more nuanced and composable steering of the generation process with better mathematical basis."
authors = [{name = "Bryan O'Malley", email = "bo122081@hotmail.com"}]
version = "1.0.0"
version = "1.1.0"
requires-python = ">=3.10"
license = {file = "LICENSE"}
readme = "README.md"
@@ -15,8 +15,8 @@ class NRSScript(scripts.Script):
super().__init__()
self.enabled = False
self.skew = 2.00
self.stretch = 5.00
self.squash = 0.75
self.stretch = 4.00
self.squash = 0.50
sorting_priority = 5
@@ -100,7 +100,7 @@ class NRSScript(scripts.Script):
)
logging.debug(
f"NRS: Enabled: {self.enabled}, Squash: {self.skew}, Stretch: {self.stretch}, Squash: {self.squash}"
f"NRS: Enabled: {self.enabled}, Skew: {self.skew}, Stretch: {self.stretch}, Squash: {self.squash}"
)
return
+2 -6
View File
@@ -151,9 +151,7 @@ def test_fallback_no_latent_shapes_matches_single_stream_shape():
# Regression check: manually compute the single-stream result the same
# way the pre-split code path did, and confirm equality.
node = nrs_module.NRS()
expected = node._apply_guidance(
x_orig, cond, uncond, sigma, 2.0, 5.0, 0.75, nrs_module.PredictionType.EPS
)
expected = node._apply_guidance(x_orig, cond, uncond, sigma, 2.0, 5.0, 0.75, nrs_module.PredictionType.EPS)
assert torch.allclose(result, expected)
@@ -168,9 +166,7 @@ def test_single_stream_latent_shapes_also_matches():
result = _run_nrs(model, cond, uncond, x_orig, sigma)
node = nrs_module.NRS()
expected = node._apply_guidance(
x_orig, cond, uncond, sigma, 2.0, 5.0, 0.75, nrs_module.PredictionType.EPS
)
expected = node._apply_guidance(x_orig, cond, uncond, sigma, 2.0, 5.0, 0.75, nrs_module.PredictionType.EPS)
assert torch.allclose(result, expected)
+31 -42
View File
@@ -5,8 +5,12 @@ PR-3 reclassified the flow-matching family (flux, chroma, flow, wan, const)
from PredictionType.EPS onto a new PredictionType.FLOW, which is operated
natively (identity conversion, no VP ε<->v algebra). These tests pin that
post-reclassification behavior at both detection sites (the _RAW_TO_ENUM
dict and the enhanced-detection fallback in _get_pred_type), and cover the
FLOW identity round-trip through _convert_to_v_space / _finalize_from_v_space.
dict and the enhanced-detection fallback in _get_pred_type).
FLOW is the sole native path; every VP parameterization (EPS, V, X0, and the
UNKNOWN fallback) shares one ε/v/x0 -> v-space conversion through
_convert_to_v_space / _finalize_from_v_space. These tests cover the FLOW
identity round-trip and confirm the VP branches actually transform their inputs.
"""
import enum
@@ -153,64 +157,49 @@ class TestGetPredTypeEnhancedDetectionFallback:
assert NRS()._get_pred_type(model) == PredictionType.EPS
class TestConvertToVSpaceIdentityBranches:
"""FLOW and V are both pure identity conversions -- no ε<->v algebra runs,
so plain sentinel objects (no real tensor math) are enough to prove it.
class TestConvertToVSpaceBranches:
"""FLOW is the only native (identity) parameterization; every VP type
(EPS, V, X0, and the UNKNOWN fallback) now runs the shared ε/v/x0 -> v-space
algebra. FLOW identity needs no tensor math, so sentinel objects prove it;
the VP branches use MagicMock to confirm the algebra actually transforms.
"""
def test_flow_convert_is_identity(self):
node = NRS()
cond, uncond = object(), object()
x_div, v_cond, v_uncond = node._convert_to_v_space(
object(), object(), object(), cond, uncond, PredictionType.FLOW
)
assert x_div is None
v_cond, v_uncond = node._convert_to_v_space(object(), object(), object(), cond, uncond, PredictionType.FLOW)
assert v_cond is cond
assert v_uncond is uncond
def test_flow_finalize_is_identity(self):
node = NRS()
x_final = object()
result = node._finalize_from_v_space(object(), None, x_final, object(), object(), PredictionType.FLOW)
result = node._finalize_from_v_space(object(), x_final, object(), object(), PredictionType.FLOW)
assert result is x_final
def test_v_convert_is_identity(self):
"""Regression guard: PR-3 must not disturb the existing V path."""
node = NRS()
cond, uncond = object(), object()
x_div, v_cond, v_uncond = node._convert_to_v_space(
object(), object(), object(), cond, uncond, PredictionType.V
)
assert x_div is None
assert v_cond is cond
assert v_uncond is uncond
def test_v_finalize_is_identity(self):
node = NRS()
x_final = object()
result = node._finalize_from_v_space(object(), None, x_final, object(), object(), PredictionType.V)
assert result is x_final
def test_eps_convert_still_performs_algebra(self):
"""Regression guard: EPS must still run the ε->v conversion (x_div gets
computed, and cond/uncond are transformed rather than passed through).
"""
@pytest.mark.parametrize("pred_type", [PredictionType.EPS, PredictionType.V, PredictionType.X0])
def test_vp_convert_performs_algebra(self, pred_type):
"""EPS/V/X0 all run the ε->v conversion (cond/uncond are transformed,
not passed through)."""
node = NRS()
x_orig, sig_root, sigma = MagicMock(), MagicMock(), MagicMock()
cond, uncond = MagicMock(), MagicMock()
x_div, v_cond, v_uncond = node._convert_to_v_space(x_orig, sig_root, sigma, cond, uncond, PredictionType.EPS)
assert x_div is not None
v_cond, v_uncond = node._convert_to_v_space(x_orig, sig_root, sigma, cond, uncond, pred_type)
assert v_cond is not cond
assert v_uncond is not uncond
def test_eps_finalize_still_performs_algebra(self):
@pytest.mark.parametrize("pred_type", [PredictionType.EPS, PredictionType.V, PredictionType.X0])
def test_vp_finalize_performs_algebra(self, pred_type):
node = NRS()
x_orig, x_div, x_final, sig_root, sigma = (
MagicMock(),
MagicMock(),
MagicMock(),
MagicMock(),
MagicMock(),
)
result = node._finalize_from_v_space(x_orig, x_div, x_final, sig_root, sigma, PredictionType.EPS)
x_orig, x_final, sig_root, sigma = MagicMock(), MagicMock(), MagicMock(), MagicMock()
result = node._finalize_from_v_space(x_orig, x_final, sig_root, sigma, pred_type)
assert result is not x_final
def test_unknown_convert_falls_back_to_vp(self):
"""UNKNOWN (and any unhandled type) is treated as VP -> runs the algebra."""
node = NRS()
x_orig, sig_root, sigma = MagicMock(), MagicMock(), MagicMock()
cond, uncond = MagicMock(), MagicMock()
v_cond, v_uncond = node._convert_to_v_space(x_orig, sig_root, sigma, cond, uncond, PredictionType.UNKNOWN)
assert v_cond is not cond
assert v_uncond is not uncond