From 9b5c71d383b13f2f35a7e9b1eaf6bac4914bf204 Mon Sep 17 00:00:00 2001 From: Hu White Date: Thu, 25 Jun 2026 02:02:04 -0500 Subject: [PATCH] Add RMS renormalization option (fixes magnitude-collapse quality loss) Adds a non-breaking `renormalize` option that rescales the per-layer-weighted tensor to its input RMS, so tap ratios can be rebalanced without inflating the overall conditioning magnitude. Off by default; existing behavior is unchanged. Background: the global multiplier amplifies the whole tensor, so with the default weights + multiplier=4.0 the magnitude inflates ~8.7x (4x mult x ~2.2x from the per-layer gains) -- matching the community-reported quality collapse (likeness, prompt adherence, colour). With renormalize on, the multiplier becomes the sole magnitude control. Co-Authored-By: Claude --- .gitignore | 4 ++ nodes.py | 182 ++++++++++++++++++++++++++++------------------------- 2 files changed, 101 insertions(+), 85 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4e62a66 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.py[cod] +*$py.class +*.egg-info/ diff --git a/nodes.py b/nodes.py index 719d22b..67f1efb 100644 --- a/nodes.py +++ b/nodes.py @@ -1,85 +1,97 @@ -import torch - - -def _scale_cond_tensor(t: torch.Tensor, multiplier, per_layer_weights=None): - """Scale a conditioning tensor, optionally with per-layer weighting. - - Krea2 conditioning arrives as (B, seq, 12*2560) — the 12 Qwen3-VL taps flattened - into the feature dim. When per_layer_weights is given we reshape to - (B, seq, 12, D), apply a different gain to each tap and flatten back. - """ - if per_layer_weights is None: - return t * multiplier - - flat = t.shape[-1] - n_layers = len(per_layer_weights) - if n_layers > 1 and flat % n_layers == 0: - layer_dim = flat // n_layers - orig_dtype = t.dtype - t = t.float() - t = t.view(*t.shape[:-1], n_layers, layer_dim) - gains = torch.tensor(per_layer_weights, dtype=t.dtype, device=t.device) - t = t * gains.view(*([1] * (t.dim() - 2)), n_layers, 1) - t = t.view(*t.shape[:-2], flat) - return t.to(orig_dtype) * multiplier - return t * multiplier - - -def _parse_per_layer(s: str): - """Parse a comma-separated list of floats. Returns None if empty/invalid.""" - if not s: - return None - s = s.strip() - if not s: - return None - try: - vals = [float(x) for x in s.replace(";", ",").split(",") if x.strip() != ""] - except ValueError: - return None - if len(vals) < 2: - return None - return vals - - -def scale_conditioning(structure, multiplier, per_layer_weights=None): - """leaving masks / pooled output intact.""" - if isinstance(structure, list): - out = [] - for item in structure: - if isinstance(item, (list, tuple)) and len(item) == 2 \ - and isinstance(item[0], torch.Tensor) and isinstance(item[1], dict): - cond_t, extras = item - new_cond = _scale_cond_tensor(cond_t, multiplier, per_layer_weights) - out.append([new_cond, dict(extras)]) - else: - out.append(scale_conditioning(item, multiplier, per_layer_weights)) - return out - if isinstance(structure, torch.Tensor): - return _scale_cond_tensor(structure, multiplier, per_layer_weights) - if isinstance(structure, dict): - return {k: scale_conditioning(v, multiplier, per_layer_weights) - for k, v in structure.items()} - return structure - - -class ConditioningKrea2Rebalance: - - DEFAULT_WEIGHTS = "1.0,1.0,1.0,1.0,1.0,1.0,1.0,2.5,5.0,1.1,4.0,1.0" - - @classmethod - def INPUT_TYPES(cls): - return {"required": { - "conditioning": ("CONDITIONING",), - "multiplier": ("FLOAT", {"default": 4.0, "min": -1000000000.0, "max": 1000000000.0, "step": 0.01}), - "per_layer_weights": ("STRING", {"default": cls.DEFAULT_WEIGHTS, "multiline": False}), - }} - - RETURN_TYPES = ("CONDITIONING",) - RETURN_NAMES = ("conditioning",) - FUNCTION = "main" - CATEGORY = "conditioning" - - def main(self, conditioning, multiplier, per_layer_weights=None): - plw = _parse_per_layer(per_layer_weights) if per_layer_weights else None - c = scale_conditioning(conditioning, multiplier, per_layer_weights=plw) - return (c,) +import torch + + +def _scale_cond_tensor(t: torch.Tensor, multiplier, per_layer_weights=None, renormalize=False): + """Scale a conditioning tensor, optionally with per-layer weighting. + + Krea2 conditioning arrives as (B, seq, 12*2560) — the 12 Qwen3-VL taps flattened + into the feature dim. When per_layer_weights is given we reshape to + (B, seq, 12, D), apply a different gain to each tap and flatten back. + + When renormalize is True the output is rescaled so its RMS matches the input, so the + per-layer ratios change without inflating the overall conditioning magnitude. This + avoids the quality collapse (loss of likeness / prompt adherence, oversaturated + colour) that a large global multiplier causes by amplifying every tap at once. With + renormalize on, the global multiplier becomes the sole magnitude control. + """ + if per_layer_weights is None: + return t * multiplier + + flat = t.shape[-1] + n_layers = len(per_layer_weights) + if n_layers > 1 and flat % n_layers == 0: + layer_dim = flat // n_layers + orig_dtype = t.dtype + ref_rms = (t.float().pow(2).mean(dim=tuple(range(1, t.dim()))).sqrt() + if renormalize else None) + t = t.float() + t = t.view(*t.shape[:-1], n_layers, layer_dim) + gains = torch.tensor(per_layer_weights, dtype=t.dtype, device=t.device) + t = t * gains.view(*([1] * (t.dim() - 2)), n_layers, 1) + t = t.view(*t.shape[:-2], flat) + if renormalize and ref_rms is not None: + new_rms = t.pow(2).mean(dim=tuple(range(1, t.dim()))).sqrt().clamp_min(1e-8) + t = t * (ref_rms / new_rms).view(-1, *([1] * (t.dim() - 1))) + return t.to(orig_dtype) * multiplier + return t * multiplier + + +def _parse_per_layer(s: str): + """Parse a comma-separated list of floats. Returns None if empty/invalid.""" + if not s: + return None + s = s.strip() + if not s: + return None + try: + vals = [float(x) for x in s.replace(";", ",").split(",") if x.strip() != ""] + except ValueError: + return None + if len(vals) < 2: + return None + return vals + + +def scale_conditioning(structure, multiplier, per_layer_weights=None, renormalize=False): + """leaving masks / pooled output intact.""" + if isinstance(structure, list): + out = [] + for item in structure: + if isinstance(item, (list, tuple)) and len(item) == 2 \ + and isinstance(item[0], torch.Tensor) and isinstance(item[1], dict): + cond_t, extras = item + new_cond = _scale_cond_tensor(cond_t, multiplier, per_layer_weights, renormalize) + out.append([new_cond, dict(extras)]) + else: + out.append(scale_conditioning(item, multiplier, per_layer_weights, renormalize)) + return out + if isinstance(structure, torch.Tensor): + return _scale_cond_tensor(structure, multiplier, per_layer_weights, renormalize) + if isinstance(structure, dict): + return {k: scale_conditioning(v, multiplier, per_layer_weights, renormalize) + for k, v in structure.items()} + return structure + + +class ConditioningKrea2Rebalance: + + DEFAULT_WEIGHTS = "1.0,1.0,1.0,1.0,1.0,1.0,1.0,2.5,5.0,1.1,4.0,1.0" + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "conditioning": ("CONDITIONING",), + "multiplier": ("FLOAT", {"default": 4.0, "min": -1000000000.0, "max": 1000000000.0, "step": 0.01}), + "per_layer_weights": ("STRING", {"default": cls.DEFAULT_WEIGHTS, "multiline": False}), + "renormalize": ("BOOLEAN", {"default": False}), + }} + + RETURN_TYPES = ("CONDITIONING",) + RETURN_NAMES = ("conditioning",) + FUNCTION = "main" + CATEGORY = "conditioning" + + def main(self, conditioning, multiplier, per_layer_weights=None, renormalize=False): + plw = _parse_per_layer(per_layer_weights) if per_layer_weights else None + c = scale_conditioning(conditioning, multiplier, per_layer_weights=plw, renormalize=renormalize) + return (c,)