Files
hot-step-cpp-ROCm/plugins/schedulers/md_ht_scheduler_V3.lua
T
2026-08-16 18:24:52 +07:00

393 lines
18 KiB
Lua

-- ============================================================================
-- 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 HT Scheduler v5.0 — HAP + TPT Thermodynamic Timestep Schedule
-- MDMAchine | A&E Concepts © 2026
--
-- Plugin version: V3 (HOT-Step UI / filename — what users see)
-- Internal version: v5.0 (math/changelog — what developers track)
-- These are different: plugin version bumps on breaking changes or major
-- feature drops. Internal version bumps on any code change.
--
-- Sub-index CDF interpolation, parameter caching, post-CDF smoothing,
-- density floor, LINA warp, SNR-space mode, poly slope, uniformity blend,
-- verbose step-size diagnostics.
--
-- WHAT THIS DOES:
-- Standard schedulers space timesteps linearly or with a simple power curve.
-- HT uses two coupled density functions to place steps where they matter:
--
-- HAP (Hamiltonian Action-Principle):
-- density_hap = 1 / ((1 + KE * t) * exp(-DF * t))
-- High KE = front-loaded steps (aggressive early denoising)
-- High DF = fast exponential damping toward formation
--
-- TPT (Thermodynamic Phase Transition):
-- density_tpt = 1 / (|sigma - Tc| + well_width)
-- Creates a "gravity well" that clusters steps near the critical temp
-- where latent structure crystallizes.
--
-- Combined: density = density_hap + phase_intensity * density_tpt + floor
-- Result: non-uniform sigma sequence. Works at any step count (12-150+).
--
-- POST-PROCESSING CHAIN (all optional, all default off):
-- 1. Shift warp (native HOT-Step sigma warp)
-- 2. LINA warp (time-axis resampling from MD Causal)
-- 3. Poly slope (power curve on sigma values)
-- 4. Uniformity blend (blend with linear uniform schedule)
-- 5. Schedule smoothing (moving average on final sigmas)
--
-- CHANGELOG:
-- v5.0: Density floor, LINA warp, SNR-space mode, poly slope, uniformity
-- blend, post-CDF smoothing, verbose diagnostics. Additive HAP+TPT
-- blending. Exposed well width. Restored descriptive header. Complete
-- rework from v4.0 baseline.
-- ============================================================================
scheduler = {
name = "md_ht_scheduler V3",
display = "MD HT Scheduler (HAP+TPT) V3",
description = "Thermodynamic timestep schedule: HAP + TPT additive density, density floor, LINA warp, SNR-space, poly slope, uniformity blend, smoothing. 12 to 150+ steps.",
params = {
-- ── HAP ─────────────────────────────────────────────────────────────
{ key = "kinetic_energy", type = "slider", label = "Kinetic Energy",
default = 0.3, min = 0.0, max = 3.0, step = 0.05,
hint = "HAP leading-edge sharpness. 0=uniform, 0.3=standard, 2+=aggressive front-loading." },
{ key = "damping_friction", type = "slider", label = "Damping Friction",
default = 2.2, min = 0.0, max = 6.0, step = 0.1,
hint = "HAP tail compression. Higher = steps cluster toward the front." },
-- ── TPT ─────────────────────────────────────────────────────────────
{ key = "critical_temp", type = "slider", label = "Critical Temp",
default = 0.6, min = 0.05, max = 0.95, step = 0.05,
hint = "TPT phase transition center (sigma fraction). Steps cluster here." },
{ key = "phase_intensity", type = "slider", label = "Phase Intensity",
default = 1.0, min = 0.0, max = 3.0, step = 0.1,
hint = "TPT clustering strength. 0=off (pure HAP). 1=moderate. 2+=strong." },
{ key = "well_width", type = "slider", label = "Well Width",
default = 0.25, min = 0.05, max = 0.5, step = 0.05,
hint = "TPT softening radius. 0.1=tight. 0.25=balanced. 0.4+=broad." },
-- ── Density Floor ───────────────────────────────────────────────────
{ key = "density_floor", type = "slider", label = "Density Floor",
default = 0.1, min = 0.0, max = 1.0, step = 0.05,
hint = "Minimum density everywhere. Prevents sparse gaps. 0=off. 0.1=gentle. 0.3+=uniform-leaning." },
-- ── SNR Space ───────────────────────────────────────────────────────
{ key = "snr_space", type = "toggle", label = "SNR Space",
default = false,
hint = "Compute density on an SNR-uniform grid instead of sigma-uniform. Steps track perceptual importance. Better for audio at high step counts." },
-- ── Post-Processing ─────────────────────────────────────────────────
{ key = "lina_shift", type = "slider", label = "LINA Warp",
default = 1.0, min = 0.5, max = 2.0, step = 0.05,
hint = "Time-axis resampling (from MD Causal). 1.0=off. <1=front-load (more high-sigma steps). >1=back-load (more low-sigma steps). Different from shift warp — this resamples WHERE on the curve, not the sigma VALUES." },
{ key = "poly_slope", type = "slider", label = "Poly Slope",
default = 1.0, min = 0.5, max = 2.0, step = 0.05,
hint = "Power curve on sigma values. 1.0=off. >1=compress toward zero (more detail steps, good for long runs). <1=compress toward one (more structure steps, good for 12-step turbo)." },
{ key = "uniform_blend", type = "slider", label = "Uniformity Blend",
default = 0.0, min = 0.0, max = 1.0, step = 0.05,
hint = "Blend with linear uniform schedule. 0=pure HT. 0.3=gentle uniformity. 1.0=pure uniform. Tames HT clustering for stabilization solvers (Trajectory Anchor)." },
{ key = "smooth_window", type = "slider", label = "Schedule Smoothing",
default = 0, min = 0, max = 7, step = 1,
hint = "Post-CDF moving average on final sigmas. 0=off. 3=mild. 5+=heavy. Smooths step-size transitions." },
-- ── Engine ──────────────────────────────────────────────────────────
{ key = "dense_steps", type = "slider", label = "CDF Resolution",
default = 1000, min = 200, max = 5000, step = 100,
hint = "Resolution of the integration grid." },
{ key = "shift", type = "slider", label = "Shift Warp",
default = 1.0, min = 0.5, max = 8.0, step = 0.1,
hint = "Native HOT-Step sigma warp. Applied first in the post-processing chain." },
{ key = "verbose", type = "toggle", label = "Verbose",
default = false,
hint = "Print per-step sigma values, step sizes, and gap ratio to console." },
},
}
-- ── Hoisted Buffers & Cache ──────────────────────────────────────────────────
local EPSILON = 1e-6
local _cache = { ke = -1, df = -1, tc = -1, pi = -1, ww = -1, fl = -1, snr = -1, dense_n = -1 }
local _dense = {}
local _cdf = {}
local _sigmas = {}
local _smooth_buf = {}
local _last_num_steps = -1
local function clamp(v, lo, hi)
if v < lo then return lo end
if v > hi then return hi end
return v
end
-- ── SNR-Space Grid ──────────────────────────────────────────────────────────
-- Build dense grid uniform in SNR space: snr = log(sigma / (1 - sigma)).
-- Maps back to sigma via sigmoid: sigma = 1 / (1 + exp(-snr)).
-- Endpoints clamped to avoid inf at sigma=0 and sigma=1.
local function build_snr_grid(dense, dense_n, sigma_max, sigma_min)
local s_hi = clamp(sigma_max, 0.001, 0.999)
local s_lo = clamp(sigma_min + 0.001, 0.001, 0.999)
local snr_hi = math.log(s_hi / (1.0 - s_hi))
local snr_lo = math.log(s_lo / (1.0 - s_lo))
for i = 0, dense_n - 1 do
local snr = snr_hi + (snr_lo - snr_hi) * i / (dense_n - 1)
dense[i] = 1.0 / (1.0 + math.exp(-snr))
end
dense[0] = sigma_max
dense[dense_n - 1] = sigma_min
end
-- ── LINA Warp (from MD Causal) ──────────────────────────────────────────────
-- Time-axis resampling: t_warped = t^shift, then interpolate into the sigma
-- array at the warped position. Shift < 1 front-loads, shift > 1 back-loads.
-- Different from native shift warp which transforms sigma values directly.
local function apply_lina_warp(sigmas, n, shift)
if shift == 1.0 then return end
-- Read into scratch buffer first
for i = 0, n do _smooth_buf[i] = sigmas[i] end
for i = 0, n do
local t = i / n
local t_w = t ^ shift
local raw_idx = t_w * n
local idx_lo = clamp(math.floor(raw_idx), 0, n)
local idx_hi = clamp(idx_lo + 1, 0, n)
local frac = raw_idx - idx_lo
local s_lo = _smooth_buf[idx_lo]
local s_hi = _smooth_buf[idx_hi] or _smooth_buf[n]
sigmas[i] = s_lo * (1.0 - frac) + s_hi * frac
end
sigmas[0] = _smooth_buf[0]
sigmas[n] = _smooth_buf[n]
end
-- ── Post-CDF Schedule Smoothing ─────────────────────────────────────────────
local function smooth_schedule(sigmas, n, window)
if window < 2 or n < window then return end
local half = math.floor(window / 2)
for i = 1, n - 1 do
local sum = 0.0
local count = 0
for j = math.max(0, i - half), math.min(n, i + half) do
sum = sum + sigmas[j]
count = count + 1
end
_smooth_buf[i] = sum / count
end
for i = 1, n - 1 do
sigmas[i] = _smooth_buf[i]
end
-- Enforce monotonically decreasing
for i = 1, n do
if sigmas[i] >= sigmas[i - 1] then
sigmas[i] = sigmas[i - 1] - EPSILON
end
end
end
-- ── Core Schedule Builder ────────────────────────────────────────────────────
local function build_ht_schedule(num_steps, ke, df, tc_frac, pi, ww, fl, snr_mode, dense_n)
local sigma_max = 1.0
local sigma_min = 0.0
local snr_flag = snr_mode and 1 or 0
-- Only rebuild CDF if params changed
if ke ~= _cache.ke or df ~= _cache.df or tc_frac ~= _cache.tc
or pi ~= _cache.pi or ww ~= _cache.ww or fl ~= _cache.fl
or snr_flag ~= _cache.snr or dense_n ~= _cache.dense_n then
-- Build dense grid (sigma-uniform or SNR-uniform)
if snr_mode then
build_snr_grid(_dense, dense_n, sigma_max, sigma_min)
else
for i = 0, dense_n - 1 do
_dense[i] = sigma_max - (sigma_max - sigma_min) * i / (dense_n - 1)
end
end
local critical_temp = sigma_min + tc_frac * (sigma_max - sigma_min)
local running = 0.0
for i = 0, dense_n - 1 do
local s = _dense[i]
local t = (sigma_max - s) / (sigma_max - sigma_min + EPSILON)
-- HAP density
local v_hap = (1.0 + ke * t) * math.exp(-df * t)
local d_hap = 1.0 / (v_hap + EPSILON)
-- TPT density
local dist_tc = math.abs(s - critical_temp)
local d_tpt = 1.0 / (dist_tc + ww)
-- Additive blending + density floor
running = running + d_hap + pi * d_tpt + fl
_cdf[i] = running
end
-- Normalize CDF to [0, 1]
local cdf0 = _cdf[0]
local cdf_n1 = _cdf[dense_n - 1]
local range = cdf_n1 - cdf0 + EPSILON
for i = 0, dense_n - 1 do
_cdf[i] = (_cdf[i] - cdf0) / range
end
_cache.ke = ke
_cache.df = df
_cache.tc = tc_frac
_cache.pi = pi
_cache.ww = ww
_cache.fl = fl
_cache.snr = snr_flag
_cache.dense_n = dense_n
end
-- Pre-allocate
if num_steps > _last_num_steps then
for i = 0, num_steps do _sigmas[i] = 0.0 end
for i = 0, num_steps do _smooth_buf[i] = 0.0 end
_last_num_steps = num_steps
end
-- Binary search
local function searchsorted(target)
local lo, hi = 0, dense_n - 1
while lo < hi do
local mid = math.floor((lo + hi) / 2)
if _cdf[mid] < target then lo = mid + 1 else hi = mid end
end
return clamp(lo, 0, dense_n - 1)
end
-- Sub-index interpolation
for i = 0, num_steps do
local tgt = i / num_steps
local idx = searchsorted(tgt)
if idx == 0 then
_sigmas[i] = _dense[0]
else
local c0 = _cdf[idx - 1]
local c1 = _cdf[idx]
local t_interp = (c1 > c0) and ((tgt - c0) / (c1 - c0)) or 0.0
local d0 = _dense[idx - 1]
local d1 = _dense[idx]
_sigmas[i] = d0 + t_interp * (d1 - d0)
end
end
-- Force exact endpoints
_sigmas[0] = sigma_max
_sigmas[num_steps] = sigma_min
return _sigmas
end
-- ── Required schedule() function ─────────────────────────────────────────────
function schedule(output, num_steps, shift_val)
local ke = (params and params.kinetic_energy) or 0.3
local df = (params and params.damping_friction) or 2.2
local tc_frac = (params and params.critical_temp) or 0.6
local pi_ = (params and params.phase_intensity) or 1.0
local ww = (params and params.well_width) or 0.25
local fl = (params and params.density_floor) or 0.1
local snr_mode = (params and params.snr_space) or false
local lina = (params and params.lina_shift) or 1.0
local poly = (params and params.poly_slope) or 1.0
local u_blend = (params and params.uniform_blend) or 0.0
local sm_win = math.floor((params and params.smooth_window) or 0)
local dense_n = math.floor((params and params.dense_steps) or 1000)
local sh = (params and params.shift) or shift_val
local verbose = (params and params.verbose) or false
local sigmas = build_ht_schedule(num_steps, ke, df, tc_frac, pi_, ww, fl, snr_mode, dense_n)
-- ── POST-PROCESSING CHAIN ───────────────────────────────────────────
-- Order: shift → LINA → poly → uniformity → smooth
-- Each is independent and default-off. Compose cleanly at any step count.
-- 1. Native shift warp (sigma-value transform)
if sh ~= 1.0 then
for i = 0, num_steps do
local t = sigmas[i]
sigmas[i] = sh * t / (1.0 + (sh - 1.0) * t)
end
end
-- 2. LINA warp (time-axis resampling)
if lina ~= 1.0 then
apply_lina_warp(sigmas, num_steps, lina)
end
-- 3. Poly slope (power curve on sigma values)
-- >1 = compress toward zero (more detail steps)
-- <1 = compress toward one (more structure steps, good for 12-step turbo)
if poly ~= 1.0 then
for i = 1, num_steps - 1 do
sigmas[i] = sigmas[i] ^ poly
end
-- Endpoints stay exact
end
-- 4. Uniformity blend (blend with linear schedule)
-- Tames HT clustering for stabilization solvers
if u_blend > 0.0 then
local inv = 1.0 - u_blend
for i = 0, num_steps do
local uniform_sigma = 1.0 - (i / num_steps)
sigmas[i] = inv * sigmas[i] + u_blend * uniform_sigma
end
end
-- 5. Schedule smoothing (moving average)
if sm_win >= 2 then
smooth_schedule(sigmas, num_steps, sm_win)
sigmas[0] = 1.0
sigmas[num_steps] = 0.0
end
-- ── WRITE OUTPUT (no trailing zero — engine contract) ───────────────
for i = 0, num_steps - 1 do
output[i] = sigmas[i]
end
-- ── VERBOSE ─────────────────────────────────────────────────────────
if verbose then
print(string.format(
"[HT V5] %d steps | ke=%.2f df=%.1f tc=%.2f pi=%.1f ww=%.2f fl=%.2f | snr=%s lina=%.2f poly=%.2f ub=%.2f sm=%d",
num_steps, ke, df, tc_frac, pi_, ww, fl,
snr_mode and "on" or "off", lina, poly, u_blend, sm_win))
local min_gap, max_gap = 1.0, 0.0
for i = 0, num_steps - 1 do
local sigma_next = (i < num_steps - 1) and sigmas[i + 1] or 0.0
local gap = sigmas[i] - sigma_next
if gap < min_gap then min_gap = gap end
if gap > max_gap then max_gap = gap end
print(string.format(" step %02d: sigma=%.5f gap=%.5f", i, sigmas[i], gap))
end
print(string.format("[HT V5] Gap range: min=%.5f max=%.5f ratio=%.1f:1",
min_gap, max_gap, max_gap / (min_gap + EPSILON)))
end
end