nodes/Reithan-negative_rejection_steering
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix detection of model type (#11)

Browse Source 
- fix some scaling issues
This commit is contained in:
Reithan
2025-07-20 04:04:54 -07:00
committed by GitHub
parent 932c7b2136
commit e0988d3b24
2 changed files with 33 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
NRS/nodes_NRS.py
+23 -31
View File
@@ -1,4 +1,3 @@
import inspect
import logging
import torch
from enum import Enum, auto
@@ -52,36 +51,26 @@ class NRS:
return str(p).strip().lower()
# Breadth-first search through a few well-known wrappers.
queue = [model]
visited = set()
queue, seen = [model], set()
while queue:
obj = queue.pop(0)
if id(obj) in visited:
continue
visited.add(id(obj))
# 1) direct hit on this object ---------------------------------
p = _canon(getattr(obj, "parameterization", None)) # k-diffusion
if p:
return _RAW_TO_ENUM.get(p, PredictionType.UNKNOWN)
p = _canon(getattr(getattr(obj, "config", None), "prediction_type", None)) # diffusers
if p:
return _RAW_TO_ENUM.get(p, PredictionType.UNKNOWN)
p = _canon(getattr(obj, "prediction_type", None)) # rare misc
if p:
return _RAW_TO_ENUM.get(p, PredictionType.UNKNOWN)
for attr in ("model_type", "prediction_type", "parameterization"):
p = _canon(getattr(obj, attr, None))
if p:
return _RAW_TO_ENUM.get(p, PredictionType.UNKNOWN)
# 2) enqueue child containers we care about -------------------
for attr in ("model_sampling", "model", "diffusion_model", "scheduler"):
for attr in ("model", "diffusion_model", "config", "scheduler", "inner_model", "model_sampling"):
child = getattr(obj, attr, None)
if child is not None:
if child is not None and id(child) not in seen:
seen.add(id(child))
queue.append(child)
# 3) default ------------------------------------------------------
return PredictionType.EPS
return PredictionType.UNKNOWN
def _pre_scale_conditioning(self, x_orig, sigma, cond, uncond):
@@ -129,10 +118,14 @@ class NRS:
def nrs(args):
cond = args["cond"]
uncond = args["uncond"]
sigma = args["sigma"]
sigma = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1))
x_orig = args["input"]
self.__pred_type = self.__pred_type if self.__pred_type is not None else self._get_pred_type(model)
sigma = None
if self.__pred_type == PredictionType.V:
sigma = args["sigma"]
sigma = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1))
logging.debug(f"NRS.nrs: Skew: {skew}, Stretch: {stretch}, Squash: {squash}")
@@ -281,30 +274,29 @@ class NRS:
x_final = skewed * squash_scale
case "v0.5.0":
def _dot(a, b):
return (a*b).flatten(1).sum(dim=1, keepdim=True) # [B,1]
return (a*b).flatten(1).sum(dim=1, keepdim=True) # [B,C,W,H] => [B,1]
def _nrm2(v):
return _dot(v, v)
eps = eps = torch.finfo(eps_cond.dtype).eps
eps = torch.finfo(eps_cond.dtype).eps
c_dot_c = _nrm2(eps_cond) + eps # [B,1]
u_dot_c = _dot(eps_uncond, eps_cond) # [B,1]
u_on_c = (u_dot_c / c_dot_c).view(-1, 1, 1, 1) * eps_cond # [B,1,1,1] * [B,C,H,W]
u_rej_c = eps_uncond - u_on_c
proj_diff = eps_cond - u_on_c
# Amplify Cond based on length compared to projection of uncond
proj_diff = eps_cond - u_on_c
stretched = eps_cond + (stretch * proj_diff)
# Skew/Steer Conf based on rejection of uncond on cond
skewed = stretched - skew * u_rej_c
u_rej_c = eps_uncond - u_on_c
skewed = stretched - (skew * u_rej_c)
# Squash final length back down to original length of cond
cond_len = torch.sqrt(c_dot_c) # [B,1]
nrs_len = torch.sqrt(_nrm2(skewed) + eps) # [B,1]
nrs_len = torch.sqrt(_nrm2(skewed)) + eps # [B,1]
squash_scale = (1 - squash) + squash * (cond_len / nrs_len)
squash_scale = (1 - squash) + (squash * (cond_len / nrs_len))
x_final = skewed * squash_scale.view(-1, 1, 1, 1)
return self._post_scale_conditioning(x_orig, x_div, x_final, sigma)
note.md
+10
View File
@@ -0,0 +1,10 @@
# The 'purpose' or 'intention' behind our 3 knobs
## SKEW
This is the primary 'steering' knob. This will 'turn' the 'direction' the current denoising step is traveling in the latent space. If we define the 'default' (cond) direction as 'prior step -> cond' then we're just applying a 'lateral' skew to that direction to 'turn' it 'away' from the 'unintended' direction (uncond)
## STRETCH
This is the sister knob to Skew. This is the accelerator. We want to go 'faster' into the intended direction (cond) the less aligned it is with the unintended direction (uncond). Think of this like a combination of brakes + gas. If we're headed directly for a brick wall (uncond is in the same direction as cond), we want to apply no acceleration, or negative acceleration. If we're traveling directly away from danger (uncond is in the opposite direction of cond) then we want to stomp the gas and get as far away as we can. There's only 1 problem with this BASIC-level description: as we get further into generation, regardless of pos/neg promp, cond & uncon will naturally align to be the same vector. In the last stop of inference, cond and uncond will be basically identical if nothing has fucked up. So whatever math we apply here needs to take the progressive alignment of cond & uncond into account. That's why were/are scaling only on the projection difference right now, rather than the full projection.
## SQUASH
This is out 'safety' knob. Think of this like a 'limiter' in a car. This sets the 'top speed' we can go to some multiple of the 'default' speed the model would 'like to' go. i.e. whatever length of directional vector the model produces prior to any skewing or stretching is treated as the 'default' length with Squash=1.0 ensuring we only every go that 'speed' and no more, while Squash=0.0 lets us go any speed we want based on the other 2 knobs. GENERALLY we'll be leaving Squash at 0.0 unless we need it for specific generations.