Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1be05ad7d2 | ||
|
|
244f8c656f |
@@ -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
@@ -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))
|
||||
]
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user