Add Post-CFG SHIFT model patcher

WASasquatch · WASasquatch · commit cf3714d80ced · 2025-09-29T08:45:59.000-07:00
diff --git a/README.md b/README.md
@@ -1,4 +1,4 @@
-# FreeU Advanced Plus
+# FreeU Advanced Plus (And Post-CFG SHIFT)
 Let's say you and I grab dinner, and movie after lunch? 🌃📺😏
  
 ![image](https://github.com/WASasquatch/FreeU_Advanced/assets/1151589/c1dc2ec9-e6a3-4d2d-bf81-697e5d5aabcb)
@@ -52,3 +52,36 @@ Let's say you and I grab dinner, and movie after lunch? 🌃📺😏
 ## :newspaper_roll: License
 
 Distributed under the MIT License. See `LICENSE` for more information.
+
+---
+
+## Post-CFG SHIFT (Flux)
+
+Post-CFG Stepwise Hybrid Inject + Fourier Tuning.
+
+- Runs after classifier-free guidance (CFG) merges cond/uncond each sampler step.
+- Modifies the sampler’s current denoised tensor (in VAE latent space in typical pipelines), not model weights.
+- Applies a hybrid blend between the denoised tensor and a scaled version, with optional frequency-domain shaping.
+
+### How it works
+1) Model predicts noise; CFG produces a denoised tensor for the current step.
+2) SHIFT blends `denoised` with `denoised * b` using the chosen `mode` and `blend`.
+3) Optionally applies `Fourier_filter` with per-scale controls.
+4) Applies a final `force_gain` multiplier.
+
+### Parameters
+- `mode` (combo): Blend strategy for `denoised` vs `denoised*b`.
+  - Useful: `inject` (strong), `stable_slerp` (smooth), `lerp` (linear), etc.
+- `blend` (float): Blend amount between base and scaled tensors.
+- `b` (float): Scale factor for the injected path. Higher = stronger effect.
+- `apply_fourier` (bool): Enable frequency-domain shaping.
+- `multiscale_mode` (combo): Preset shaping curves. Use stable options (e.g., Default, Pass-Through, Sharpen).
+- `multiscale_strength` (float): Intensity of multi-scale shaping.
+- `threshold` (int): Base radius in frequency mask.
+- `s` (float): Base scale value applied at `threshold` radius.
+- `force_gain` (float): Final multiplier to boost or attenuate the overall effect.
+- `debug_log` (bool): Prints one-time registration and periodic fire logs.
+
+### Notes
+- SHIFT is always-on in Flux; attention/forward-timestep/wrapper paths are disabled for stability.
+- If a multiscale preset yields flat/gray output, switch to a stable preset (e.g., Sharpen, Pass-Through) or tune `threshold`/`s`.
diff --git a/nodes.py b/nodes.py
@@ -3,6 +3,7 @@
 import torch
 import torch as th
 import torch.fft as fft
+import torch.nn as nn
 import math
 
 def normalize(latent, target_min=None, target_max=None):
@@ -129,17 +130,9 @@ def stable_slerp(a, b, t: float, eps: float = 1e-6):
 
 mscales = {
     "Default": None,
-    "Bandpass": [
-        (5, 0.0),    # Low-pass filter
-        (15, 1.0),   # Pass-through filter (allows mid-range frequencies)
-        (25, 0.0),   # High-pass filter
-    ],
     "Low-Pass": [
         (10, 1.0),   # Allows low-frequency components, suppresses high-frequency components
     ],
-    "High-Pass": [
-        (10, 0.0),   # Suppresses low-frequency components, allows high-frequency components
-    ],
     "Pass-Through": [
         (10, 1.0),   # Passes all frequencies unchanged, no filtering
     ],
@@ -228,27 +221,27 @@ class WAS_FreeU:
     def INPUT_TYPES(s):
         return {"required": {
                     "model": ("MODEL",),
-                    "target_block": (["output_block", "middle_block", "input_block", "all"],),
-                    "multiscale_mode": (list(mscales.keys()),),
-                    "multiscale_strength": ("FLOAT", {"default": 1.0, "max": 1.0, "min": 0, "step": 0.001}),
-                    "slice_b1": ("INT", {"default": 640, "min": 64, "max": 1280, "step": 1}),
-                    "slice_b2": ("INT", {"default": 320, "min": 64, "max": 640, "step": 1}),
-                    "b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.001}),
-                    "b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.001}),
-                    "s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.001}),
-                    "s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.001}),
+                    "target_block": (["output_block", "middle_block", "input_block", "all"], {"tooltip": "Which UNet block(s) to patch."}),
+                    "multiscale_mode": (list(mscales.keys()), {"tooltip": "Frequency shaping preset used by Fourier_filter."}),
+                    "multiscale_strength": ("FLOAT", {"default": 1.0, "max": 1.0, "min": 0, "step": 0.001, "tooltip": "Intensity of multi-scale shaping [0-1]."}),
+                    "slice_b1": ("INT", {"default": 640, "min": 64, "max": 1280, "step": 1, "tooltip": "Slice width (channels) affected in 1280-channel features."}),
+                    "slice_b2": ("INT", {"default": 320, "min": 64, "max": 640, "step": 1, "tooltip": "Slice width (channels) affected in 640-channel features."}),
+                    "b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Gain multiplier applied to the b1 slice (1280-ch)."}),
+                    "b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Gain multiplier applied to the b2 slice (640-ch)."}),
+                    "s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Fourier scale at threshold for 1280-ch features."}),
+                    "s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Fourier scale at threshold for 640-ch features."}),
                 },
                 "optional": {
-                    "b1_mode": (list(blending_modes.keys()),),
-                    "b1_blend": ("FLOAT", {"default": 1.0, "max": 100, "min": 0, "step": 0.001}),
-                    "b2_mode": (list(blending_modes.keys()),),
-                    "b2_blend": ("FLOAT", {"default": 1.0, "max": 100, "min": 0, "step": 0.001}),
-                    "threshold": ("INT", {"default": 1.0, "max": 10, "min": 1, "step": 1}),
-                    "use_override_scales": (["false", "true"],),
+                    "b1_mode": (list(blending_modes.keys()), {"tooltip": "Blending mode for b1 path."}),
+                    "b1_blend": ("FLOAT", {"default": 1.0, "max": 100, "min": 0, "step": 0.001, "tooltip": "Blend strength for b1 path."}),
+                    "b2_mode": (list(blending_modes.keys()), {"tooltip": "Blending mode for b2 path."}),
+                    "b2_blend": ("FLOAT", {"default": 1.0, "max": 100, "min": 0, "step": 0.001, "tooltip": "Blend strength for b2 path."}),
+                    "threshold": ("INT", {"default": 1.0, "max": 10, "min": 1, "step": 1, "tooltip": "Base radius for the Fourier mask."}),
+                    "use_override_scales": (["false", "true"], {"tooltip": "Enable manual override of scale presets."}),
                     "override_scales": ("STRING", {"default": '''# OVERRIDE SCALES
 
 # Sharpen
-# 10, 1.5''', "multiline": True}),
+# 10, 1.5''', "multiline": True, "tooltip": "Custom scale lines: '<radius>, <scale>'. Comments with #,//,!"}),
                 }
         }
 
@@ -307,6 +300,7 @@ def block_patch_hsp(h, hsp, transformer_options):
         print(f"Patching {target_block}")
 
         m = model.clone()
+
         if target_block == "all" or target_block == "output_block":
             m.set_model_output_block_patch(block_patch_hsp)
         if target_block == "all" or target_block == "input_block":
@@ -315,30 +309,97 @@ def block_patch_hsp(h, hsp, transformer_options):
             m.set_model_patch(block_patch, "middle_block_patch")
         return (m, )
 
+
+
+class WAS_PostCFGShift:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                    "model": ("MODEL",),
+                    "steps": ("INT", {"default": 20, "min": 1, "max": 1000, "step": 1, "tooltip": "Number of steps to apply SHIFT."}),
+                    "mode": (list(blending_modes.keys()), {"tooltip": "Blend strategy for denoised vs denoised*b (e.g., inject, stable_slerp)."}),
+                    "blend": ("FLOAT", {"default": 1.0, "max": 100.0, "min": 0.0, "step": 0.001, "tooltip": "Blend amount between base and scaled tensors."}),
+                    "b": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Scale factor for the injected path (higher = stronger)."}),
+                    "apply_fourier": ("BOOLEAN", {"default": False, "tooltip": "Apply frequency-domain shaping (Fourier_filter)."}),
+                    "multiscale_mode": (list(mscales.keys()), {"tooltip": "Preset shaping curves for Fourier_filter."}),
+                    "multiscale_strength": ("FLOAT", {"default": 1.0, "max": 1.0, "min": 0.0, "step": 0.001, "tooltip": "Intensity of multi-scale shaping [0-1]."}),
+                    "threshold": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1, "tooltip": "Base radius for frequency mask."}),
+                    "s": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Base scale value applied at threshold radius."}),
+                    "force_gain": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Final multiplier to boost or attenuate effect."}),
+                }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+
+    CATEGORY = "_for_testing"
+
+    def patch(self, model, steps, mode, blend, b, apply_fourier, multiscale_mode, multiscale_strength, threshold, s, force_gain):
+
+        scales = mscales[multiscale_mode]
+        steps = max(1, min(1000, steps))
+        current_step = 0
+
+        print(
+            "[FluxU] inputs:",
+            f"mode={mode}", f"blend={blend}", f"b={b}",
+            f"apply_fourier={apply_fourier}", f"multiscale_mode={multiscale_mode}", f"multiscale_strength={multiscale_strength}",
+            f"threshold={threshold}", f"s={s}", f"force_gain={force_gain}"
+            )
+
+        m = model.clone()
+
+        def post_cfg_function(args):
+
+            nonlocal current_step
+            current_step += 1
+            if current_step > steps:
+                return args.get("denoised")
+
+            denoised = args.get("denoised")
+            eff_blend = float(blend)
+            t_scaled = denoised * b
+            y = blending_modes[mode](denoised, t_scaled, eff_blend)
+
+            if apply_fourier:
+                y = Fourier_filter(y, threshold=threshold, scale=s, scales=scales, strength=multiscale_strength)
+
+            if force_gain != 1.0:
+                y = y * float(force_gain)
+            return y
+
+
+        try:
+            m.set_model_sampler_post_cfg_function(post_cfg_function)
+            print("[FluxU] set_model_sampler_post_cfg_function registered")
+        except Exception as e:
+            print(f"[FluxU] set_model_sampler_post_cfg_function failed: {e}")
+        return (m, )
+
 class WAS_FreeU_V2:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
                     "model": ("MODEL",),
-                    "input_block": ("BOOLEAN", {"default": False}),
-                    "middle_block": ("BOOLEAN", {"default": False}),
-                    "output_block": ("BOOLEAN", {"default": False}),
-                    "multiscale_mode": (list(mscales.keys()),),
-                    "multiscale_strength": ("FLOAT", {"default": 1.0, "max": 1.0, "min": 0, "step": 0.001}),
-                    "slice_b1": ("INT", {"default": 640, "min": 64, "max": 1280, "step": 1}),
-                    "slice_b2": ("INT", {"default": 320, "min": 64, "max": 640, "step": 1}),
-                    "b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.001}),
-                    "b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.001}),
-                    "s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.001}),
-                    "s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.001}),
+                    "input_block": ("BOOLEAN", {"default": False, "tooltip": "Enable patching on the input block."}),
+                    "middle_block": ("BOOLEAN", {"default": False, "tooltip": "Enable patching on the middle block."}),
+                    "output_block": ("BOOLEAN", {"default": False, "tooltip": "Enable patching on the output block."}),
+                    "multiscale_mode": (list(mscales.keys()), {"tooltip": "Frequency shaping preset used by Fourier_filter."}),
+                    "multiscale_strength": ("FLOAT", {"default": 1.0, "max": 1.0, "min": 0, "step": 0.001, "tooltip": "Intensity of multi-scale shaping [0-1]."}),
+                    "slice_b1": ("INT", {"default": 640, "min": 64, "max": 1280, "step": 1, "tooltip": "Slice width (channels) affected in 1280-channel features."}),
+                    "slice_b2": ("INT", {"default": 320, "min": 64, "max": 640, "step": 1, "tooltip": "Slice width (channels) affected in 640-channel features."}),
+                    "b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Gain multiplier for 1280-channel slice."}),
+                    "b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Gain multiplier for 640-channel slice."}),
+                    "s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Fourier scale at threshold for 1280-ch features."}),
+                    "s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Fourier scale at threshold for 640-ch features."}),
                 },
                 "optional": {
-                    "threshold": ("INT", {"default": 1.0, "max": 10, "min": 1, "step": 1}),
-                    "use_override_scales": (["false", "true"],),
+                    "threshold": ("INT", {"default": 1.0, "max": 10, "min": 1, "step": 1, "tooltip": "Base radius for the Fourier mask."}),
+                    "use_override_scales": (["false", "true"], {"tooltip": "Enable manual override of scale presets."}),
                     "override_scales": ("STRING", {"default": '''# OVERRIDE SCALES
 
 # Sharpen
-# 10, 1.5''', "multiline": True}),
+# 10, 1.5''', "multiline": True, "tooltip": "Custom scale lines: '<radius>, <scale>'. Comments with #,//,!"}),
                 }
         }
 
@@ -410,11 +471,14 @@ def block_patch_hsp(h, hsp, transformer_options):
         return (m, )
 
 NODE_CLASS_MAPPINGS = {
-    "FreeU (Advanced)": WAS_FreeU,
-    "FreeU_V2 (Advanced)": WAS_FreeU_V2,
+    "WAS_FreeU": WAS_FreeU,
+    "WAS_FreeU_V2": WAS_FreeU_V2,
+    "WAS_PostCFGShift": WAS_PostCFGShift,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "FreeU (Advanced)": "FreeU (Advanced Plus)",
-    "FreeU_V2 (Advanced)": "FreeU V2 (Advanced Plus)",
+    "WAS_FreeU": "FreeU (Advanced Plus)",
+    "WAS_FreeU_V2": "FreeU V2 (Advanced Plus)",
+    "WAS_PostCFGShift": "Post-CFG SHIFT",
 }
+
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "freeu_advanced"
 description = "This custom node provides advanced settings for FreeU."
-version = "1.1.0"
+version = "1.2.0"
 license = { file = "LICENSE" }
 
 [project.urls]
