From 5f0bd6825d5812194bea3987c88f43d84cd1baf8 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Wed, 4 Mar 2026 17:43:07 +0100 Subject: [PATCH] Fix: normalize to noise-prediction space before SMC MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous fix (denoised space) still had the problem: K * cond_scale produced a constant ±2.4 perturbation per element at cfg=12, destroying the image at every step. The paper's K=0.2 is calibrated for unit-variance noise predictions. ComfyUI's cond/uncond are sigma-scaled (x - denoised ≈ sigma * epsilon). Now we divide by sigma to recover epsilon-space, apply SMC there, then multiply back by sigma. This gives natural dampening at late steps: - sigma=14 (early): correction ±33 in latent space (image is noise anyway) - sigma=0.01 (late): correction ±0.024 in latent space (negligible) This matches the paper's behavior where the scheduler conversion inherently dampens the noise-space correction at low sigma values. Co-Authored-By: Claude Opus 4.6 --- nodes.py | 30 ++++++++++++++---------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/nodes.py b/nodes.py index bbe62e6..df849d5 100644 --- a/nodes.py +++ b/nodes.py @@ -47,14 +47,9 @@ class SMCCFGCtrl: K = smc_cfg_K def smc_cfg_function(args): - # Use denoised-space predictions — these have consistent magnitude - # across sigma values. ComfyUI's args["cond"]/["uncond"] are - # (x - denoised), which are sigma-scaled and would make the fixed - # K correction dominate at low sigma (late steps), destroying the image. - cond_denoised = args["cond_denoised"] - uncond_denoised = args["uncond_denoised"] + cond = args["cond"] # x - cond_denoised (sigma-scaled noise) + uncond = args["uncond"] # x - uncond_denoised (sigma-scaled noise) cond_scale = args["cond_scale"] - x = args["input"] sigma = args["sigma"] # Detect new generation: sigma should decrease monotonically during @@ -70,10 +65,16 @@ class SMCCFGCtrl: # Warmup: pure conditional prediction (no guidance) if warmup_steps > 0 and step < warmup_steps: - return x - cond_denoised + return cond - # Guidance error in denoised space (consistent magnitude across sigma) - guidance_eps = cond_denoised - uncond_denoised + # Normalize to noise-prediction space by dividing out sigma. + # The paper's K is calibrated for unit-variance noise predictions. + # ComfyUI's cond/uncond are (x - denoised) ≈ sigma * epsilon, + # so dividing by sigma recovers epsilon-space where K=0.2 is correct. + # Crucially, when converting back, the sigma factor naturally dampens + # the correction at late steps (small sigma), preventing noise injection. + sigma_val = max(curr_sigma, 1e-8) + guidance_eps = (cond - uncond) / sigma_val # Initialize prev_eps on first SMC step (matches original paper # where SMC correction is applied from the very first step) @@ -88,17 +89,14 @@ class SMCCFGCtrl: # Switching control: u_sw = -K * sign(s_t) u_sw = -K * torch.sign(s) - # Corrected guidance error + # Corrected guidance error (in normalized noise space) guidance_eps = guidance_eps + u_sw # Store corrected guidance for next step's sliding surface state["prev_eps"] = guidance_eps.detach().clone() - # Guided denoised output - denoised = uncond_denoised + cond_scale * guidance_eps - - # Return noise residual (framework computes cfg_result = x - return) - return x - denoised + # Convert back to sigma-scaled space and apply CFG + return uncond + cond_scale * guidance_eps * sigma_val m = model.clone() m.set_model_sampler_cfg_function(smc_cfg_function, disable_cfg1_optimization=True)