Initial release

engine/plugins/solvers/jkass_fast.lua

@@ -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