Initial release
This commit is contained in:
civ
@@ -0,0 +1,24 @@
|
||||
-- apg.lua: Adaptive Perpendicular Guidance (default)
|
||||
-- Routes through the native C++ apg_forward() for momentum smoothing,
|
||||
-- per-channel norm thresholding, and perpendicular projection.
|
||||
-- In practice, the engine takes the native C++ path for APG directly,
|
||||
-- but this Lua implementation exists as a correct fallback.
|
||||
|
||||
guidance = {
|
||||
name = "apg",
|
||||
display = "APG",
|
||||
description = "Adaptive perpendicular guidance (default)",
|
||||
params = {
|
||||
{ key = "momentum", type = "slider", label = "Momentum",
|
||||
default = 0.75, min = -1, max = 1, step = 0.01,
|
||||
hint = "APG momentum coefficient (negative = adaptive)" },
|
||||
{ key = "norm_threshold", type = "slider", label = "Norm Threshold",
|
||||
default = 2.5, min = 0, max = 10, step = 0.1,
|
||||
hint = "APG norm clipping threshold" },
|
||||
},
|
||||
}
|
||||
|
||||
function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
||||
-- Route through native APG C++ implementation
|
||||
apg(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
||||
end
|
||||
@@ -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
|
||||
@@ -0,0 +1,21 @@
|
||||
-- cfg_pp.lua: CFG++ — Step-scaled guidance for large steps
|
||||
-- Reduces effective scale by the step-size-to-sigma ratio.
|
||||
-- Routes through native APG (momentum + projection + norm thresholding).
|
||||
|
||||
guidance = {
|
||||
name = "cfg_pp",
|
||||
display = "CFG++",
|
||||
description = "Step-scaled guidance for few-step models",
|
||||
}
|
||||
|
||||
function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
||||
-- Effective scale adjusted by step ratio (t_curr and dt are globals)
|
||||
local effective_scale = guidance_scale
|
||||
if t_curr and t_curr > 1e-6 and dt then
|
||||
local step_scale = math.abs(dt) / t_curr
|
||||
effective_scale = 1.0 + (guidance_scale - 1.0) * step_scale
|
||||
end
|
||||
|
||||
-- Route through native APG with the adjusted scale
|
||||
apg(pred_cond, pred_uncond, effective_scale, result, Oc, T, norm_threshold)
|
||||
end
|
||||
@@ -0,0 +1,39 @@
|
||||
-- cfg_zero_star.lua: CFG-Zero⋆ — Zero-Init Guidance
|
||||
-- Paper: "CFG-Zero⋆: Improved Classifier-Free Guidance for Flow Matching Models"
|
||||
-- Fan et al., 2025 (arXiv:2503.18886)
|
||||
--
|
||||
-- The paper proposes two improvements: optimised scale (s⋆) and zero-init.
|
||||
-- Since our engine uses APG (perpendicular projection + momentum), which
|
||||
-- already corrects for the underfitting that s⋆ addresses, only zero-init
|
||||
-- is applied here. Combining both would double-correct.
|
||||
--
|
||||
-- Zero-init: zeroes out velocity for the first N ODE steps, since early-step
|
||||
-- CFG predictions in flow matching are often worse than doing nothing.
|
||||
-- All subsequent steps use the standard APG pipeline.
|
||||
|
||||
guidance = {
|
||||
name = "cfg_zero_star",
|
||||
display = "CFG-Zero⋆",
|
||||
description = "Zero-init + APG guidance (Fan et al. 2025)",
|
||||
params = {
|
||||
{ key = "zero_init_steps", type = "slider", label = "Zero-Init Steps",
|
||||
default = 1, min = 0, max = 5, step = 1,
|
||||
hint = "Number of initial ODE steps to zero out (paper recommends 1)" },
|
||||
},
|
||||
}
|
||||
|
||||
function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
||||
local n = Oc * T
|
||||
local zero_init_steps = (params and params.zero_init_steps) or 1
|
||||
|
||||
-- Zero-init: zero out velocity for the first N steps
|
||||
if (step_idx or 0) < zero_init_steps then
|
||||
for i = 0, n - 1 do
|
||||
result[i] = 0.0
|
||||
end
|
||||
return
|
||||
end
|
||||
|
||||
-- Standard APG for all other steps
|
||||
apg(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
||||
end
|
||||
@@ -0,0 +1,20 @@
|
||||
-- dynamic_cfg.lua: Cosine-decaying guidance schedule
|
||||
-- Full guidance early (structure), reduced guidance late (fine detail).
|
||||
-- Routes through native APG (momentum + projection + norm thresholding).
|
||||
|
||||
guidance = {
|
||||
name = "dynamic_cfg",
|
||||
display = "Dynamic CFG",
|
||||
description = "Cosine-decaying guidance schedule with APG projection",
|
||||
}
|
||||
|
||||
function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
||||
local power = 0.5
|
||||
local progress = (step_idx or 0) / math.max((total_steps or 1) - 1, 1)
|
||||
local cos_val = math.max(math.cos(math.pi / 2 * progress), 0)
|
||||
local decay = cos_val ^ power
|
||||
local effective_scale = 1.0 + (guidance_scale - 1.0) * decay
|
||||
|
||||
-- Route through native APG with the decayed scale
|
||||
apg(pred_cond, pred_uncond, effective_scale, result, Oc, T, norm_threshold)
|
||||
end
|
||||
@@ -0,0 +1,41 @@
|
||||
-- rescaled_cfg.lua: Std-matched post-processing guidance
|
||||
-- Runs APG at full scale, then rescales output to match conditional std.
|
||||
-- Routes through native APG (momentum + projection + norm thresholding),
|
||||
-- then applies std-matching post-processing.
|
||||
|
||||
guidance = {
|
||||
name = "rescaled_cfg",
|
||||
display = "Rescaled CFG",
|
||||
description = "Std-matched to prevent saturation",
|
||||
}
|
||||
|
||||
function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
||||
local n = Oc * T
|
||||
local phi = (guidance_scale > 4.0) and 0.95 or 0.7
|
||||
|
||||
-- Run APG at full guidance scale first
|
||||
apg(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
||||
|
||||
-- Compute std of conditional prediction and guided output
|
||||
local sum_c, sum2_c = 0, 0
|
||||
local sum_g, sum2_g = 0, 0
|
||||
for i = 0, n - 1 do
|
||||
local c = pred_cond[i]
|
||||
local g = result[i]
|
||||
sum_c = sum_c + c; sum2_c = sum2_c + c * c
|
||||
sum_g = sum_g + g; sum2_g = sum2_g + g * g
|
||||
end
|
||||
local mean_c = sum_c / n
|
||||
local mean_g = sum_g / n
|
||||
local var_c = sum2_c / n - mean_c * mean_c
|
||||
local var_g = sum2_g / n - mean_g * mean_g
|
||||
local std_c = (var_c > 0) and math.sqrt(var_c) or 1e-5
|
||||
local std_g = (var_g > 0) and math.sqrt(var_g) or 1e-5
|
||||
|
||||
-- Rescale to match conditional std, blend with raw APG output
|
||||
local factor = std_c / (std_g + 1e-5)
|
||||
for i = 0, n - 1 do
|
||||
local rescaled = result[i] * factor
|
||||
result[i] = phi * rescaled + (1 - phi) * result[i]
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,84 @@
|
||||
-- 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
|
||||
@@ -0,0 +1,42 @@
|
||||
-- beta57.lua: Beta(0.5, 0.7) distribution scheduler
|
||||
-- Requires beta_math companion for the inverse CDF computation.
|
||||
|
||||
local beta_math = require("beta_math")
|
||||
|
||||
scheduler = {
|
||||
name = "beta57",
|
||||
display = "Beta 57",
|
||||
description = "Beta(0.5,0.7) — smooth S-curve from RES4LYF",
|
||||
}
|
||||
|
||||
function schedule(output, num_steps, shift)
|
||||
local alpha = 0.5
|
||||
local beta = 0.7
|
||||
|
||||
for i = 0, num_steps - 1 do
|
||||
local u = (i + 0.5) / num_steps
|
||||
local t = 1.0 - beta_math.ppf(u, alpha, beta)
|
||||
output[i] = t
|
||||
end
|
||||
-- Sort descending
|
||||
local vals = {}
|
||||
for i = 0, num_steps - 1 do vals[i+1] = output[i] end
|
||||
table.sort(vals, function(a,b) return a > b end)
|
||||
for i = 0, num_steps - 1 do output[i] = vals[i+1] end
|
||||
clamp(output, num_steps)
|
||||
apply_shift(output, num_steps, shift)
|
||||
end
|
||||
|
||||
function apply_shift(ts, n, shift)
|
||||
if shift == 1.0 then return end
|
||||
for i = 0, n - 1 do
|
||||
local t = ts[i]; ts[i] = shift * t / (1.0 + (shift - 1.0) * t)
|
||||
end
|
||||
end
|
||||
|
||||
function clamp(ts, n)
|
||||
for i = 0, n - 1 do
|
||||
if ts[i] < 1e-6 then ts[i] = 1e-6 end
|
||||
if ts[i] > 1.0 then ts[i] = 1.0 end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,114 @@
|
||||
-- beta_math.lua: Beta distribution math helpers (companion data file)
|
||||
-- Provides regularized incomplete beta function and its inverse (ppf).
|
||||
-- Ported from engine/src/schedulers/scheduler-implementations.h
|
||||
|
||||
local M = {}
|
||||
|
||||
-- Log-gamma (uses Lua's built-in math library)
|
||||
local function lgamma(x)
|
||||
-- Lanczos approximation for log-gamma
|
||||
if x <= 0 then return 0 end
|
||||
local g = 7
|
||||
local c = {
|
||||
0.99999999999980993,
|
||||
676.5203681218851,
|
||||
-1259.1392167224028,
|
||||
771.32342877765313,
|
||||
-176.61502916214059,
|
||||
12.507343278686905,
|
||||
-0.13857109526572012,
|
||||
9.9843695780195716e-6,
|
||||
1.5056327351493116e-7
|
||||
}
|
||||
if x < 0.5 then
|
||||
return math.log(math.pi / math.sin(math.pi * x)) - lgamma(1 - x)
|
||||
end
|
||||
x = x - 1
|
||||
local a = c[1]
|
||||
local t = x + g + 0.5
|
||||
for i = 2, #c do
|
||||
a = a + c[i] / (x + i - 1)
|
||||
end
|
||||
return 0.5 * math.log(2 * math.pi) + (x + 0.5) * math.log(t) - t + math.log(a)
|
||||
end
|
||||
|
||||
-- Log of beta function: B(a,b) = Gamma(a)*Gamma(b)/Gamma(a+b)
|
||||
local function lbeta(a, b)
|
||||
return lgamma(a) + lgamma(b) - lgamma(a + b)
|
||||
end
|
||||
|
||||
-- Regularized incomplete beta function via continued fraction (Lentz's method)
|
||||
local function betainc(a, b, x)
|
||||
if x <= 0 then return 0 end
|
||||
if x >= 1 then return 1 end
|
||||
|
||||
-- Use symmetry for convergence
|
||||
if x > (a + 1) / (a + b + 2) then
|
||||
return 1 - betainc(b, a, 1 - x)
|
||||
end
|
||||
|
||||
local ln_pre = a * math.log(x) + b * math.log(1 - x) - lbeta(a, b)
|
||||
local qab = a + b
|
||||
local qap = a + 1
|
||||
local qam = a - 1
|
||||
local c = 1
|
||||
local d = 1 - qab * x / qap
|
||||
if math.abs(d) < 1e-30 then d = 1e-30 end
|
||||
d = 1 / d
|
||||
local h = d
|
||||
|
||||
for m = 1, 200 do
|
||||
local m2 = 2 * m
|
||||
-- Even numerator
|
||||
local aa = m * (b - m) * x / ((qam + m2) * (a + m2))
|
||||
d = 1 + aa * d; if math.abs(d) < 1e-30 then d = 1e-30 end
|
||||
c = 1 + aa / c; if math.abs(c) < 1e-30 then c = 1e-30 end
|
||||
d = 1 / d; h = h * d * c
|
||||
|
||||
-- Odd numerator
|
||||
aa = -((a + m) * (qab + m) * x) / ((a + m2) * (qap + m2))
|
||||
d = 1 + aa * d; if math.abs(d) < 1e-30 then d = 1e-30 end
|
||||
c = 1 + aa / c; if math.abs(c) < 1e-30 then c = 1e-30 end
|
||||
d = 1 / d
|
||||
local del = d * c; h = h * del
|
||||
|
||||
if math.abs(del - 1) < 3e-14 then break end
|
||||
end
|
||||
|
||||
return math.exp(ln_pre) * h / a
|
||||
end
|
||||
|
||||
-- Beta PDF
|
||||
local function beta_pdf(x, a, b)
|
||||
if x <= 0 or x >= 1 then return 0 end
|
||||
return math.exp((a - 1) * math.log(x) + (b - 1) * math.log(1 - x) - lbeta(a, b))
|
||||
end
|
||||
|
||||
-- Inverse CDF (ppf) via Newton's method
|
||||
function M.ppf(p, a, b)
|
||||
if p <= 0 then return 0 end
|
||||
if p >= 1 then return 1 end
|
||||
|
||||
-- Initial guess
|
||||
local mu = a / (a + b)
|
||||
local var = a * b / ((a + b)^2 * (a + b + 1))
|
||||
local sigma = math.sqrt(var)
|
||||
local x = mu + sigma * (2 * p - 1)
|
||||
if x < 0.001 then x = 0.001 end
|
||||
if x > 0.999 then x = 0.999 end
|
||||
|
||||
-- Newton-Raphson
|
||||
for _ = 1, 50 do
|
||||
local F = betainc(a, b, x) - p
|
||||
local f = beta_pdf(x, a, b)
|
||||
if math.abs(f) < 1e-30 then break end
|
||||
local dx = -F / f
|
||||
x = x + dx
|
||||
if x < 1e-10 then x = 1e-10 end
|
||||
if x > 1 - 1e-10 then x = 1 - 1e-10 end
|
||||
if math.abs(dx) < 1e-12 then break end
|
||||
end
|
||||
return x
|
||||
end
|
||||
|
||||
return M
|
||||
@@ -0,0 +1,38 @@
|
||||
-- bong_tangent.lua: Tangent-based scheduler, concentrates at high noise
|
||||
|
||||
scheduler = {
|
||||
name = "bong_tangent",
|
||||
display = "Tangent",
|
||||
description = "Front-loaded (structural focus)",
|
||||
}
|
||||
|
||||
function schedule(output, num_steps, shift)
|
||||
local scale = 1.5
|
||||
for i = 0, num_steps - 1 do
|
||||
local frac = (i + 0.5) / num_steps
|
||||
local angle = frac * math.pi / 2.0
|
||||
local tan_val = math.tan(angle)
|
||||
output[i] = 1.0 - (2.0 / math.pi) * math.atan(tan_val * scale)
|
||||
end
|
||||
-- Sort descending
|
||||
local vals = {}
|
||||
for i = 0, num_steps - 1 do vals[i+1] = output[i] end
|
||||
table.sort(vals, function(a,b) return a > b end)
|
||||
for i = 0, num_steps - 1 do output[i] = vals[i+1] end
|
||||
clamp(output, num_steps)
|
||||
apply_shift(output, num_steps, shift)
|
||||
end
|
||||
|
||||
function apply_shift(ts, n, shift)
|
||||
if shift == 1.0 then return end
|
||||
for i = 0, n - 1 do
|
||||
local t = ts[i]; ts[i] = shift * t / (1.0 + (shift - 1.0) * t)
|
||||
end
|
||||
end
|
||||
|
||||
function clamp(ts, n)
|
||||
for i = 0, n - 1 do
|
||||
if ts[i] < 1e-6 then ts[i] = 1e-6 end
|
||||
if ts[i] > 1.0 then ts[i] = 1.0 end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,31 @@
|
||||
-- cosine.lua: Cosine scheduler — half-cosine S-curve
|
||||
|
||||
scheduler = {
|
||||
name = "cosine",
|
||||
display = "Cosine",
|
||||
description = "Cosine annealing — balanced S-curve",
|
||||
}
|
||||
|
||||
function schedule(output, num_steps, shift)
|
||||
for i = 0, num_steps - 1 do
|
||||
local frac = i / num_steps
|
||||
output[i] = 0.5 * (1.0 + math.cos(math.pi * frac))
|
||||
end
|
||||
clamp(output, num_steps)
|
||||
apply_shift(output, num_steps, shift)
|
||||
end
|
||||
|
||||
function apply_shift(ts, n, shift)
|
||||
if shift == 1.0 then return end
|
||||
for i = 0, n - 1 do
|
||||
local t = ts[i]
|
||||
ts[i] = shift * t / (1.0 + (shift - 1.0) * t)
|
||||
end
|
||||
end
|
||||
|
||||
function clamp(ts, n)
|
||||
for i = 0, n - 1 do
|
||||
if ts[i] < 1e-6 then ts[i] = 1e-6 end
|
||||
if ts[i] > 1.0 then ts[i] = 1.0 end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,37 @@
|
||||
-- ddim_uniform.lua: DDIM Uniform — log-SNR uniform (S-shaped)
|
||||
|
||||
scheduler = {
|
||||
name = "ddim_uniform",
|
||||
display = "DDIM Uniform",
|
||||
description = "Log-SNR uniform (S-shaped)",
|
||||
}
|
||||
|
||||
function schedule(output, num_steps, shift)
|
||||
local t_max = 0.9986
|
||||
local t_min = 0.0014
|
||||
local logit_max = math.log(t_max / (1 - t_max))
|
||||
local logit_min = math.log(t_min / (1 - t_min))
|
||||
|
||||
for i = 0, num_steps - 1 do
|
||||
local frac = i / num_steps
|
||||
local logit_t = logit_max + (logit_min - logit_max) * frac
|
||||
output[i] = 1.0 / (1.0 + math.exp(-logit_t))
|
||||
end
|
||||
clamp(output, num_steps)
|
||||
apply_shift(output, num_steps, shift)
|
||||
end
|
||||
|
||||
function apply_shift(ts, n, shift)
|
||||
if shift == 1.0 then return end
|
||||
for i = 0, n - 1 do
|
||||
local t = ts[i]
|
||||
ts[i] = shift * t / (1.0 + (shift - 1.0) * t)
|
||||
end
|
||||
end
|
||||
|
||||
function clamp(ts, n)
|
||||
for i = 0, n - 1 do
|
||||
if ts[i] < 1e-6 then ts[i] = 1e-6 end
|
||||
if ts[i] > 1.0 then ts[i] = 1.0 end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,31 @@
|
||||
-- linear.lua: Linear (uniform) scheduler — the ACE-Step default
|
||||
|
||||
scheduler = {
|
||||
name = "linear",
|
||||
display = "Linear",
|
||||
description = "Uniform spacing (default)",
|
||||
}
|
||||
|
||||
function schedule(output, num_steps, shift)
|
||||
for i = 0, num_steps - 1 do
|
||||
output[i] = 1.0 - i / num_steps
|
||||
end
|
||||
apply_shift(output, num_steps, shift)
|
||||
end
|
||||
|
||||
-- Standard shift warp: t' = shift*t / (1 + (shift-1)*t)
|
||||
function apply_shift(ts, n, shift)
|
||||
if shift == 1.0 then return end
|
||||
for i = 0, n - 1 do
|
||||
local t = ts[i]
|
||||
ts[i] = shift * t / (1.0 + (shift - 1.0) * t)
|
||||
end
|
||||
end
|
||||
|
||||
-- Clamp to [1e-6, 1.0]
|
||||
function clamp(ts, n)
|
||||
for i = 0, n - 1 do
|
||||
if ts[i] < 1e-6 then ts[i] = 1e-6 end
|
||||
if ts[i] > 1.0 then ts[i] = 1.0 end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,38 @@
|
||||
-- linear_quadratic.lua: Linear start, quadratic finish
|
||||
|
||||
scheduler = {
|
||||
name = "linear_quadratic",
|
||||
display = "Linear-Quadratic",
|
||||
description = "Linear start, quadratic finish",
|
||||
}
|
||||
|
||||
function schedule(output, num_steps, shift)
|
||||
local crossover = 0.5
|
||||
local n_linear = math.max(math.floor(num_steps * crossover), 1)
|
||||
local n_quad = num_steps - n_linear
|
||||
local t_cross = 1.0 - crossover
|
||||
|
||||
for i = 0, n_linear - 1 do
|
||||
output[i] = 1.0 - i * crossover / n_linear
|
||||
end
|
||||
for i = 0, n_quad - 1 do
|
||||
local frac = (i + 1) / n_quad
|
||||
output[n_linear + i] = t_cross * (1.0 - frac * frac)
|
||||
end
|
||||
clamp(output, num_steps)
|
||||
apply_shift(output, num_steps, shift)
|
||||
end
|
||||
|
||||
function apply_shift(ts, n, shift)
|
||||
if shift == 1.0 then return end
|
||||
for i = 0, n - 1 do
|
||||
local t = ts[i]; ts[i] = shift * t / (1.0 + (shift - 1.0) * t)
|
||||
end
|
||||
end
|
||||
|
||||
function clamp(ts, n)
|
||||
for i = 0, n - 1 do
|
||||
if ts[i] < 1e-6 then ts[i] = 1e-6 end
|
||||
if ts[i] > 1.0 then ts[i] = 1.0 end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,36 @@
|
||||
-- power.lua: Power-law scheduler with configurable exponent
|
||||
|
||||
scheduler = {
|
||||
name = "power",
|
||||
display = "Power (p=2)",
|
||||
description = "Power-law t^p, front-loaded",
|
||||
params = {
|
||||
{ key = "exponent", type = "slider", label = "Exponent",
|
||||
default = 2.0, min = 0.5, max = 5.0, step = 0.1,
|
||||
hint = "Higher values front-load more steps at high noise" },
|
||||
},
|
||||
}
|
||||
|
||||
function schedule(output, num_steps, shift)
|
||||
local p = (params and params.exponent) or 2.0
|
||||
for i = 0, num_steps - 1 do
|
||||
local frac = i / num_steps
|
||||
output[i] = (1.0 - frac) ^ p
|
||||
end
|
||||
clamp(output, num_steps)
|
||||
apply_shift(output, num_steps, shift)
|
||||
end
|
||||
|
||||
function apply_shift(ts, n, shift)
|
||||
if shift == 1.0 then return end
|
||||
for i = 0, n - 1 do
|
||||
local t = ts[i]; ts[i] = shift * t / (1.0 + (shift - 1.0) * t)
|
||||
end
|
||||
end
|
||||
|
||||
function clamp(ts, n)
|
||||
for i = 0, n - 1 do
|
||||
if ts[i] < 1e-6 then ts[i] = 1e-6 end
|
||||
if ts[i] > 1.0 then ts[i] = 1.0 end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,41 @@
|
||||
-- sgm_uniform.lua: SGM Uniform (Karras) — uniform in σ^(1/ρ) space
|
||||
|
||||
scheduler = {
|
||||
name = "sgm_uniform",
|
||||
display = "SGM-Uniform (Karras)",
|
||||
description = "Karras σ-ramp (ρ=7), front-loads structural steps",
|
||||
}
|
||||
|
||||
function schedule(output, num_steps, shift)
|
||||
local t_max = 0.999
|
||||
local t_min = 0.001
|
||||
local sigma_max = t_max / (1 - t_max)
|
||||
local sigma_min = t_min / (1 - t_min)
|
||||
local rho = 7.0
|
||||
|
||||
local inv_rho = 1.0 / rho
|
||||
local s_max = sigma_max ^ inv_rho
|
||||
local s_min = sigma_min ^ inv_rho
|
||||
|
||||
for i = 0, num_steps - 1 do
|
||||
local frac = i / num_steps
|
||||
local sigma = (s_max + frac * (s_min - s_max)) ^ rho
|
||||
output[i] = sigma / (1 + sigma)
|
||||
end
|
||||
clamp(output, num_steps)
|
||||
apply_shift(output, num_steps, shift)
|
||||
end
|
||||
|
||||
function apply_shift(ts, n, shift)
|
||||
if shift == 1.0 then return end
|
||||
for i = 0, n - 1 do
|
||||
local t = ts[i]; ts[i] = shift * t / (1.0 + (shift - 1.0) * t)
|
||||
end
|
||||
end
|
||||
|
||||
function clamp(ts, n)
|
||||
for i = 0, n - 1 do
|
||||
if ts[i] < 1e-6 then ts[i] = 1e-6 end
|
||||
if ts[i] > 1.0 then ts[i] = 1.0 end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,57 @@
|
||||
-- aflops.lua: A-FloPS (1 NFE, stateful multistep)
|
||||
-- Adaptive Flow Path Sampler with velocity decomposition.
|
||||
|
||||
solver = {
|
||||
name = "aflops",
|
||||
display = "A-FloPS (1 NFE)",
|
||||
description = "Exponential integrator with residual tracking",
|
||||
nfe = 1,
|
||||
order = 2,
|
||||
needs_model = false,
|
||||
stateful = true,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
local prev_w = nil
|
||||
local prev_t = 0
|
||||
local prev_t_dst = 0
|
||||
|
||||
local function clamp_alpha(t)
|
||||
local a = 1 - t
|
||||
return math.max(1e-6, math.min(a, 1 - 1e-6))
|
||||
end
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
if (step_index or 0) == 0 then prev_w = nil; prev_t = 0; prev_t_dst = 0 end
|
||||
local alpha_curr = clamp_alpha(t_curr)
|
||||
local alpha_prev = clamp_alpha(t_prev)
|
||||
local alpha_ratio = alpha_prev / alpha_curr
|
||||
local log_ratio = math.log(alpha_ratio)
|
||||
|
||||
-- Compute residual velocity: w = v + x/(1-t)
|
||||
local w = {}
|
||||
local inv_alpha = 1 / alpha_curr
|
||||
for i = 0, n-1 do w[i] = vt[i] + xt[i] * inv_alpha end
|
||||
|
||||
if prev_w then
|
||||
-- 2nd order: AB-like correction
|
||||
local dt_curr = t_curr - t_prev
|
||||
local dt_prev = prev_t - prev_t_dst
|
||||
local r = (dt_prev > 1e-8) and (dt_curr / dt_prev) or 1
|
||||
local c1 = 1 + 0.5 * r
|
||||
local c0 = 0.5 * r
|
||||
for i = 0, n-1 do
|
||||
local w_eff = c1 * w[i] - c0 * prev_w[i]
|
||||
xt[i] = alpha_ratio * xt[i] - alpha_prev * w_eff * log_ratio
|
||||
end
|
||||
else
|
||||
-- 1st order: exponential Euler
|
||||
for i = 0, n-1 do
|
||||
xt[i] = alpha_ratio * xt[i] - alpha_prev * w[i] * log_ratio
|
||||
end
|
||||
end
|
||||
|
||||
prev_w = w
|
||||
prev_t = t_curr
|
||||
prev_t_dst = t_prev
|
||||
end
|
||||
@@ -0,0 +1,65 @@
|
||||
-- aflops2.lua: A-FloPS Midpoint (2 NFE, stateless)
|
||||
-- Midpoint-corrected exponential integrator.
|
||||
|
||||
solver = {
|
||||
name = "aflops2",
|
||||
display = "A-FloPS Midpoint (2 NFE)",
|
||||
description = "Midpoint-corrected exponential integrator",
|
||||
nfe = 2,
|
||||
order = 2,
|
||||
needs_model = true,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
local function clamp_alpha(t)
|
||||
local a = 1 - t
|
||||
return math.max(1e-6, math.min(a, 1 - 1e-6))
|
||||
end
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
if t_curr < 1e-8 then
|
||||
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
|
||||
return
|
||||
end
|
||||
|
||||
local alpha_curr = clamp_alpha(t_curr)
|
||||
local alpha_prev = clamp_alpha(t_prev)
|
||||
|
||||
-- Save v_curr
|
||||
local v_curr = {}
|
||||
for i = 0, n-1 do v_curr[i] = vt[i] end
|
||||
|
||||
-- Euler half-step to midpoint
|
||||
local half_dt = dt * 0.5
|
||||
local t_mid = t_curr - half_dt
|
||||
local alpha_mid = clamp_alpha(t_mid)
|
||||
|
||||
local x_mid = {}
|
||||
for i = 0, n-1 do x_mid[i] = xt[i] - v_curr[i] * half_dt end
|
||||
|
||||
-- Evaluate at midpoint
|
||||
model_fn(xt, t_mid) -- xt used as scratch, but we need x_mid...
|
||||
-- Actually we need to pass x_mid to model_fn. Fix:
|
||||
-- Store xt, use x_mid for model_fn
|
||||
local xt_save = {}
|
||||
for i = 0, n-1 do xt_save[i] = xt[i]; xt[i] = x_mid[i] end
|
||||
model_fn(xt, t_mid)
|
||||
|
||||
-- Compute w_mid from midpoint
|
||||
local inv_alpha_mid = 1 / alpha_mid
|
||||
local w_mid = {}
|
||||
for i = 0, n-1 do
|
||||
w_mid[i] = vt_buf[i] + x_mid[i] * inv_alpha_mid
|
||||
end
|
||||
|
||||
-- Full step using midpoint residual
|
||||
local alpha_ratio = alpha_prev / alpha_curr
|
||||
local log_ratio = math.log(alpha_ratio)
|
||||
|
||||
for i = 0, n-1 do
|
||||
xt[i] = alpha_ratio * xt_save[i] - alpha_prev * w_mid[i] * log_ratio
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,77 @@
|
||||
-- dop853.lua: Dormand-Prince 8th order (13 NFE, fixed step)
|
||||
|
||||
solver = {
|
||||
name = "dop853",
|
||||
display = "DOP853 (13 NFE)",
|
||||
description = "8th order Dormand-Prince (maximum accuracy)",
|
||||
nfe = 13,
|
||||
order = 8,
|
||||
needs_model = true,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
local C = {
|
||||
1/18, 1/12, 1/8, 5/16, 3/8,
|
||||
59/400, 93/200, 5490023248/9719169821,
|
||||
13/20, 1201146811/1299019798,
|
||||
1, 1,
|
||||
}
|
||||
local A = {
|
||||
{1/18},
|
||||
{1/48, 1/16},
|
||||
{1/32, 0, 3/32},
|
||||
{5/16, 0, -75/64, 75/64},
|
||||
{3/80, 0, 0, 3/16, 3/20},
|
||||
{29443841/614563906, 0, 0, 77736538/692538347, -28693883/1125000000, 23124283/1800000000},
|
||||
{16016141/946692911, 0, 0, 61564180/158732637, 22789713/633445777, 545815736/2771057229, -180193667/1043307555},
|
||||
{39632708/573591083, 0, 0, -433636366/683701615, -421739975/2616292301, 100302831/723423059, 790204164/839813087, 800635310/3783071287},
|
||||
{246121993/1340847787, 0, 0, -37695042795/15268766246, -309121744/1061227803, -12992083/490766935, 6005943493/2108947869, 393006217/1396673457, 123872331/1001029789},
|
||||
{-1028468189/846180014, 0, 0, 8478235783/508512852, 1311729495/1432422823, -10304129995/1701304382, -48777925059/3047939560, 15336726248/1032824649, -45442868181/3398467696, 3065993473/597172653},
|
||||
{185892177/718116043, 0, 0, -3185094517/667107341, -477755414/1098053517, -703635378/230739211, 5731566787/1027545527, 5232866602/850066563, -4093664535/808688257, 3962137247/1805957418, 65686358/487910083},
|
||||
{403863854/491063109, 0, 0, -5068492393/434740067, -411421997/543043805, 652783627/914296604, 11173962825/925320556, -13158990841/6184727034, 3936647629/1978049680, -160528059/685178525, 248638103/1413531060, 0},
|
||||
}
|
||||
local B = {
|
||||
14005451/335480064, 0, 0, 0, 0,
|
||||
-59238493/1068277825, 181606767/758867731, 561292985/797845732,
|
||||
-1041891430/1371343529, 760417239/1151165299, 118820643/751138087,
|
||||
-528747749/2220607170, 1/4,
|
||||
}
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
local dt = t_curr - t_prev
|
||||
local xt_orig = {}
|
||||
local k1 = {}
|
||||
for i = 0, n-1 do xt_orig[i] = xt[i]; k1[i] = vt[i] end
|
||||
|
||||
local ks = {k1}
|
||||
local x_tmp = {}
|
||||
|
||||
-- 12 extra stages
|
||||
for s = 1, 12 do
|
||||
local a_row = A[s]
|
||||
for i = 0, n-1 do
|
||||
local combo = 0
|
||||
for j = 1, #a_row do
|
||||
if a_row[j] ~= 0 and ks[j] then
|
||||
combo = combo + a_row[j] * ks[j][i]
|
||||
end
|
||||
end
|
||||
x_tmp[i] = xt_orig[i] - dt * combo
|
||||
end
|
||||
-- Write x_tmp to xt for model_fn
|
||||
for i = 0, n-1 do xt[i] = x_tmp[i] end
|
||||
model_fn(xt, t_curr - C[s] * dt)
|
||||
ks[s+1] = {}
|
||||
for i = 0, n-1 do ks[s+1][i] = vt_buf[i] end
|
||||
end
|
||||
|
||||
-- Combine
|
||||
for i = 0, n-1 do
|
||||
local sol = 0
|
||||
for j = 1, 13 do
|
||||
if B[j] ~= 0 and ks[j] then sol = sol + B[j] * ks[j][i] end
|
||||
end
|
||||
xt[i] = xt_orig[i] - dt * sol
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,120 @@
|
||||
-- dopri5.lua: Dormand-Prince 5(4) adaptive solver (7+ NFE)
|
||||
-- Adaptive sub-stepping with error estimation for optimal accuracy.
|
||||
|
||||
solver = {
|
||||
name = "dopri5",
|
||||
display = "DOPRI5 (7+ NFE)",
|
||||
description = "Adaptive Dormand-Prince 5th order with error control",
|
||||
nfe = 0, -- variable
|
||||
order = 5,
|
||||
needs_model = true,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
-- Butcher tableau constants
|
||||
local C = {1/5, 3/10, 4/5, 8/9, 1, 1}
|
||||
local A = {
|
||||
{1/5},
|
||||
{3/40, 9/40},
|
||||
{44/45, -56/15, 32/9},
|
||||
{19372/6561, -25360/2187, 64448/6561, -212/729},
|
||||
{9017/3168, -355/33, 46732/5247, 49/176, -5103/18656},
|
||||
{35/384, 0, 500/1113, 125/192, -2187/6784, 11/84},
|
||||
}
|
||||
local B = {35/384, 0, 500/1113, 125/192, -2187/6784, 11/84, 0}
|
||||
local E = {
|
||||
35/384 - 1951/21600,
|
||||
0,
|
||||
500/1113 - 22642/50085,
|
||||
125/192 - 451/720,
|
||||
-2187/6784 + 12231/42400,
|
||||
11/84 - 649/6300,
|
||||
-1/60,
|
||||
}
|
||||
|
||||
-- Generic ERK step: compute all stages, return result in xt_out
|
||||
-- xt_fa is a FloatArray used as scratch space for model_fn calls
|
||||
local function erk_step(x, k1, t, h, n_elem, model_fn, vt_buf, xt_fa, num_extra, a_rows, c_vals, b_vals)
|
||||
local ks = {k1}
|
||||
for s = 1, num_extra do
|
||||
local a_row = a_rows[s]
|
||||
for i = 0, n_elem-1 do
|
||||
local combo = 0
|
||||
for j = 1, #a_row do
|
||||
if a_row[j] ~= 0 then combo = combo + a_row[j] * ks[j][i] end
|
||||
end
|
||||
xt_fa[i] = x[i] - h * combo
|
||||
end
|
||||
-- model_fn expects a FloatArray, writes result into vt_buf
|
||||
model_fn(xt_fa, t - c_vals[s] * h)
|
||||
ks[s+1] = {}
|
||||
for i = 0, n_elem-1 do ks[s+1][i] = vt_buf[i] end
|
||||
end
|
||||
|
||||
local result = {}
|
||||
for i = 0, n_elem-1 do
|
||||
local sol = 0
|
||||
for j = 1, #b_vals do
|
||||
if b_vals[j] ~= 0 and ks[j] then sol = sol + b_vals[j] * ks[j][i] end
|
||||
end
|
||||
result[i] = x[i] - h * sol
|
||||
end
|
||||
return result, ks
|
||||
end
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
local atol = 1e-3
|
||||
local rtol = 1e-2
|
||||
local max_sub = 8
|
||||
local safety = 0.9
|
||||
|
||||
local t = t_curr
|
||||
local t_end = t_prev
|
||||
local h = t - t_end
|
||||
|
||||
-- Working copy (Lua tables for intermediate math)
|
||||
local x_cur = {}
|
||||
local v_cur = {}
|
||||
for i = 0, n-1 do x_cur[i] = xt[i]; v_cur[i] = vt[i] end
|
||||
|
||||
local sub = 0
|
||||
while sub < max_sub and (t - t_end) > 1e-10 do
|
||||
h = math.min(h, t - t_end)
|
||||
local k1 = v_cur
|
||||
|
||||
-- Full DOPRI5 step (6 extra stages for 7 total including FSAL)
|
||||
-- Pass xt as scratch FloatArray for model_fn calls
|
||||
local x_next, ks = erk_step(x_cur, k1, t, h, n, model_fn, vt_buf, xt, 6, A, C, B)
|
||||
|
||||
-- Error estimate
|
||||
local err_sq_sum = 0
|
||||
for i = 0, n-1 do
|
||||
local err_i = 0
|
||||
for j = 1, 7 do
|
||||
if E[j] ~= 0 and ks[j] then err_i = err_i + E[j] * ks[j][i] end
|
||||
end
|
||||
err_i = err_i * h
|
||||
local scale = atol + rtol * math.max(math.abs(x_cur[i]), math.abs(x_next[i]))
|
||||
local ratio = err_i / scale
|
||||
err_sq_sum = err_sq_sum + ratio * ratio
|
||||
end
|
||||
local err_norm = math.sqrt(err_sq_sum / n)
|
||||
|
||||
if err_norm <= 1 then
|
||||
t = t - h
|
||||
x_cur = x_next
|
||||
v_cur = ks[7] -- FSAL
|
||||
if err_norm > 1e-10 then
|
||||
h = h * math.min(5, safety * err_norm ^ (-0.2))
|
||||
else
|
||||
h = h * 5
|
||||
end
|
||||
else
|
||||
h = h * math.max(0.2, safety * err_norm ^ (-0.2))
|
||||
end
|
||||
sub = sub + 1
|
||||
end
|
||||
|
||||
for i = 0, n-1 do xt[i] = x_cur[i] end
|
||||
end
|
||||
@@ -0,0 +1,39 @@
|
||||
-- dpm2m.lua: DPM++ 2M solver (Adams-Bashforth 2, 1 NFE, stateful)
|
||||
-- Uses previous step's velocity for 2nd order correction.
|
||||
-- First step falls back to Euler.
|
||||
|
||||
solver = {
|
||||
name = "dpm2m",
|
||||
display = "DPM++ 2M",
|
||||
description = "2nd order Adams-Bashforth multistep (1 NFE)",
|
||||
nfe = 1,
|
||||
order = 2,
|
||||
needs_model = false,
|
||||
stateful = true,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
-- Persistent state across steps (Lua tables survive between calls)
|
||||
local prev_vt = nil
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
if (step_index or 0) == 0 then prev_vt = nil end
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
if prev_vt then
|
||||
-- AB2: v_eff = 1.5 * vt - 0.5 * prev_vt
|
||||
for i = 0, n - 1 do
|
||||
local v_eff = 1.5 * vt[i] - 0.5 * prev_vt[i]
|
||||
xt[i] = xt[i] - v_eff * dt
|
||||
end
|
||||
else
|
||||
-- First step: Euler
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt[i] - vt[i] * dt
|
||||
end
|
||||
end
|
||||
|
||||
-- Save velocity for next step
|
||||
prev_vt = {}
|
||||
for i = 0, n - 1 do prev_vt[i] = vt[i] end
|
||||
end
|
||||
@@ -0,0 +1,41 @@
|
||||
-- dpm2m_ada.lua: DPM++ 2M Adaptive (step-ratio-corrected AB2)
|
||||
-- Adjusts AB2 coefficients based on step size ratio for non-uniform schedules.
|
||||
|
||||
solver = {
|
||||
name = "dpm2m_ada",
|
||||
display = "DPM++ 2M Adaptive",
|
||||
description = "Step-ratio-corrected AB2 for non-uniform schedules",
|
||||
nfe = 1,
|
||||
order = 2,
|
||||
needs_model = false,
|
||||
stateful = true,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
local prev_vt = nil
|
||||
local prev_dt = 0
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
if (step_index or 0) == 0 then prev_vt = nil; prev_dt = 0 end
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
if prev_vt and prev_dt > 0 then
|
||||
-- Step-ratio corrected AB2
|
||||
local r = dt / prev_dt
|
||||
local c1 = 1.0 + r / 2.0
|
||||
local c0 = r / 2.0
|
||||
for i = 0, n - 1 do
|
||||
local v_eff = c1 * vt[i] - c0 * prev_vt[i]
|
||||
xt[i] = xt[i] - v_eff * dt
|
||||
end
|
||||
else
|
||||
-- First step: Euler
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt[i] - vt[i] * dt
|
||||
end
|
||||
end
|
||||
|
||||
prev_vt = {}
|
||||
for i = 0, n - 1 do prev_vt[i] = vt[i] end
|
||||
prev_dt = dt
|
||||
end
|
||||
@@ -0,0 +1,45 @@
|
||||
-- dpm3m.lua: DPM++ 3M solver (Adams-Bashforth 3, 1 NFE, stateful)
|
||||
-- Uses two previous velocities for 3rd order correction.
|
||||
|
||||
solver = {
|
||||
name = "dpm3m",
|
||||
display = "DPM++ 3M",
|
||||
description = "3rd order Adams-Bashforth multistep (1 NFE)",
|
||||
nfe = 1,
|
||||
order = 3,
|
||||
needs_model = false,
|
||||
stateful = true,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
local prev_vt = nil
|
||||
local prev_prev_vt = nil
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
if (step_index or 0) == 0 then prev_vt = nil; prev_prev_vt = nil end
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
if prev_vt and prev_prev_vt then
|
||||
-- AB3: v_eff = (23/12)*vt - (16/12)*prev - (5/12)*prev_prev
|
||||
for i = 0, n - 1 do
|
||||
local v_eff = (23/12) * vt[i] - (16/12) * prev_vt[i] + (5/12) * prev_prev_vt[i]
|
||||
xt[i] = xt[i] - v_eff * dt
|
||||
end
|
||||
elseif prev_vt then
|
||||
-- AB2 fallback
|
||||
for i = 0, n - 1 do
|
||||
local v_eff = 1.5 * vt[i] - 0.5 * prev_vt[i]
|
||||
xt[i] = xt[i] - v_eff * dt
|
||||
end
|
||||
else
|
||||
-- Euler fallback
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt[i] - vt[i] * dt
|
||||
end
|
||||
end
|
||||
|
||||
-- Shift history
|
||||
prev_prev_vt = prev_vt
|
||||
prev_vt = {}
|
||||
for i = 0, n - 1 do prev_vt[i] = vt[i] end
|
||||
end
|
||||
@@ -0,0 +1,21 @@
|
||||
-- euler.lua: First-order Euler ODE solver
|
||||
-- Single evaluation, simplest possible solver.
|
||||
-- xt_next = xt - vt * (t_curr - t_prev)
|
||||
|
||||
solver = {
|
||||
name = "euler",
|
||||
display = "Euler (ODE)",
|
||||
description = "First-order Euler step (default)",
|
||||
nfe = 1,
|
||||
order = 1,
|
||||
needs_model = false,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
local dt = t_curr - t_prev
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt[i] - vt[i] * dt
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,60 @@
|
||||
-- gl2s.lua: Gauss-Legendre 2-stage implicit Runge-Kutta (4th order, 6 NFE)
|
||||
-- A-stable, symplectic. Fixed-point iteration solves the implicit system.
|
||||
|
||||
solver = {
|
||||
name = "gl2s",
|
||||
display = "Gauss-Legendre 2s (6 NFE)",
|
||||
description = "Implicit 4th-order A-stable symplectic integrator",
|
||||
nfe = 6,
|
||||
order = 4,
|
||||
needs_model = true,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
-- Butcher tableau
|
||||
local SQRT3_6 = 0.28867513459481287
|
||||
local C1 = 0.5 - SQRT3_6 -- ≈ 0.2113
|
||||
local C2 = 0.5 + SQRT3_6 -- ≈ 0.7887
|
||||
local A11 = 0.25
|
||||
local A12 = 0.25 - SQRT3_6 -- ≈ -0.0387
|
||||
local A21 = 0.25 + SQRT3_6 -- ≈ 0.5387
|
||||
local A22 = 0.25
|
||||
local ITERATIONS = 3
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
-- Initialize k1 = k2 = vt
|
||||
local k1 = {}
|
||||
local k2 = {}
|
||||
for i = 0, n-1 do k1[i] = vt[i]; k2[i] = vt[i] end
|
||||
|
||||
local xt_orig = {}
|
||||
for i = 0, n-1 do xt_orig[i] = xt[i] end
|
||||
|
||||
local t1 = t_curr - C1 * dt
|
||||
local t2 = t_curr - C2 * dt
|
||||
|
||||
-- Fixed-point iteration
|
||||
for iter = 1, ITERATIONS do
|
||||
-- Stage 1: x1 = xt - dt*(A11*k1 + A12*k2)
|
||||
for i = 0, n-1 do
|
||||
xt[i] = xt_orig[i] - dt * (A11 * k1[i] + A12 * k2[i])
|
||||
end
|
||||
model_fn(xt, t1)
|
||||
for i = 0, n-1 do k1[i] = vt_buf[i] end
|
||||
|
||||
-- Stage 2: x2 = xt - dt*(A21*k1 + A22*k2)
|
||||
for i = 0, n-1 do
|
||||
xt[i] = xt_orig[i] - dt * (A21 * k1[i] + A22 * k2[i])
|
||||
end
|
||||
model_fn(xt, t2)
|
||||
for i = 0, n-1 do k2[i] = vt_buf[i] end
|
||||
end
|
||||
|
||||
-- Final: xt = xt_orig - dt * 0.5 * (k1 + k2)
|
||||
for i = 0, n-1 do
|
||||
xt[i] = xt_orig[i] - dt * 0.5 * (k1[i] + k2[i])
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,34 @@
|
||||
-- heun.lua: Heun's method (improved Euler / explicit trapezoidal)
|
||||
-- 2 NFE: evaluate at t_curr, predict, evaluate at t_prev, average.
|
||||
|
||||
solver = {
|
||||
name = "heun",
|
||||
display = "Heun (2 NFE)",
|
||||
description = "Second-order predictor-corrector",
|
||||
nfe = 2,
|
||||
order = 2,
|
||||
needs_model = true,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
-- Predict: xt_pred = xt - vt * dt
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt[i] - vt[i] * dt
|
||||
end
|
||||
|
||||
-- Correct: evaluate at (xt_pred, t_prev)
|
||||
model_fn(xt, t_prev)
|
||||
|
||||
-- Average: xt = xt + 0.5 * (vt_buf - vt) * dt
|
||||
-- Note: xt is already xt_pred = xt_orig - vt * dt
|
||||
-- We want: xt_orig - 0.5*(vt + vt_buf)*dt
|
||||
-- = (xt + vt*dt) - 0.5*(vt + vt_buf)*dt
|
||||
-- = xt + 0.5*(vt - vt_buf)*dt
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt[i] + 0.5 * (vt[i] - vt_buf[i]) * dt
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,111 @@
|
||||
-- jkass_fast.lua: JKASS Fast solver (1 NFE, stateful)
|
||||
-- Euler with momentum blending, frequency damping, and temporal smoothing.
|
||||
-- Port from jeankassio/JK-AceStep-Nodes.
|
||||
|
||||
solver = {
|
||||
name = "jkass_fast",
|
||||
display = "JKASS Fast",
|
||||
description = "Euler with beat stability, frequency damping, and temporal smoothing",
|
||||
accent = "amber",
|
||||
nfe = 1,
|
||||
order = 1,
|
||||
needs_model = false,
|
||||
stateful = true,
|
||||
stochastic = false,
|
||||
params = {
|
||||
{ key = "beat_stability", type = "slider", label = "Beat Stability",
|
||||
default = 0.25, min = 0, max = 1, step = 0.01,
|
||||
hint = "Momentum blend with previous step (0=off, 1=full momentum)" },
|
||||
{ key = "frequency_damping", type = "slider", label = "Frequency Damping",
|
||||
default = 0.4, min = 0, max = 5, step = 0.1,
|
||||
hint = "Attenuate high-frequency bins (0=off)" },
|
||||
{ key = "temporal_smoothing", type = "slider", label = "Temporal Smoothing",
|
||||
default = 0.13, min = 0, max = 1, step = 0.01,
|
||||
hint = "1D blur across time axis (0=off)" },
|
||||
},
|
||||
}
|
||||
|
||||
local prev_delta = nil
|
||||
|
||||
-- Frequency damping: exponential decay across channel dimension
|
||||
local function apply_frequency_damping(data, offset, T, Oc, damping)
|
||||
if damping <= 0 then return end
|
||||
local freq_mult = {}
|
||||
for c = 0, Oc - 1 do
|
||||
local freq = c / (Oc - 1)
|
||||
freq_mult[c] = math.exp(-damping * freq * freq)
|
||||
end
|
||||
for t = 0, T - 1 do
|
||||
for c = 0, Oc - 1 do
|
||||
local idx = offset + t * Oc + c
|
||||
data[idx] = data[idx] * freq_mult[c]
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
-- Temporal smoothing: [0.25, 0.5, 0.25] blur across time axis
|
||||
local function apply_temporal_smoothing(data, offset, T, Oc, strength)
|
||||
if strength <= 0 or T < 3 then return end
|
||||
local smoothed = {}
|
||||
for c = 0, Oc - 1 do
|
||||
for t = 0, T - 1 do
|
||||
local t_prev = (t > 0) and (t - 1) or 1
|
||||
local t_next = (t < T - 1) and (t + 1) or (T - 2)
|
||||
local v_prev = data[offset + t_prev * Oc + c]
|
||||
local v_curr = data[offset + t * Oc + c]
|
||||
local v_next = data[offset + t_next * Oc + c]
|
||||
smoothed[t * Oc + c] = 0.25 * v_prev + 0.5 * v_curr + 0.25 * v_next
|
||||
end
|
||||
end
|
||||
for i = 0, T * Oc - 1 do
|
||||
data[offset + i] = (1 - strength) * data[offset + i] + strength * smoothed[i]
|
||||
end
|
||||
end
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
if (step_index or 0) == 0 then prev_delta = nil end
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
-- Read params (injected by C++ before each call)
|
||||
local bs = params and params.beat_stability or 0.25
|
||||
local fd = params and params.frequency_damping or 0.4
|
||||
local ts = params and params.temporal_smoothing or 0.13
|
||||
|
||||
-- Copy velocity as working delta
|
||||
local delta = {}
|
||||
for i = 0, n - 1 do delta[i] = vt[i] end
|
||||
|
||||
-- Beat stability: momentum blend
|
||||
if prev_delta and bs > 0 then
|
||||
for i = 0, n - 1 do
|
||||
delta[i] = (1 - bs) * delta[i] + bs * prev_delta[i]
|
||||
end
|
||||
end
|
||||
|
||||
-- Save for next step
|
||||
prev_delta = {}
|
||||
for i = 0, n - 1 do prev_delta[i] = delta[i] end
|
||||
|
||||
-- Frequency damping (per batch item, Oc=64 for ACE-Step)
|
||||
if fd > 0 and n_per and n_per > 0 then
|
||||
local Oc = 64
|
||||
local T = n_per / Oc
|
||||
for b = 0, batch_n - 1 do
|
||||
apply_frequency_damping(delta, b * n_per, T, Oc, fd)
|
||||
end
|
||||
end
|
||||
|
||||
-- Temporal smoothing (per batch item)
|
||||
if ts > 0 and n_per and n_per > 0 then
|
||||
local Oc = 64
|
||||
local T = n_per / Oc
|
||||
for b = 0, batch_n - 1 do
|
||||
apply_temporal_smoothing(delta, b * n_per, T, Oc, ts)
|
||||
end
|
||||
end
|
||||
|
||||
-- Euler step with modified delta
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt[i] - delta[i] * dt
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,40 @@
|
||||
-- jkass_quality.lua: JKASS Quality solver (2 NFE)
|
||||
-- Heun with derivative averaging.
|
||||
|
||||
solver = {
|
||||
name = "jkass_quality",
|
||||
display = "JKASS Quality (2 NFE)",
|
||||
description = "Heun predictor-corrector with derivative averaging",
|
||||
accent = "amber",
|
||||
nfe = 2,
|
||||
order = 2,
|
||||
needs_model = true,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
if t_prev <= 0 then
|
||||
-- Fallback to Euler for final step
|
||||
for i = 0, n - 1 do xt[i] = xt[i] - vt[i] * dt end
|
||||
return
|
||||
end
|
||||
|
||||
-- Save k1
|
||||
local k1 = {}
|
||||
for i = 0, n - 1 do k1[i] = vt[i] end
|
||||
|
||||
-- Euler predictor
|
||||
for i = 0, n - 1 do xt[i] = xt[i] - vt[i] * dt end
|
||||
|
||||
-- Second evaluation
|
||||
model_fn(xt, t_prev)
|
||||
|
||||
-- Correct: undo Euler, apply averaged derivative
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt[i] + k1[i] * dt -- undo Euler
|
||||
xt[i] = xt[i] - 0.5 * (k1[i] + vt_buf[i]) * dt -- Heun average
|
||||
end
|
||||
end
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,40 @@
|
||||
-- rfsolver.lua: RF-Solver (2 NFE) — Rectified Flow specific
|
||||
-- Exploits the RF ODE structure for higher accuracy than generic midpoint.
|
||||
|
||||
solver = {
|
||||
name = "rfsolver",
|
||||
display = "RF-Solver (2 NFE)",
|
||||
description = "Rectified-flow-aware midpoint solver",
|
||||
nfe = 2,
|
||||
order = 2,
|
||||
needs_model = true,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
if t_curr < 1e-8 then
|
||||
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
|
||||
return
|
||||
end
|
||||
|
||||
-- Save v_t
|
||||
local v_t = {}
|
||||
for i = 0, n-1 do v_t[i] = vt[i] end
|
||||
|
||||
-- Half-step to midpoint
|
||||
local half_dt = dt * 0.5
|
||||
local t_mid = t_curr - half_dt
|
||||
for i = 0, n-1 do xt[i] = xt[i] - v_t[i] * half_dt end
|
||||
|
||||
-- Evaluate at midpoint
|
||||
model_fn(xt, t_mid)
|
||||
|
||||
-- RF-specific: reconstruct via x_0 prediction from midpoint
|
||||
for i = 0, n-1 do
|
||||
local x_0_mid = xt[i] - t_mid * vt_buf[i]
|
||||
xt[i] = x_0_mid + t_prev * vt_buf[i]
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,54 @@
|
||||
-- rk4.lua: Classic 4th-order Runge-Kutta solver
|
||||
-- 4 NFE per step, excellent accuracy for smooth flows.
|
||||
|
||||
solver = {
|
||||
name = "rk4",
|
||||
display = "RK4 (4 NFE)",
|
||||
description = "Classic 4th-order Runge-Kutta",
|
||||
nfe = 4,
|
||||
order = 4,
|
||||
needs_model = true,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
local dt = t_curr - t_prev
|
||||
local t_mid = t_curr - 0.5 * dt
|
||||
|
||||
-- k1 = vt (already evaluated)
|
||||
-- Save k1 and original xt
|
||||
local k1 = {}
|
||||
local xt_orig = {}
|
||||
for i = 0, n - 1 do
|
||||
k1[i] = vt[i]
|
||||
xt_orig[i] = xt[i]
|
||||
end
|
||||
|
||||
-- k2: evaluate at midpoint using k1
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt_orig[i] - 0.5 * k1[i] * dt
|
||||
end
|
||||
model_fn(xt, t_mid)
|
||||
local k2 = {}
|
||||
for i = 0, n - 1 do k2[i] = vt_buf[i] end
|
||||
|
||||
-- k3: evaluate at midpoint using k2
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt_orig[i] - 0.5 * k2[i] * dt
|
||||
end
|
||||
model_fn(xt, t_mid)
|
||||
local k3 = {}
|
||||
for i = 0, n - 1 do k3[i] = vt_buf[i] end
|
||||
|
||||
-- k4: evaluate at endpoint using k3
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt_orig[i] - k3[i] * dt
|
||||
end
|
||||
model_fn(xt, t_prev)
|
||||
|
||||
-- Combine: xt = xt_orig - (k1 + 2*k2 + 2*k3 + k4) * dt / 6
|
||||
for i = 0, n - 1 do
|
||||
xt[i] = xt_orig[i] - (k1[i] + 2*k2[i] + 2*k3[i] + vt_buf[i]) * dt / 6.0
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,64 @@
|
||||
-- rk5.lua: 5th-order Runge-Kutta (Cash-Karp, 6 NFE per step)
|
||||
|
||||
solver = {
|
||||
name = "rk5",
|
||||
display = "RK5 (6 NFE)",
|
||||
description = "5th-order Cash-Karp Runge-Kutta",
|
||||
nfe = 6,
|
||||
order = 5,
|
||||
needs_model = true,
|
||||
stateful = false,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
-- Cash-Karp Butcher tableau
|
||||
local a2 = 1/5
|
||||
local a3 = 3/10
|
||||
local a4 = 3/5
|
||||
local a5 = 1
|
||||
local a6 = 7/8
|
||||
|
||||
local b21 = 1/5
|
||||
local b31 = 3/40; local b32 = 9/40
|
||||
local b41 = 3/10; local b42 = -9/10; local b43 = 6/5
|
||||
local b51 = -11/54; local b52 = 5/2; local b53 = -70/27; local b54 = 35/27
|
||||
local b61 = 1631/55296; local b62 = 175/512; local b63 = 575/13824; local b64 = 44275/110592; local b65 = 253/4096
|
||||
|
||||
-- 5th order weights
|
||||
local c1 = 37/378; local c3 = 250/621; local c4 = 125/594; local c6 = 512/1771
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
local dt = t_curr - t_prev
|
||||
local xt_orig = {}
|
||||
local k1 = {}
|
||||
for i = 0, n-1 do xt_orig[i] = xt[i]; k1[i] = vt[i] end
|
||||
|
||||
-- k2
|
||||
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * b21 * k1[i] end
|
||||
model_fn(xt, t_curr - a2 * dt)
|
||||
local k2 = {}; for i = 0, n-1 do k2[i] = vt_buf[i] end
|
||||
|
||||
-- k3
|
||||
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * (b31*k1[i] + b32*k2[i]) end
|
||||
model_fn(xt, t_curr - a3 * dt)
|
||||
local k3 = {}; for i = 0, n-1 do k3[i] = vt_buf[i] end
|
||||
|
||||
-- k4
|
||||
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * (b41*k1[i] + b42*k2[i] + b43*k3[i]) end
|
||||
model_fn(xt, t_curr - a4 * dt)
|
||||
local k4 = {}; for i = 0, n-1 do k4[i] = vt_buf[i] end
|
||||
|
||||
-- k5
|
||||
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * (b51*k1[i] + b52*k2[i] + b53*k3[i] + b54*k4[i]) end
|
||||
model_fn(xt, t_curr - a5 * dt)
|
||||
local k5 = {}; for i = 0, n-1 do k5[i] = vt_buf[i] end
|
||||
|
||||
-- k6
|
||||
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * (b61*k1[i] + b62*k2[i] + b63*k3[i] + b64*k4[i] + b65*k5[i]) end
|
||||
model_fn(xt, t_curr - a6 * dt)
|
||||
|
||||
-- Final 5th-order result
|
||||
for i = 0, n-1 do
|
||||
xt[i] = xt_orig[i] - dt * (c1*k1[i] + c3*k3[i] + c4*k4[i] + c6*vt_buf[i])
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,28 @@
|
||||
-- sde.lua: Stochastic Differential Equation Euler solver (1 NFE)
|
||||
-- Predicts x0 then re-noises with Philox random numbers.
|
||||
-- Note: requires philox_randn C helper exposed to Lua.
|
||||
-- Falls back to ODE Euler if no seed info available.
|
||||
|
||||
solver = {
|
||||
name = "sde",
|
||||
display = "SDE (Stochastic)",
|
||||
description = "SDE Euler with Philox re-noising for diversity",
|
||||
nfe = 1,
|
||||
order = 1,
|
||||
needs_model = false,
|
||||
stateful = false,
|
||||
stochastic = true,
|
||||
}
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
-- SDE needs philox_randn which is a C helper.
|
||||
-- For now, fall back to ODE Euler. The C++ wrapper will handle
|
||||
-- the SDE-specific logic (philox noise injection) when it detects
|
||||
-- this plugin is stochastic and seeds are available.
|
||||
local dt = t_curr - t_prev
|
||||
for i = 0, n - 1 do
|
||||
local x0 = xt[i] - vt[i] * t_curr
|
||||
-- Without philox, do deterministic ODE step
|
||||
xt[i] = xt[i] - vt[i] * dt
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,139 @@
|
||||
-- stork2.lua: STORK 2 — Stabilized Taylor Orthogonal RK (2nd order, 1 NFE)
|
||||
-- Uses RKG2 Chebyshev sub-stepping with velocity derivatives from history.
|
||||
|
||||
solver = {
|
||||
name = "stork2",
|
||||
display = "STORK 2",
|
||||
description = "2nd-order stabilized Taylor-Chebyshev (1 NFE)",
|
||||
nfe = 1,
|
||||
order = 2,
|
||||
needs_model = false,
|
||||
stateful = true,
|
||||
stochastic = false,
|
||||
params = {
|
||||
{ key = "substeps", type = "slider", label = "Substeps",
|
||||
default = 10, min = 2, max = 50, step = 1,
|
||||
hint = "Number of Chebyshev sub-steps (more = more stable)" },
|
||||
},
|
||||
}
|
||||
|
||||
local velocity_history = {} -- {vt={}, dt=float}
|
||||
|
||||
local function rms(data, n)
|
||||
local sum = 0
|
||||
for i = 0, n-1 do sum = sum + data[i] * data[i] end
|
||||
return math.sqrt(sum / n)
|
||||
end
|
||||
|
||||
local function has_nan_inf(data, n)
|
||||
for i = 0, n-1 do
|
||||
local v = data[i]
|
||||
if v ~= v or v == math.huge or v == -math.huge then return true end
|
||||
end
|
||||
return false
|
||||
end
|
||||
|
||||
local function compute_derivatives(vt, n)
|
||||
local hist = velocity_history
|
||||
if #hist == 0 then return 0, nil, nil end
|
||||
local v_prev = hist[#hist].vt
|
||||
local h1 = hist[#hist].dt
|
||||
if h1 == 0 then return 0, nil, nil end
|
||||
local dv = {}
|
||||
for i = 0, n-1 do dv[i] = (v_prev[i] - vt[i]) / h1 end
|
||||
local vt_rms = rms(vt, n)
|
||||
local dv_rms = rms(dv, n)
|
||||
if vt_rms > 0 and dv_rms * math.abs(h1) > 5 * vt_rms then return 0, nil, nil end
|
||||
if #hist < 2 then return 1, dv, nil end
|
||||
local v_prev2 = hist[#hist - 1].vt
|
||||
local h2 = hist[#hist - 1].dt
|
||||
if h2 == 0 then return 1, dv, nil end
|
||||
local denom = h1 * h2 * (h1 + h2)
|
||||
if math.abs(denom) < 1e-30 then return 1, dv, nil end
|
||||
local coeff = 2 / denom
|
||||
local d2v = {}
|
||||
for i = 0, n-1 do
|
||||
d2v[i] = coeff * (v_prev2[i] * h1 - v_prev[i] * (h1 + h2) + vt[i] * h2)
|
||||
end
|
||||
local d2v_rms = rms(d2v, n)
|
||||
if vt_rms > 0 and d2v_rms * h1 * h1 > 5 * vt_rms then return 1, dv, nil end
|
||||
return 2, dv, d2v
|
||||
end
|
||||
|
||||
local function taylor_approx(vt, deriv_order, diff, dv, d2v, n)
|
||||
local out = {}
|
||||
if deriv_order >= 2 and d2v then
|
||||
local half_d2 = 0.5 * diff * diff
|
||||
for i = 0, n-1 do out[i] = vt[i] + diff * dv[i] + half_d2 * d2v[i] end
|
||||
elseif deriv_order >= 1 and dv then
|
||||
for i = 0, n-1 do out[i] = vt[i] + diff * dv[i] end
|
||||
else
|
||||
for i = 0, n-1 do out[i] = vt[i] end
|
||||
end
|
||||
return out
|
||||
end
|
||||
|
||||
local function rkg2_b(j)
|
||||
if j <= 0 then return 1 end
|
||||
if j == 1 then return 1/3 end
|
||||
return 4 * (j - 1) * (j + 4) / (3 * j * (j + 1) * (j + 2) * (j + 3))
|
||||
end
|
||||
|
||||
local function rkg2_substep(xt, vt, s, t_curr, t_prev, deriv_order, dv, d2v, n)
|
||||
local dt = t_curr - t_prev
|
||||
local Y_j_2 = {}; local Y_j_1 = {}; local Y_j = {}
|
||||
for i = 0, n-1 do Y_j_2[i] = xt[i]; Y_j_1[i] = xt[i]; Y_j[i] = xt[i] end
|
||||
local s2ps = s * s + s - 2
|
||||
for j = 1, s do
|
||||
if j == 1 then
|
||||
local mu_t = 6 / ((s + 4) * (s - 1))
|
||||
for i = 0, n-1 do Y_j[i] = Y_j_1[i] - dt * mu_t * vt[i] end
|
||||
else
|
||||
local frac = (j == 2) and (4 / (3 * s2ps)) or ((j-1)*(j-1)+(j-1)-2) / s2ps
|
||||
local bj = rkg2_b(j); local bj1 = rkg2_b(j-1); local bj2 = rkg2_b(j-2)
|
||||
local mu = (2*j+1) * bj / (j * bj1)
|
||||
local nu = -(j+1) * bj / (j * bj2)
|
||||
local mu_t = mu * 6 / ((s+4)*(s-1))
|
||||
local gamma_t = -mu_t * (1 - j*(j+1)*bj1/2)
|
||||
local diff = -frac * dt
|
||||
local vel = taylor_approx(vt, deriv_order, diff, dv, d2v, n)
|
||||
for i = 0, n-1 do
|
||||
Y_j[i] = mu*Y_j_1[i] + nu*Y_j_2[i] + (1-mu-nu)*xt[i]
|
||||
- dt*mu_t*vel[i] - dt*gamma_t*vt[i]
|
||||
end
|
||||
end
|
||||
for i = 0, n-1 do Y_j_2[i] = Y_j_1[i]; Y_j_1[i] = Y_j[i] end
|
||||
end
|
||||
return Y_j, not has_nan_inf(Y_j, n)
|
||||
end
|
||||
|
||||
local function update_history(vt, n, dt)
|
||||
local rec = {vt = {}, dt = dt}
|
||||
for i = 0, n-1 do rec.vt[i] = vt[i] end
|
||||
table.insert(velocity_history, rec)
|
||||
while #velocity_history > 3 do table.remove(velocity_history, 1) end
|
||||
end
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
local dt = t_curr - t_prev
|
||||
if step_index == 0 then
|
||||
velocity_history = {}
|
||||
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
|
||||
update_history(vt, n, dt)
|
||||
return
|
||||
end
|
||||
local deriv_order, dv, d2v = compute_derivatives(vt, n)
|
||||
local s = math.max((params and params.substeps) or 10, 2)
|
||||
local success = false; local result
|
||||
while s >= 2 do
|
||||
result, success = rkg2_substep(xt, vt, s, t_curr, t_prev, deriv_order, dv, d2v, n)
|
||||
if success then break end
|
||||
s = math.floor(s / 2)
|
||||
end
|
||||
if success then
|
||||
for i = 0, n-1 do xt[i] = result[i] end
|
||||
else
|
||||
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
|
||||
end
|
||||
update_history(vt, n, dt)
|
||||
end
|
||||
@@ -0,0 +1,186 @@
|
||||
-- stork4.lua: STORK 4 — Stabilized Taylor Orthogonal RK (4th order, 1 NFE)
|
||||
-- Uses ROCK4 Chebyshev sub-stepping with precomputed coefficients.
|
||||
-- Requires companion stork4_constants.lua.
|
||||
|
||||
local C = require("stork4_constants")
|
||||
|
||||
solver = {
|
||||
name = "stork4",
|
||||
display = "STORK 4",
|
||||
description = "4th-order stabilized Taylor-ROCK4 (1 NFE)",
|
||||
nfe = 1,
|
||||
order = 4,
|
||||
needs_model = false,
|
||||
stateful = true,
|
||||
stochastic = false,
|
||||
params = {
|
||||
{ key = "substeps", type = "slider", label = "Substeps",
|
||||
default = 10, min = 2, max = 50, step = 1,
|
||||
hint = "Number of ROCK4 sub-steps (more = more stable)" },
|
||||
},
|
||||
}
|
||||
|
||||
local velocity_history = {}
|
||||
|
||||
local function rms(data, n)
|
||||
local sum = 0
|
||||
for i = 0, n-1 do sum = sum + data[i] * data[i] end
|
||||
return math.sqrt(sum / n)
|
||||
end
|
||||
|
||||
local function has_nan_inf(data, n)
|
||||
for i = 0, n-1 do
|
||||
local v = data[i]
|
||||
if v ~= v or v == math.huge or v == -math.huge then return true end
|
||||
end
|
||||
return false
|
||||
end
|
||||
|
||||
local function compute_derivatives(vt, n)
|
||||
local hist = velocity_history
|
||||
if #hist == 0 then return 0, nil, nil end
|
||||
local v_prev = hist[#hist].vt
|
||||
local h1 = hist[#hist].dt
|
||||
if h1 == 0 then return 0, nil, nil end
|
||||
local dv = {}
|
||||
for i = 0, n-1 do dv[i] = (v_prev[i] - vt[i]) / h1 end
|
||||
local vt_rms = rms(vt, n)
|
||||
local dv_rms = rms(dv, n)
|
||||
if vt_rms > 0 and dv_rms * math.abs(h1) > 5 * vt_rms then return 0, nil, nil end
|
||||
if #hist < 2 then return 1, dv, nil end
|
||||
local v_prev2 = hist[#hist - 1].vt
|
||||
local h2 = hist[#hist - 1].dt
|
||||
if h2 == 0 then return 1, dv, nil end
|
||||
local denom = h1 * h2 * (h1 + h2)
|
||||
if math.abs(denom) < 1e-30 then return 1, dv, nil end
|
||||
local coeff = 2 / denom
|
||||
local d2v = {}
|
||||
for i = 0, n-1 do
|
||||
d2v[i] = coeff * (v_prev2[i] * h1 - v_prev[i] * (h1 + h2) + vt[i] * h2)
|
||||
end
|
||||
local d2v_rms = rms(d2v, n)
|
||||
if vt_rms > 0 and d2v_rms * h1 * h1 > 5 * vt_rms then return 1, dv, nil end
|
||||
return 2, dv, d2v
|
||||
end
|
||||
|
||||
local function taylor_approx(vt, order, diff, dv, d2v, n)
|
||||
local out = {}
|
||||
if order >= 2 and d2v then
|
||||
local half_d2 = 0.5 * diff * diff
|
||||
for i = 0, n-1 do out[i] = vt[i] + diff * dv[i] + half_d2 * d2v[i] end
|
||||
elseif order >= 1 and dv then
|
||||
for i = 0, n-1 do out[i] = vt[i] + diff * dv[i] end
|
||||
else
|
||||
for i = 0, n-1 do out[i] = vt[i] end
|
||||
end
|
||||
return out
|
||||
end
|
||||
|
||||
local function rock4_mdegr(s)
|
||||
local mp1 = 1
|
||||
for i = 1, C.MS_LEN do
|
||||
if C.MS[i] >= s then
|
||||
return C.MS[i], i, mp1 - 1
|
||||
end
|
||||
mp1 = mp1 + C.MS[i] * 2 - 1
|
||||
end
|
||||
return C.MS[C.MS_LEN], C.MS_LEN, mp1 - 1
|
||||
end
|
||||
|
||||
local function rock4_substep(xt, vt, s, t_curr, t_prev, deriv_order, dv, d2v, n)
|
||||
local max_rock4 = C.MS[C.MS_LEN]
|
||||
if s > max_rock4 then s = max_rock4 end
|
||||
local mdeg, mz, mr = rock4_mdegr(s)
|
||||
local dt = t_curr - t_prev
|
||||
|
||||
local Y_j_2 = {}; local Y_j_1 = {}; local Y_j = {}
|
||||
for i = 0, n-1 do Y_j_2[i] = xt[i]; Y_j_1[i] = xt[i] end
|
||||
|
||||
local ci1 = t_curr
|
||||
|
||||
-- ROCK4 Chebyshev recurrence (1-indexed RECF)
|
||||
for j = 1, mdeg do
|
||||
if j == 1 then
|
||||
local temp1 = -dt * C.RECF[mr + 1]
|
||||
ci1 = t_curr + temp1
|
||||
for i = 0, n-1 do Y_j_1[i] = xt[i] + temp1 * vt[i] end
|
||||
else
|
||||
local diff = ci1 - t_curr
|
||||
local vel = taylor_approx(vt, deriv_order, diff, dv, d2v, n)
|
||||
local idx1 = mr + 2 * (j - 2) + 2
|
||||
local idx2 = mr + 2 * (j - 2) + 3
|
||||
local temp1 = -dt * C.RECF[idx1]
|
||||
local temp3 = -C.RECF[idx2]
|
||||
local temp2 = 1 - temp3
|
||||
for i = 0, n-1 do
|
||||
Y_j[i] = temp1 * vel[i] + temp2 * Y_j_1[i] + temp3 * Y_j_2[i]
|
||||
end
|
||||
for i = 0, n-1 do Y_j_2[i] = Y_j_1[i]; Y_j_1[i] = Y_j[i] end
|
||||
ci1 = temp1 + temp2 * ci1 + temp3 * ci1 -- simplified
|
||||
end
|
||||
end
|
||||
|
||||
-- ROCK4 finishing procedure (4 stages)
|
||||
local Y_base = Y_j_1
|
||||
local fpa = C.FPA[mz]; local fpb = C.FPB[mz]
|
||||
|
||||
-- F1
|
||||
local diff1 = ci1 - t_curr
|
||||
local F1 = taylor_approx(vt, deriv_order, diff1, dv, d2v, n)
|
||||
local fpa0 = -dt * fpa[1]
|
||||
local Yf = {}
|
||||
for i = 0, n-1 do Yf[i] = Y_base[i] + fpa0 * F1[i] end
|
||||
|
||||
-- F2
|
||||
local diff2 = ci1 + fpa0 - t_curr
|
||||
local F2 = taylor_approx(vt, deriv_order, diff2, dv, d2v, n)
|
||||
local fpa1 = -dt * fpa[2]; local fpa2 = -dt * fpa[3]
|
||||
for i = 0, n-1 do Yf[i] = Y_base[i] + fpa1 * F1[i] + fpa2 * F2[i] end
|
||||
|
||||
-- F3
|
||||
local diff3 = ci1 + fpa1 + fpa2 - t_curr
|
||||
local F3 = taylor_approx(vt, deriv_order, diff3, dv, d2v, n)
|
||||
local fpa3 = -dt * fpa[4]; local fpa4 = -dt * fpa[5]; local fpa5 = -dt * fpa[6]
|
||||
|
||||
-- F4
|
||||
local diff4 = ci1 + fpa3 + fpa4 + fpa5 - t_curr
|
||||
local F4 = taylor_approx(vt, deriv_order, diff4, dv, d2v, n)
|
||||
local fpb0 = -dt * fpb[1]; local fpb1 = -dt * fpb[2]; local fpb2 = -dt * fpb[3]; local fpb3 = -dt * fpb[4]
|
||||
|
||||
local result = {}
|
||||
for i = 0, n-1 do
|
||||
result[i] = Y_base[i] + fpb0*F1[i] + fpb1*F2[i] + fpb2*F3[i] + fpb3*F4[i]
|
||||
end
|
||||
return result, not has_nan_inf(result, n)
|
||||
end
|
||||
|
||||
local function update_history(vt, n, dt)
|
||||
local rec = {vt = {}, dt = dt}
|
||||
for i = 0, n-1 do rec.vt[i] = vt[i] end
|
||||
table.insert(velocity_history, rec)
|
||||
while #velocity_history > 3 do table.remove(velocity_history, 1) end
|
||||
end
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
local dt = t_curr - t_prev
|
||||
if step_index == 0 then
|
||||
velocity_history = {}
|
||||
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
|
||||
update_history(vt, n, dt)
|
||||
return
|
||||
end
|
||||
local deriv_order, dv, d2v = compute_derivatives(vt, n)
|
||||
local s = math.max((params and params.substeps) or 10, 2)
|
||||
local success = false; local result
|
||||
while s >= 2 do
|
||||
result, success = rock4_substep(xt, vt, s, t_curr, t_prev, deriv_order, dv, d2v, n)
|
||||
if success then break end
|
||||
s = math.floor(s / 2)
|
||||
end
|
||||
if success then
|
||||
for i = 0, n-1 do xt[i] = result[i] end
|
||||
else
|
||||
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
|
||||
end
|
||||
update_history(vt, n, dt)
|
||||
end
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,158 @@
|
||||
-- unipc.lua: UniPC (Unified Predictor-Corrector, 2 NFE)
|
||||
-- B(h)1 variant with data prediction in log-SNR space.
|
||||
|
||||
solver = {
|
||||
name = "unipc",
|
||||
display = "UniPC (2 NFE)",
|
||||
description = "Unified predictor-corrector in log-SNR space",
|
||||
nfe = 2,
|
||||
order = 2,
|
||||
needs_model = true,
|
||||
stateful = true,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
local history = {} -- {model_output={}, t=float}
|
||||
local max_order = 2
|
||||
|
||||
local function lambda(t)
|
||||
t = math.max(t, 1e-7); t = math.min(t, 1 - 1e-7)
|
||||
return math.log((1 - t) / t)
|
||||
end
|
||||
|
||||
local function expm1(x) return math.exp(x) - 1 end
|
||||
|
||||
local function solve_1x1(R, b) return {b[1] / (math.abs(R[1]) > 1e-12 and R[1] or 1)} end
|
||||
|
||||
local function solve_2x2(R, b)
|
||||
local det = R[1]*R[4] - R[2]*R[3]
|
||||
if math.abs(det) < 1e-12 then return {0, 0} end
|
||||
local inv = 1 / det
|
||||
return {(R[4]*b[1] - R[2]*b[2]) * inv, (R[1]*b[2] - R[3]*b[1]) * inv}
|
||||
end
|
||||
|
||||
local function solve(K, R, b)
|
||||
if K == 1 then return solve_1x1(R, b)
|
||||
elseif K == 2 then return solve_2x2(R, b)
|
||||
else return {0} end
|
||||
end
|
||||
|
||||
local function bh1_update(xt, vt, t_curr, t_next, n, model_fn, vt_buf, use_corrector)
|
||||
-- Data prediction: D_n = x - t * v
|
||||
local D_n = {}
|
||||
for i = 0, n-1 do D_n[i] = xt[i] - t_curr * vt[i] end
|
||||
|
||||
local lam_curr = lambda(t_curr)
|
||||
local lam_next = lambda(t_next)
|
||||
local h = lam_next - lam_curr
|
||||
local alpha_next = 1 - t_next
|
||||
local sigma_next = t_next
|
||||
local sigma_curr = math.max(t_curr, 1e-7)
|
||||
local hh = -h
|
||||
local B_h = hh
|
||||
local h_phi_1 = expm1(hh)
|
||||
|
||||
local avail = #history
|
||||
local order = math.min(max_order, avail + 1)
|
||||
local K = order
|
||||
local n_D1 = order - 1
|
||||
|
||||
-- h_phi_k sequence
|
||||
local h_phi_k_vals = {h_phi_1}
|
||||
local fact = 1; local hpk = h_phi_1
|
||||
for k = 1, order do
|
||||
hpk = hpk / hh - 1 / fact
|
||||
h_phi_k_vals[k+1] = hpk
|
||||
fact = fact * (k + 1)
|
||||
end
|
||||
|
||||
-- rks, R matrix, b vector
|
||||
local rks = {}
|
||||
for i = 1, n_D1 do
|
||||
local hist_idx = avail - i + 1
|
||||
local lam_hist = lambda(history[hist_idx].t)
|
||||
rks[i] = (lam_hist - lam_curr) / h
|
||||
end
|
||||
rks[n_D1 + 1] = 1
|
||||
|
||||
local R_mat = {}
|
||||
for row = 1, K do
|
||||
for col = 1, K do
|
||||
R_mat[(row-1)*K + col] = rks[col] ^ (row - 1)
|
||||
end
|
||||
end
|
||||
|
||||
local b_vec = {}
|
||||
fact = 1; hpk = h_phi_1
|
||||
for i = 1, K do
|
||||
hpk = hpk / hh - 1 / fact
|
||||
b_vec[i] = hpk * fact / B_h
|
||||
fact = fact * (i + 1)
|
||||
end
|
||||
|
||||
-- D1 differences
|
||||
local d1 = {}
|
||||
for i = 1, n_D1 do
|
||||
local hist_idx = avail - i + 1
|
||||
local D_hist = history[hist_idx].model_output
|
||||
local rk_inv = (math.abs(rks[i]) > 1e-12) and (1 / rks[i]) or 0
|
||||
d1[i] = {}
|
||||
for j = 0, n-1 do d1[i][j] = (D_hist[j] - D_n[j]) * rk_inv end
|
||||
end
|
||||
|
||||
-- Base term
|
||||
local sigma_ratio = (math.abs(sigma_curr) > 1e-7) and (sigma_next / sigma_curr) or 0
|
||||
local x_t_ = {}
|
||||
for i = 0, n-1 do
|
||||
x_t_[i] = sigma_ratio * xt[i] - alpha_next * h_phi_1 * D_n[i]
|
||||
end
|
||||
|
||||
-- Predictor
|
||||
if n_D1 > 0 then
|
||||
local rhos_p
|
||||
if order == 2 then rhos_p = {0.5}
|
||||
else
|
||||
local Kp = K - 1; local R_p = {}
|
||||
for row = 1, Kp do for col = 1, Kp do R_p[(row-1)*Kp+col] = R_mat[(row-1)*K+col] end end
|
||||
rhos_p = solve(Kp, R_p, b_vec)
|
||||
end
|
||||
for i = 0, n-1 do
|
||||
local pred = 0
|
||||
for k = 1, n_D1 do pred = pred + rhos_p[k] * d1[k][i] end
|
||||
xt[i] = x_t_[i] - alpha_next * B_h * pred
|
||||
end
|
||||
else
|
||||
for i = 0, n-1 do xt[i] = x_t_[i] end
|
||||
end
|
||||
|
||||
-- Corrector
|
||||
if use_corrector and model_fn then
|
||||
model_fn(xt, t_next)
|
||||
local D_corr_diff = {}
|
||||
for i = 0, n-1 do
|
||||
local D_corr = xt[i] - t_next * vt_buf[i]
|
||||
D_corr_diff[i] = D_corr - D_n[i]
|
||||
end
|
||||
local rhos_c
|
||||
if order == 1 then rhos_c = {0.5}
|
||||
else rhos_c = solve(K, R_mat, b_vec) end
|
||||
|
||||
for i = 0, n-1 do
|
||||
local corr = 0
|
||||
for k = 1, n_D1 do corr = corr + rhos_c[k] * d1[k][i] end
|
||||
corr = corr + rhos_c[K] * D_corr_diff[i]
|
||||
xt[i] = x_t_[i] - alpha_next * B_h * corr
|
||||
end
|
||||
end
|
||||
|
||||
-- Update history
|
||||
table.insert(history, {model_output = D_n, t = t_curr})
|
||||
while #history > max_order do table.remove(history, 1) end
|
||||
end
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
|
||||
-- Reset state on first step of a new generation
|
||||
if (step_index or 0) == 0 then history = {} end
|
||||
|
||||
bh1_update(xt, vt, t_curr, t_prev, n, model_fn, vt_buf, true)
|
||||
end
|
||||
@@ -0,0 +1,89 @@
|
||||
-- unipc_p.lua: UniPC Predictor only (1 NFE, stateful)
|
||||
-- Same as UniPC but without the corrector step.
|
||||
|
||||
solver = {
|
||||
name = "unipc_p",
|
||||
display = "UniPC Predictor (1 NFE)",
|
||||
description = "UniPC predictor-only (no corrector, 1 NFE)",
|
||||
nfe = 1,
|
||||
order = 2,
|
||||
needs_model = false,
|
||||
stateful = true,
|
||||
stochastic = false,
|
||||
}
|
||||
|
||||
-- Shares the same logic as unipc.lua but with use_corrector = false
|
||||
-- For brevity, we duplicate the core with corrector disabled.
|
||||
|
||||
local history = {}
|
||||
local max_order = 2
|
||||
|
||||
local function lambda(t)
|
||||
t = math.max(t, 1e-7); t = math.min(t, 1 - 1e-7)
|
||||
return math.log((1 - t) / t)
|
||||
end
|
||||
|
||||
local function expm1(x) return math.exp(x) - 1 end
|
||||
|
||||
local function solve_1x1(R, b) return {b[1] / (math.abs(R[1]) > 1e-12 and R[1] or 1)} end
|
||||
local function solve_2x2(R, b)
|
||||
local det = R[1]*R[4] - R[2]*R[3]
|
||||
if math.abs(det) < 1e-12 then return {0, 0} end
|
||||
local inv = 1 / det
|
||||
return {(R[4]*b[1] - R[2]*b[2]) * inv, (R[1]*b[2] - R[3]*b[1]) * inv}
|
||||
end
|
||||
local function solve(K, R, b)
|
||||
if K == 1 then return solve_1x1(R, b) else return solve_2x2(R, b) end
|
||||
end
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
-- Reset state on first step of a new generation
|
||||
if (step_index or 0) == 0 then history = {} end
|
||||
|
||||
local D_n = {}
|
||||
for i = 0, n-1 do D_n[i] = xt[i] - t_curr * vt[i] end
|
||||
|
||||
local lam_curr = lambda(t_curr)
|
||||
local lam_next = lambda(t_prev)
|
||||
local h = lam_next - lam_curr
|
||||
local alpha_next = 1 - t_prev
|
||||
local sigma_next = t_prev
|
||||
local sigma_curr = math.max(t_curr, 1e-7)
|
||||
local hh = -h
|
||||
local h_phi_1 = expm1(hh)
|
||||
|
||||
local avail = #history
|
||||
local order = math.min(max_order, avail + 1)
|
||||
local n_D1 = order - 1
|
||||
|
||||
local rks = {}
|
||||
for i = 1, n_D1 do
|
||||
local hist_idx = avail - i + 1
|
||||
rks[i] = (lambda(history[hist_idx].t) - lam_curr) / h
|
||||
end
|
||||
|
||||
local d1 = {}
|
||||
for i = 1, n_D1 do
|
||||
local D_hist = history[avail - i + 1].model_output
|
||||
local rk_inv = (math.abs(rks[i]) > 1e-12) and (1 / rks[i]) or 0
|
||||
d1[i] = {}
|
||||
for j = 0, n-1 do d1[i][j] = (D_hist[j] - D_n[j]) * rk_inv end
|
||||
end
|
||||
|
||||
local sigma_ratio = (math.abs(sigma_curr) > 1e-7) and (sigma_next / sigma_curr) or 0
|
||||
for i = 0, n-1 do
|
||||
xt[i] = sigma_ratio * xt[i] - alpha_next * h_phi_1 * D_n[i]
|
||||
end
|
||||
|
||||
if n_D1 > 0 then
|
||||
local rhos_p = (order == 2) and {0.5} or solve(n_D1, {1}, {0.5})
|
||||
for i = 0, n-1 do
|
||||
local pred = 0
|
||||
for k = 1, n_D1 do pred = pred + rhos_p[k] * d1[k][i] end
|
||||
xt[i] = xt[i] - alpha_next * hh * pred
|
||||
end
|
||||
end
|
||||
|
||||
table.insert(history, {model_output = D_n, t = t_curr})
|
||||
while #history > max_order do table.remove(history, 1) end
|
||||
end
|
||||
Reference in New Issue
Block a user