From 9a616b81c15cec7f5ddcbc12e349f1adc03fad67 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 4 Dec 2024 19:25:50 -0500 Subject: [PATCH] Add rescaling_scale from STG to SkipLayerGuidanceDiT. --- comfy_extras/nodes_slg.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_slg.py b/comfy_extras/nodes_slg.py index 8a1181fcfbd..2fa09e2505e 100644 --- a/comfy_extras/nodes_slg.py +++ b/comfy_extras/nodes_slg.py @@ -16,7 +16,8 @@ def INPUT_TYPES(s): "single_layers": ("STRING", {"default": "7, 8, 9", "multiline": False}), "scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 10.0, "step": 0.1}), "start_percent": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 1.0, "step": 0.001}), - "end_percent": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.001}) + "end_percent": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.001}), + "rescaling_scale": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01}), }} RETURN_TYPES = ("MODEL",) FUNCTION = "skip_guidance" @@ -26,7 +27,7 @@ def INPUT_TYPES(s): CATEGORY = "advanced/guidance" - def skip_guidance(self, model, scale, start_percent, end_percent, double_layers="", single_layers=""): + def skip_guidance(self, model, scale, start_percent, end_percent, double_layers="", single_layers="", rescaling_scale=0): # check if layer is comma separated integers def skip(args, extra_args): return args @@ -65,6 +66,11 @@ def post_cfg_function(args): if scale > 0 and sigma_ >= sigma_end and sigma_ <= sigma_start: (slg,) = comfy.samplers.calc_cond_batch(model, [cond], x, sigma, model_options) cfg_result = cfg_result + (cond_pred - slg) * scale + if rescaling_scale != 0: + factor = cond_pred.std() / cfg_result.std() + factor = rescaling_scale * factor + (1 - rescaling_scale) + cfg_result *= factor + return cfg_result m = model.clone()