112 lines
3.7 KiB
Lua
112 lines
3.7 KiB
Lua
-- 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
|