Initial release

engine/plugins/guidance/rescaled_cfg.lua

@@ -0,0 +1,41 @@
+-- rescaled_cfg.lua: Std-matched post-processing guidance
+-- Runs APG at full scale, then rescales output to match conditional std.
+-- Routes through native APG (momentum + projection + norm thresholding),
+-- then applies std-matching post-processing.
+
+guidance = {
+    name        = "rescaled_cfg",
+    display     = "Rescaled CFG",
+    description = "Std-matched to prevent saturation",
+}
+
+function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
+    local n = Oc * T
+    local phi = (guidance_scale > 4.0) and 0.95 or 0.7
+
+    -- Run APG at full guidance scale first
+    apg(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
+
+    -- Compute std of conditional prediction and guided output
+    local sum_c, sum2_c = 0, 0
+    local sum_g, sum2_g = 0, 0
+    for i = 0, n - 1 do
+        local c = pred_cond[i]
+        local g = result[i]
+        sum_c = sum_c + c;  sum2_c = sum2_c + c * c
+        sum_g = sum_g + g;  sum2_g = sum2_g + g * g
+    end
+    local mean_c = sum_c / n
+    local mean_g = sum_g / n
+    local var_c = sum2_c / n - mean_c * mean_c
+    local var_g = sum2_g / n - mean_g * mean_g
+    local std_c = (var_c > 0) and math.sqrt(var_c) or 1e-5
+    local std_g = (var_g > 0) and math.sqrt(var_g) or 1e-5
+
+    -- Rescale to match conditional std, blend with raw APG output
+    local factor = std_c / (std_g + 1e-5)
+    for i = 0, n - 1 do
+        local rescaled = result[i] * factor
+        result[i] = phi * rescaled + (1 - phi) * result[i]
+    end
+end