Initial release

engine/plugins/guidance/cfg_mp.lua

@@ -0,0 +1,64 @@
+-- cfg_mp.lua: CFG-MP — Manifold Projection Guidance
+-- Paper: "Improving CFG of Flow Matching via Manifold Projection"
+--        Su et al., 2025 (arXiv:2601.21892)
+--
+-- After each solver step, projects the latent back onto a manifold where the
+-- prediction gap (cond - uncond) is minimised. Uses iterative fixed-point
+-- iteration of the operator G(x, t):
+--
+--   z = x - a * v_uncond(t, x)      -- push away from unconditioned manifold
+--   x = z + a * v_cond(t, z)        -- pull toward conditioned manifold
+--
+-- where a = |dt| / 2 (validated in paper Appendix C.2.1).
+--
+-- The guide() function applies standard linear CFG for the base velocity.
+-- The post_step() function performs K iterations of manifold projection using
+-- real model evaluations at the post-solver latent position.
+--
+-- Performance note: each iteration = 2 extra NFEs (one cond, one uncond).
+-- K=2 adds ~3x total compute; K=1 adds ~2x.
+
+guidance = {
+    name        = "cfg_mp",
+    display     = "CFG-MP",
+    description = "Manifold projection guidance (Su et al. 2025)",
+    params      = {
+        { key = "iterations", type = "slider", label = "Projection Iterations (K)",
+          default = 1, min = 1, max = 5, step = 1,
+          hint = "Fixed-point iterations per step. Paper recommends 2." },
+    },
+}
+
+-- Standard linear CFG for the base velocity step
+function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
+    -- Route through native APG for momentum/projection consistency
+    apg(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
+end
+
+-- Post-step manifold projection: called AFTER the solver updates xt
+-- Args:
+--   xt           : mutable FloatArray — current latent state (modified in-place)
+--   t            : float — timestep (t_next, the timestep we just stepped TO)
+--   n            : int — total elements in xt
+--   eval_cond    : function(xt_arr, t) — evaluates model with conditioning → vt_cond
+--   eval_uncond  : function(xt_arr, t) — evaluates model without conditioning → vt_uncond
+--   vt_cond      : mutable FloatArray — output buffer for conditional velocity
+--   vt_uncond    : mutable FloatArray — output buffer for unconditional velocity
+function post_step(xt, t, n, eval_cond, eval_uncond, vt_cond, vt_uncond)
+    local K = (params and params.iterations) or 2
+    local a = math.abs(dt or 0.03) / 2.0  -- dt is a global from the C++ bridge
+
+    for k = 1, K do
+        -- Step 1: z = xt - a * v_uncond(t, xt)
+        eval_uncond(xt, t)
+        for i = 0, n - 1 do
+            xt[i] = xt[i] - a * vt_uncond[i]
+        end
+
+        -- Step 2: xt = z + a * v_cond(t, z)
+        eval_cond(xt, t)
+        for i = 0, n - 1 do
+            xt[i] = xt[i] + a * vt_cond[i]
+        end
+    end
+end