85 lines
2.9 KiB
Lua
85 lines
2.9 KiB
Lua
-- smc_cfg.lua: SMC-CFG — Sliding Mode Control Guidance
|
|
-- Paper: "CFG-Ctrl: Control-Based Classifier-Free Diffusion Guidance"
|
|
-- Han et al., 2025 (arXiv:2603.03281)
|
|
--
|
|
-- Reinterprets CFG as a feedback control system and applies Sliding Mode
|
|
-- Control (SMC) to stabilise guidance, especially at high scales.
|
|
--
|
|
-- Key idea: define a sliding surface s(t) = ė(t) + λ·e(t) over the
|
|
-- semantic error e = v_cond - v_uncond, then apply a switching control
|
|
-- term Δe = -k·sign(s) that enforces convergence to a stable manifold.
|
|
--
|
|
-- Implementation: routes through native APG for stability (momentum,
|
|
-- projection, norm thresholding), then applies the SMC correction as
|
|
-- a delta on top: result = APG(cond, uncond, w) + w · Δe
|
|
--
|
|
-- Stateful: stores previous error vector across steps.
|
|
|
|
guidance = {
|
|
name = "smc_cfg",
|
|
display = "SMC-CFG",
|
|
description = "Sliding mode control guidance (Han et al. 2025)",
|
|
params = {
|
|
{ key = "lambda", type = "slider", label = "λ (Surface Slope)",
|
|
default = 0.5, min = 0.01, max = 2.0, step = 0.01,
|
|
hint = "Controls the sliding surface shape. Higher = faster convergence" },
|
|
{ key = "k", type = "slider", label = "k (Switching Gain)",
|
|
default = 0.1, min = 0.01, max = 1.0, step = 0.01,
|
|
hint = "Force toward the sliding surface. Too high = vibrations" },
|
|
},
|
|
}
|
|
|
|
-- Stateful: previous error buffer
|
|
local prev_error = nil
|
|
local prev_n = 0
|
|
|
|
local function sign(x)
|
|
if x > 0 then return 1.0
|
|
elseif x < 0 then return -1.0
|
|
else return 0.0
|
|
end
|
|
end
|
|
|
|
function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
|
local n = Oc * T
|
|
local lam = (params and params.lambda) or 0.5
|
|
local k = (params and params.k) or 0.1
|
|
|
|
-- Reset state on first step of a new generation
|
|
if (step_idx or 0) == 0 then prev_error = nil; prev_n = 0 end
|
|
|
|
-- Base guidance through APG (handles momentum, projection, norm thresholding)
|
|
apg(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
|
|
|
-- Compute semantic error e(t) = cond - uncond
|
|
local error_now = {}
|
|
for i = 0, n - 1 do
|
|
error_now[i] = pred_cond[i] - pred_uncond[i]
|
|
end
|
|
|
|
-- First step or size change: no previous error, just use APG as-is
|
|
if prev_error == nil or prev_n ~= n then
|
|
prev_error = error_now
|
|
prev_n = n
|
|
return
|
|
end
|
|
|
|
-- Compute ė ≈ (e_now - e_prev) / dt
|
|
local dt_abs = math.abs(dt or 1.0)
|
|
if dt_abs < 1e-8 then dt_abs = 1e-8 end
|
|
local inv_dt = 1.0 / dt_abs
|
|
|
|
-- Apply SMC correction: Δe = -k · sign(ė + λ·e)
|
|
-- Add w · Δe as delta on top of APG result
|
|
for i = 0, n - 1 do
|
|
local e_dot = (error_now[i] - prev_error[i]) * inv_dt
|
|
local s = e_dot + lam * error_now[i]
|
|
local delta_e = -k * sign(s)
|
|
result[i] = result[i] + guidance_scale * delta_e
|
|
end
|
|
|
|
-- Store for next step
|
|
prev_error = error_now
|
|
prev_n = n
|
|
end
|