From d3bd3bda48dcaec86340f4d0b6cf3edebd1f184b Mon Sep 17 00:00:00 2001
From: Gaeros <163143799+elias-gaeros@users.noreply.github.com>
Date: Tue, 12 Mar 2024 15:48:36 +0100
Subject: [PATCH] rescale: fast path, rename param

---
 nodes/nodes_pred.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/nodes/nodes_pred.py b/nodes/nodes_pred.py
index f1ea36e..cba5161 100644
--- a/nodes/nodes_pred.py
+++ b/nodes/nodes_pred.py
@@ -335,20 +335,20 @@ class ScaledGuidancePredictor(NoisePredictor):
                 "min": 0.0,
                 "max": 100.0,
             }),
-            "rescale": ("FLOAT", {
-                "default": 0.0,
-                "step": 0.1,
-                "min": 0.0,
-                "max": 100.0,
+            "rescale_multiplier": ("FLOAT", {
+                "default": 0.0, 
+                "min": 0.0, 
+                "max": 1.0, 
+                "step": 0.01
             }),
         }
     }
 
-    def __init__(self, guidance, baseline, guidance_scale, rescale):
+    def __init__(self, guidance, baseline, guidance_scale, rescale_multiplier):
         self.lhs = guidance
         self.rhs = baseline
         self.scale = guidance_scale
-        self.rescale = rescale
+        self.rescale_multiplier = rescale_multiplier
 
     def get_conds(self):
         return self.merge_conds(self.lhs.get_conds(), self.rhs.get_conds())
@@ -363,13 +363,16 @@ class ScaledGuidancePredictor(NoisePredictor):
         lhs = self.lhs.predict_noise(x, sigma, model, conds, model_options, seed)
         rhs = self.rhs.predict_noise(x, sigma, model, conds, model_options, seed)
 
+        if self.rescale_multiplier <= 0.0:
+            return lhs * self.scale + rhs
+
         sigma = sigma.view(sigma.shape[:1] + (1,) * (lhs.ndim - 1))
         x_scaled = x * (1. / (sigma.square() + 1.0))
         x_cfg = x_scaled - rhs - self.scale * lhs
         ro_pos = torch.std(x_scaled - (lhs + rhs), dim=(1, 2, 3), keepdim=True)
         ro_cfg = torch.std(x_cfg, dim=(1, 2, 3), keepdim=True)
         x_rescaled = x_cfg * (ro_pos / ro_cfg)
-        x_final = torch.lerp(x_cfg, x_rescaled, self.rescale)
+        x_final = torch.lerp(x_cfg, x_rescaled, self.rescale_multiplier)
 
         return x_scaled - x_final