Initial release
This commit is contained in:
civ
@@ -0,0 +1,360 @@
|
||||
-- ============================================================================
|
||||
-- SPDX-License-Identifier: GPL-3.0-or-later
|
||||
-- Copyright (C) 2026 Alexander Allan (MDMAchine) -- A&E Concepts
|
||||
--
|
||||
-- This program is free software: you can redistribute it and/or modify
|
||||
-- it under the terms of the GNU General Public License as published by
|
||||
-- the Free Software Foundation, either version 3 of the License, or
|
||||
-- (at your option) any later version.
|
||||
--
|
||||
-- This program is distributed in the hope that it will be useful,
|
||||
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
-- GNU General Public License for more details: https://www.gnu.org/licenses/
|
||||
-- ============================================================================
|
||||
|
||||
-- MD PingPong Simple v1.1 — Ancestral Euler + Momentum + Look-Back Smoother
|
||||
-- MDMAchine | A&E Concepts © 2026
|
||||
--
|
||||
-- Port of MD_PingPong_Samplers.py (single-branch path) to HOT-Step-CPP Lua solver.
|
||||
--
|
||||
-- WHAT THIS DOES:
|
||||
-- Standard ODE solvers (Euler, Heun, STORM) advance the latent deterministically.
|
||||
-- PingPong injects ancestral (stochastic) noise at each step — the "ping" is
|
||||
-- the clean denoising step, the "pong" is the noise re-injection that keeps
|
||||
-- the trajectory alive and stochastic.
|
||||
--
|
||||
-- CORE STEP MATH:
|
||||
-- dt = t_prev - t_curr (negative — stepping down)
|
||||
-- x_euler = xt + dt * vt (standard Euler, matches STORM/OmniRelational)
|
||||
-- x_next = x_euler + noise * |dt| * ancestral_strength
|
||||
--
|
||||
-- MOMENTUM:
|
||||
-- Latent velocity (x - x_prev) carried forward at each step. Maintains
|
||||
-- flow continuity across the ODE trajectory — reduces erratic jumps between
|
||||
-- steps, especially at low step counts.
|
||||
--
|
||||
-- NOISE COHERENCE:
|
||||
-- Blends fresh Gaussian noise with the previous step's noise at ratio
|
||||
-- noise_coherence. 0=fully fresh, 1=fully carried. Useful for temporal
|
||||
-- smoothness in audio; keep low (0-0.2) to avoid spectral smearing.
|
||||
--
|
||||
-- LOOK-BACK SNR SMOOTHER:
|
||||
-- λ(σ) = lambda_base * (σ/σ_max)^snr_power — heavy at high sigma, fades to
|
||||
-- zero at low sigma. Suppresses ODE manifold shearing and harmonic hum.
|
||||
-- Reference: arXiv:2602.09449
|
||||
--
|
||||
-- RMS SERVO:
|
||||
-- Downward-only energy ceiling that follows a smooth curve from rms_max
|
||||
-- (high sigma) to rms_min (low sigma). Domain-tunable: image latents
|
||||
-- typically sit around 0.75-0.97; audio latents around 0.3-0.7.
|
||||
-- Servo is DOWNWARD ONLY — never boosts energy, only clamps excess.
|
||||
--
|
||||
-- SOLVER API NOTE:
|
||||
-- HOT-Step-CPP passes velocity vt where dt = t_prev - t_curr is NEGATIVE
|
||||
-- (stepping from high sigma to low sigma). Euler update is:
|
||||
-- x_next = xt + dt * vt (same convention as STORM and OmniRelational)
|
||||
-- Ancestral noise is added as: noise * |dt| * strength
|
||||
-- Do NOT use xt - t_curr * vt for denoised — sign convention mismatch.
|
||||
--
|
||||
-- PARAMS:
|
||||
-- ancestral_strength — noise injection scale. 1.0=standard ancestral, 0=pure ODE
|
||||
-- noise_coherence — step-to-step noise carry. 0=fresh, 0.2=subtle temporal link
|
||||
-- momentum_strength — latent velocity carry-over. 0.15=subtle, 0.3=strong
|
||||
-- look_back_enabled — SNR smoother toggle
|
||||
-- look_back_lambda — max smoothing weight (0.55=25-step, 0.35=35-step)
|
||||
-- look_back_snr_power — falloff exponent (1.3=25-step, 1.5=35-step)
|
||||
-- rms_servo — energy ceiling toggle
|
||||
-- rms_target_min — servo floor at low sigma (audio: ~0.3, image: ~0.75)
|
||||
-- rms_target_max — servo ceiling at high sigma (audio: ~0.7, image: ~0.97)
|
||||
-- rms_servo_gain — correction aggressiveness (0.6=default, 1.0=hard snap)
|
||||
-- seed — RNG seed
|
||||
-- ============================================================================
|
||||
|
||||
solver = {
|
||||
name = "md_pingpong_simple",
|
||||
display = "MD PingPong Simple (Ancestral)",
|
||||
description = "Ancestral Euler with momentum, noise coherence, look-back SNR smoother, and domain-tunable RMS servo. Single-branch stochastic sampler. Port of MD_PingPong_Samplers v3.5.",
|
||||
nfe = 1,
|
||||
order = 1,
|
||||
needs_model = false,
|
||||
stateful = true,
|
||||
stochastic = true,
|
||||
params = {
|
||||
{
|
||||
key = "ancestral_strength",
|
||||
type = "slider",
|
||||
label = "Ancestral Strength",
|
||||
default = 0.2,
|
||||
min = 0.0,
|
||||
max = 1.5,
|
||||
step = 0.05,
|
||||
hint = "Noise injection strength. 1.0=standard ancestral. 0=pure ODE (no noise).",
|
||||
},
|
||||
{
|
||||
key = "noise_coherence",
|
||||
type = "slider",
|
||||
label = "Noise Coherence",
|
||||
default = 0.0,
|
||||
min = 0.0,
|
||||
max = 1.0,
|
||||
step = 0.05,
|
||||
hint = "Step-to-step noise correlation. 0=fresh noise each step. 0.2=subtle temporal link. Keep low for audio to avoid smearing.",
|
||||
},
|
||||
{
|
||||
key = "momentum_strength",
|
||||
type = "slider",
|
||||
label = "Momentum",
|
||||
default = 0.1,
|
||||
min = 0.0,
|
||||
max = 0.5,
|
||||
step = 0.01,
|
||||
hint = "Latent velocity carry-over. 0.1=subtle flow continuity. 0.3=strong.",
|
||||
},
|
||||
{
|
||||
key = "look_back_enabled",
|
||||
type = "toggle",
|
||||
label = "Look-Back Smoother",
|
||||
default = true,
|
||||
hint = "SNR-adaptive latent EMA. Suppresses ODE manifold shearing and harmonic hum. arXiv:2602.09449.",
|
||||
},
|
||||
{
|
||||
key = "look_back_lambda",
|
||||
type = "slider",
|
||||
label = "Look-Back Lambda",
|
||||
default = 0.55,
|
||||
min = 0.1,
|
||||
max = 1.0,
|
||||
step = 0.05,
|
||||
hint = "Max smoothing weight. Active when Look-Back Smoother is on. 0.55=25-step, 0.35=35-step.",
|
||||
},
|
||||
{
|
||||
key = "look_back_snr_power",
|
||||
type = "slider",
|
||||
label = "SNR Power",
|
||||
default = 1.3,
|
||||
min = 0.5,
|
||||
max = 3.0,
|
||||
step = 0.1,
|
||||
hint = "Falloff exponent. Active when Look-Back Smoother is on. Higher=smoother fade at low sigma.",
|
||||
},
|
||||
{
|
||||
key = "rms_servo",
|
||||
type = "toggle",
|
||||
label = "RMS Servo",
|
||||
default = true,
|
||||
hint = "Downward-only energy ceiling. Prevents latent energy accumulation. Off by default — tune min/max for your domain before enabling.",
|
||||
},
|
||||
{
|
||||
key = "rms_target_min",
|
||||
type = "slider",
|
||||
label = "RMS Target Min",
|
||||
default = 1.0,
|
||||
min = 0.1,
|
||||
max = 3.0,
|
||||
step = 0.05,
|
||||
hint = "RMS ceiling at low sigma (late steps). Active when RMS Servo is on. ACE-Step latents ~2.0 RMS. Start at 1.2-1.8.",
|
||||
},
|
||||
{
|
||||
key = "rms_target_max",
|
||||
type = "slider",
|
||||
label = "RMS Target Max",
|
||||
default = 2.2,
|
||||
min = 0.5,
|
||||
max = 4.0,
|
||||
step = 0.05,
|
||||
hint = "RMS ceiling at high sigma (early steps). Active when RMS Servo is on. ACE-Step latents ~2.0 RMS. Start at 2.0-2.5.",
|
||||
},
|
||||
{
|
||||
key = "rms_servo_gain",
|
||||
type = "slider",
|
||||
label = "Servo Gain",
|
||||
default = 0.75,
|
||||
min = 0.1,
|
||||
max = 1.0,
|
||||
step = 0.05,
|
||||
hint = "Servo correction aggressiveness. Active when RMS Servo is on. 0.6=soft, 1.0=hard snap.",
|
||||
},
|
||||
{
|
||||
key = "seed",
|
||||
type = "slider",
|
||||
label = "Seed",
|
||||
default = 42,
|
||||
min = 0,
|
||||
max = 999999,
|
||||
step = 1,
|
||||
hint = "RNG seed for noise generation.",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
-- ── State (file-level locals, reset when n changes) ──────────────────────────
|
||||
|
||||
local _prev_x = nil -- for momentum
|
||||
local _prev_noise = nil -- for noise coherence
|
||||
local _look_back_xp = nil -- for look-back smoother
|
||||
local _sigma_max = nil -- captured at step 0
|
||||
local _last_n = 0
|
||||
|
||||
-- Hoisted scratch tables — reused every step to avoid GC pressure
|
||||
-- Initialized on first step or when n changes
|
||||
local _noise_buf = {} -- reusable noise array
|
||||
local _x_next_buf = {} -- reusable output array
|
||||
local _x_copy_buf = {} -- reusable momentum copy
|
||||
|
||||
local EPSILON = 1e-8
|
||||
|
||||
-- ── Helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
local function clamp(v, lo, hi)
|
||||
if v < lo then return lo end
|
||||
if v > hi then return hi end
|
||||
return v
|
||||
end
|
||||
|
||||
-- Seeded LCG RNG — deterministic, no dependency on math.random state
|
||||
local function make_rng(seed)
|
||||
local state = math.floor(seed) % 2147483647
|
||||
if state <= 0 then state = state + 2147483646 end
|
||||
return function()
|
||||
state = (state * 1664525 + 1013904223) % 2147483648
|
||||
return state / 2147483648.0
|
||||
end
|
||||
end
|
||||
|
||||
-- Box-Muller: two uniform [0,1] → one standard normal sample
|
||||
local function normal(u1, u2)
|
||||
return math.sqrt(-2.0 * math.log(math.max(u1, EPSILON))) * math.cos(2.0 * math.pi * u2)
|
||||
end
|
||||
|
||||
-- Array RMS
|
||||
local function rms(arr, n)
|
||||
local s = 0.0
|
||||
for i = 0, n - 1 do s = s + arr[i] * arr[i] end
|
||||
return math.sqrt(s / n + EPSILON)
|
||||
end
|
||||
|
||||
-- ── Required step() function ──────────────────────────────────────────────────
|
||||
|
||||
function step(xt, vt, t_curr, t_prev, n)
|
||||
local step_idx_ = step_index or 0
|
||||
|
||||
-- Reset state on new generation: n change OR step 0 of any new run.
|
||||
-- Must check step_idx_==0 because same-length generations won't trigger n change,
|
||||
-- causing momentum/look-back to bleed finished audio from the previous run into
|
||||
-- the noise of the new one — explosive velocity on step 1.
|
||||
if n ~= _last_n or step_idx_ == 0 then
|
||||
_prev_x = nil
|
||||
_prev_noise = nil
|
||||
_look_back_xp = nil
|
||||
_sigma_max = nil
|
||||
_last_n = n
|
||||
-- Pre-size scratch tables for this n
|
||||
for i = 0, n - 1 do
|
||||
_noise_buf[i] = 0.0
|
||||
_x_next_buf[i] = 0.0
|
||||
_x_copy_buf[i] = 0.0
|
||||
end
|
||||
end
|
||||
|
||||
-- Read params with safe fallbacks
|
||||
local anc_strength = (params and params.ancestral_strength) or 1.0
|
||||
local noise_coh = (params and params.noise_coherence) or 0.0
|
||||
local mom_str = (params and params.momentum_strength) or 0.1
|
||||
local lb_enabled = (params and params.look_back_enabled) or false
|
||||
local lb_lambda = (params and params.look_back_lambda) or 0.55
|
||||
local lb_snr_power = (params and params.look_back_snr_power) or 1.3
|
||||
local rms_servo_on = (params and params.rms_servo) or true
|
||||
local rms_tgt_min = (params and params.rms_target_min) or 1.2
|
||||
local rms_tgt_max = (params and params.rms_target_max) or 2.2
|
||||
local rms_servo_gain = (params and params.rms_servo_gain) or 0.6
|
||||
local seed = math.floor((params and params.seed) or 42)
|
||||
|
||||
-- Capture sigma_max on first step for ratio computation
|
||||
if _sigma_max == nil then _sigma_max = t_curr end
|
||||
local sigma_max = _sigma_max
|
||||
local sigma_ratio = clamp(t_curr / math.max(sigma_max, EPSILON), 0.0, 1.0)
|
||||
|
||||
-- dt = t_prev - t_curr. In flow-matching, t steps DOWN (1→0),
|
||||
-- HOT-Step API: t_curr=high sigma, t_prev=lower target. t_curr > t_prev. dt=t_prev-t_curr is NEGATIVE.
|
||||
local dt = t_prev - t_curr
|
||||
|
||||
-- Save current xt for momentum (reuse hoisted buffer)
|
||||
for i = 0, n - 1 do _x_copy_buf[i] = xt[i] end
|
||||
|
||||
-- ── NOISE GENERATION ──────────────────────────────────────────────────────
|
||||
-- Write into hoisted buffer — no table allocation per step
|
||||
local rng = make_rng(seed + step_idx_ * 7919)
|
||||
for i = 0, n - 1 do
|
||||
local u1 = math.max(rng(), EPSILON)
|
||||
local u2 = rng()
|
||||
_noise_buf[i] = normal(u1, u2)
|
||||
end
|
||||
|
||||
-- Noise coherence: blend with carried noise from previous step
|
||||
if noise_coh > 0.0 and _prev_noise ~= nil then
|
||||
for i = 0, n - 1 do
|
||||
_noise_buf[i] = _noise_buf[i] * (1.0 - noise_coh) + _prev_noise[i] * noise_coh
|
||||
end
|
||||
end
|
||||
-- Store for next step — reuse _prev_noise table if same size
|
||||
if _prev_noise == nil then _prev_noise = {} end
|
||||
for i = 0, n - 1 do _prev_noise[i] = _noise_buf[i] end
|
||||
|
||||
-- ── ANCESTRAL STEP ────────────────────────────────────────────────────────
|
||||
-- Variance-preserving SDE noise for flow matching:
|
||||
-- noise_scale = sqrt(t_prev^2 - t_curr^2) * anc_strength
|
||||
-- t_curr > t_prev, so t_curr^2 - t_prev^2 > 0. Confirmed numerically.
|
||||
local noise_scale = math.sqrt(math.max(t_curr * t_curr - t_prev * t_prev, 0.0)) * anc_strength
|
||||
if noise_scale > EPSILON then
|
||||
for i = 0, n - 1 do
|
||||
_x_next_buf[i] = xt[i] + dt * vt[i] + _noise_buf[i] * noise_scale
|
||||
end
|
||||
else
|
||||
for i = 0, n - 1 do _x_next_buf[i] = xt[i] + dt * vt[i] end
|
||||
end
|
||||
|
||||
-- ── MOMENTUM ──────────────────────────────────────────────────────────────
|
||||
if mom_str > 0.0 and _prev_x ~= nil then
|
||||
for i = 0, n - 1 do
|
||||
local vel = _x_copy_buf[i] - _prev_x[i]
|
||||
_x_next_buf[i] = _x_next_buf[i] + vel * mom_str
|
||||
end
|
||||
end
|
||||
|
||||
-- ── LOOK-BACK SNR SMOOTHER ────────────────────────────────────────────────
|
||||
-- λ(σ) = lb_lambda * (σ/σ_max)^lb_snr_power — heavy early, fades late.
|
||||
if lb_enabled then
|
||||
local lb_w = lb_lambda * (sigma_ratio ^ lb_snr_power)
|
||||
if _look_back_xp == nil then
|
||||
_look_back_xp = {}
|
||||
for i = 0, n - 1 do
|
||||
local u1 = math.max(rng(), EPSILON)
|
||||
local u2 = rng()
|
||||
_look_back_xp[i] = _x_next_buf[i] + normal(u1, u2) * sigma_max * 0.1
|
||||
end
|
||||
end
|
||||
for i = 0, n - 1 do
|
||||
_x_next_buf[i] = _x_next_buf[i] * (1.0 - lb_w) + _look_back_xp[i] * lb_w
|
||||
end
|
||||
-- Update look-back buffer in-place
|
||||
for i = 0, n - 1 do _look_back_xp[i] = _x_next_buf[i] end
|
||||
end
|
||||
|
||||
-- ── RMS SERVO (DOWNWARD ONLY) ─────────────────────────────────────────────
|
||||
if rms_servo_on then
|
||||
local rms_target = rms_tgt_min + (sigma_ratio ^ 0.6) * (rms_tgt_max - rms_tgt_min)
|
||||
local cur_rms = rms(_x_next_buf, n)
|
||||
if cur_rms > rms_target then
|
||||
local servo_rms = cur_rms + rms_servo_gain * (rms_target - cur_rms)
|
||||
local scale = servo_rms / cur_rms
|
||||
for i = 0, n - 1 do _x_next_buf[i] = _x_next_buf[i] * scale end
|
||||
end
|
||||
end
|
||||
|
||||
-- ── UPDATE STATE & WRITE OUTPUT ───────────────────────────────────────────
|
||||
-- Store momentum reference — reuse table, copy values
|
||||
if _prev_x == nil then _prev_x = {} end
|
||||
for i = 0, n - 1 do _prev_x[i] = _x_copy_buf[i] end
|
||||
for i = 0, n - 1 do xt[i] = _x_next_buf[i] end
|
||||
end
|
||||
Reference in New Issue
Block a user