commit a79c5163a1a715659d4dd0bd3310bf2639ea531a Author: Ethanfel Date: Wed Mar 4 17:10:07 2026 +0100 Initial implementation of SMC-CFG Ctrl ComfyUI node Implements the Sliding Mode Control CFG algorithm from the paper "CFG-Ctrl: A Control-Theoretic Perspective on Classifier-Free Guidance" (CVPR 2026) as a ComfyUI model patch node. Co-Authored-By: Claude Opus 4.6 diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..7f52c64 --- /dev/null +++ b/nodes.py @@ -0,0 +1,102 @@ +import torch + + +class SMCCFGCtrl: + """ + Implements SMC-CFG (Sliding Mode Control CFG) from the paper: + "CFG-Ctrl: A Control-Theoretic Perspective on Classifier-Free Guidance" (CVPR 2026) + https://github.com/hanyang-21/CFG-Ctrl + + Replaces standard linear CFG with a nonlinear sliding mode controller + that prevents instability, overshooting, and artifacts at high guidance scales. + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "smc_cfg_lambda": ("FLOAT", { + "default": 5.0, "min": 0.0, "max": 50.0, "step": 0.01, + "tooltip": "Sliding surface coefficient. Controls how much the controller weights previous error magnitude vs error derivative. Paper recommended: 5.0", + }), + "smc_cfg_K": ("FLOAT", { + "default": 0.2, "min": 0.0, "max": 5.0, "step": 0.01, + "tooltip": "Switching gain. Bounds the correction to [-K, +K] per element. Higher = stronger correction but may introduce chattering. Paper recommended: 0.2", + }), + "warmup_steps": ("INT", { + "default": 0, "min": 0, "max": 100, + "tooltip": "Number of initial steps with no guidance (pure conditional prediction). Lets the model establish structure before guidance kicks in.", + }), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "sampling/custom_sampling" + + def patch(self, model, smc_cfg_lambda, smc_cfg_K, warmup_steps): + # Mutable state persisted across denoising steps via closure + state = { + "prev_eps": None, + "step": 0, + "prev_sigma": None, + } + + lam = smc_cfg_lambda + K = smc_cfg_K + + def smc_cfg_function(args): + cond = args["cond"] + uncond = args["uncond"] + cond_scale = args["cond_scale"] + sigma = args["sigma"] + + # Detect new generation: sigma should decrease monotonically during + # denoising. If it jumps up, a new sampling run has started. + curr_sigma = sigma.max().item() if torch.is_tensor(sigma) else float(sigma) + if state["prev_sigma"] is not None and curr_sigma > state["prev_sigma"] * 1.1: + state["prev_eps"] = None + state["step"] = 0 + state["prev_sigma"] = curr_sigma + + step = state["step"] + state["step"] = step + 1 + + # Warmup: pure conditional prediction (no guidance) + if warmup_steps > 0 and step < warmup_steps: + return cond + + # Guidance error: e_t = noise_cond - noise_uncond + guidance_eps = cond - uncond + + if state["prev_eps"] is not None: + prev_eps = state["prev_eps"] + + # Sliding surface: s_t = (e_t - e_{t-1}) + lambda * e_{t-1} + s = (guidance_eps - prev_eps) + lam * prev_eps + + # Switching control: u_sw = -K * sign(s_t) + u_sw = -K * torch.sign(s) + + # Apply correction to guidance error + guidance_eps = guidance_eps + u_sw + + # Store corrected guidance for next step's sliding surface + state["prev_eps"] = guidance_eps.detach().clone() + + # v_guided = v_uncond + scale * corrected_guidance + return uncond + cond_scale * guidance_eps + + m = model.clone() + m.set_model_sampler_cfg_function(smc_cfg_function, disable_cfg1_optimization=True) + return (m,) + + +NODE_CLASS_MAPPINGS = { + "SMCCFGCtrl": SMCCFGCtrl, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SMCCFGCtrl": "SMC-CFG Ctrl", +}