# Plugin Authoring Guide

How to create custom solvers, schedulers, and guidance modes for the HOT-Step CPP engine.

---

## Quick Start

1. Create a `.lua` file in the appropriate directory:
   - `engine/plugins/solvers/` — ODE/SDE solvers
   - `engine/plugins/schedulers/` — noise schedules
   - `engine/plugins/guidance/` — CFG guidance modes
2. Declare a metadata table (`solver`, `scheduler`, or `guidance`)
3. Implement the required function (`step`, `schedule`, or `guide`)
4. Restart the app — your plugin appears in the UI automatically

No C++ rebuild required. The engine hot-loads all `.lua` files at startup.

---

## Plugin Types

### Solver

Solvers advance the latent state `xt` by one step along the ODE/SDE trajectory.

**Metadata table:**

```lua
solver = {
    name        = "my_solver",           -- unique internal ID (lowercase, underscores)
    display     = "My Solver (2 NFE)",   -- name shown in UI dropdown
    description = "A custom solver",     -- tooltip text
    nfe         = 2,                     -- number of function evaluations per step
    order       = 2,                     -- solver order (informational)
    needs_model = true,                  -- true if step() needs model_fn callback
    stateful    = false,                 -- true if solver carries state across steps
    stochastic  = false,                 -- true if solver uses randomness (SDE)
}
```

**Required function — single-eval solver:**

```lua
function step(xt, vt, t_curr, t_prev, n)
    -- xt:     mutable FloatArray — current latent state (modify in-place)
    -- vt:     read-only FloatArray — velocity at (xt, t_curr)
    -- t_curr: float — current timestep
    -- t_prev: float — next timestep (we step FROM t_curr TO t_prev)
    -- n:      int — total elements in xt/vt

    local dt = t_curr - t_prev
    for i = 0, n - 1 do
        xt[i] = xt[i] - vt[i] * dt
    end
end
```

**Required function — multi-eval solver** (when `needs_model = true`):

```lua
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
    -- Additional args when needs_model = true:
    -- model_fn(xt_array, t_val): evaluates the model at (xt_array, t_val),
    --                            writes velocity output to vt_buf
    -- vt_buf:  mutable FloatArray — receives model_fn output

    local dt = t_curr - t_prev
    local t_mid = t_curr - 0.5 * dt

    -- Save state
    local k1 = {}
    local xt_orig = {}
    for i = 0, n - 1 do
        k1[i] = vt[i]
        xt_orig[i] = xt[i]
    end

    -- Evaluate at midpoint
    for i = 0, n - 1 do
        xt[i] = xt_orig[i] - 0.5 * k1[i] * dt
    end
    model_fn(xt, t_mid)  -- result appears in vt_buf

    -- Final update using midpoint velocity
    for i = 0, n - 1 do
        xt[i] = xt_orig[i] - vt_buf[i] * dt
    end
end
```

---

### Scheduler

Schedulers produce a timestep sequence for the denoising trajectory.

**Metadata table:**

```lua
scheduler = {
    name        = "my_schedule",
    display     = "My Schedule",
    description = "Custom noise schedule",
}
```

**Required function:**

```lua
function schedule(output, num_steps, shift)
    -- output:    mutable FloatArray — write num_steps timestep values
    -- num_steps: int — number of timesteps to generate
    -- shift:     float — noise shift parameter from UI

    for i = 0, num_steps - 1 do
        output[i] = 1.0 - i / num_steps
    end
    apply_shift(output, num_steps, shift)
end
```

Timesteps go from `1.0` (pure noise) to `~0.0` (clean signal). The engine appends a final `0.0` step automatically — your schedule should produce `num_steps` values, not `num_steps + 1`.

**Common helper — shift warp:**

```lua
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
```

---

### Guidance

Guidance modes control how the conditional and unconditional model predictions are combined.

**Metadata table:**

```lua
guidance = {
    name        = "my_guidance",
    display     = "My Guidance",
    description = "Custom guidance mode",
}
```

**Required function:**

```lua
function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
    -- pred_cond:      read-only FloatArray — conditional velocity prediction
    -- pred_uncond:    read-only FloatArray — unconditional velocity prediction
    -- guidance_scale: float — the guidance scale (w) from the UI
    -- result:         mutable FloatArray — write the guided velocity here
    -- Oc:             int — output channels per timestep frame
    -- T:              int — number of timestep frames (n = Oc * T)
    -- norm_threshold: float — APG norm threshold from the UI

    -- Route through APG for stability (STRONGLY RECOMMENDED)
    apg(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
end
```

> **⚠️ Important:** Always route through `apg()` for the base guidance computation. Raw linear interpolation (`result = uncond + w * (cond - uncond)`) produces severe audio artifacts (static, underwater sound, frequency distortion). The `apg()` function provides momentum smoothing, perpendicular projection, and norm thresholding that are essential for stable audio output. If your guidance mode needs custom math, apply it as a correction on top of the APG result.

**Available globals in guidance plugins:**

| Global | Type | Description |
|--------|------|-------------|
| `step_idx` | int | Current step index (0-based) |
| `total_steps` | int | Total number of denoising steps |
| `dt` | float | Current timestep delta (t_curr - t_next) |
| `t_curr` | float | Current timestep value |
| `params` | table | User-configured parameter values |

---

## The `apg()` Bridge

The `apg()` function is the native C++ APG (Analytical Perpendicular Guidance) implementation, exposed to Lua guidance plugins. It handles:

1. **Perpendicular projection** — removes the component of `(cond - uncond)` parallel to `uncond`, keeping only the steering signal
2. **Momentum smoothing** — exponential moving average across steps to prevent jitter
3. **Norm thresholding** — caps per-channel magnitudes to prevent blowup

**Signature:**

```lua
apg(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
```

All guidance plugins have access to this function. It is registered automatically on first use.

---

## Advanced: The `post_step()` Hook

For guidance modes that need to run extra model evaluations *after* the solver step (e.g., manifold projection), guidance plugins can declare a `post_step()` function. The engine detects this at load time and calls it after each solver step.

**When to use this:**
- Your guidance technique requires iterative refinement of the latent state
- You need to evaluate the model at positions different from the main solver trajectory
- The technique calls the model with conditioning and unconditioning separately

**Performance warning:** Each call to `eval_cond` or `eval_uncond` runs a full model forward pass. This is expensive — use sparingly.

```lua
function post_step(xt, t, n, eval_cond, eval_uncond, vt_cond, vt_uncond)
    -- xt:          mutable FloatArray — latent state after solver step (modify in-place)
    -- t:           float — timestep we just stepped TO (t_next)
    -- n:           int — total elements in xt
    -- eval_cond:   function(xt_arr, t) — runs model with conditioning, writes to vt_cond
    -- eval_uncond: function(xt_arr, t) — runs model without conditioning, writes to vt_uncond
    -- vt_cond:     mutable FloatArray — output buffer for conditional velocity
    -- vt_uncond:   mutable FloatArray — output buffer for unconditional velocity

    -- Example: one iteration of manifold projection
    local a = math.abs(dt) / 2.0

    eval_uncond(xt, t)                        -- fills vt_uncond
    for i = 0, n - 1 do
        xt[i] = xt[i] - a * vt_uncond[i]     -- push away from uncond manifold
    end

    eval_cond(xt, t)                          -- fills vt_cond
    for i = 0, n - 1 do
        xt[i] = xt[i] + a * vt_cond[i]       -- pull toward cond manifold
    end
end
```

The `post_step` hook has access to the same globals as `guide()` (`step_idx`, `total_steps`, `dt`, `t_curr`, `params`).

The hook is **not called on the final step** (the latent is about to be decoded, so further projection is pointless).

---

## Parameter Schema

Plugins can declare user-facing parameters that appear in the UI. Parameters are defined in the `params` array of the metadata table.

### Slider

```lua
{ key = "strength", type = "slider", label = "Strength",
  default = 0.5, min = 0.0, max = 1.0, step = 0.01,
  hint = "Controls the effect intensity" }
```

### Select (Dropdown)

```lua
{ key = "mode", type = "select", label = "Mode",
  default = "fast",
  options = {
      { value = "fast",    label = "Fast" },
      { value = "quality", label = "Quality" },
  },
  hint = "Choose between speed and quality" }
```

### Toggle

```lua
{ key = "enabled", type = "toggle", label = "Enable Feature",
  default = false,
  hint = "Turn this feature on or off" }
```

### Conditional Visibility

Parameters can be shown/hidden based on another parameter's value:

```lua
{ key = "sub_param", type = "slider", label = "Sub-Parameter",
  default = 1.0, min = 0.0, max = 5.0, step = 0.1,
  visible_when = { key = "mode", equals = "quality" },
  hint = "Only visible when Mode is set to Quality" }
```

### Transform Expressions

The `transform` field allows the UI to apply a mathematical transformation to the displayed value before sending it to the plugin. This is useful when the internal value differs from what the user sees:

```lua
{ key = "sigma", type = "slider", label = "Noise σ",
  default = 5, min = 0, max = 100, step = 1,
  transform = "value * 0.05",
  hint = "Displayed as 0-100, sent to plugin as 0-5" }
```

### Reading Parameters

Parameters are available via the `params` global table, keyed by their `key` field:

```lua
local strength = (params and params.strength) or 0.5
local mode     = (params and params.mode) or "fast"
local enabled  = (params and params.enabled) or false
```

Always provide a fallback default with `or` — `params` may be `nil` if no parameters have been set.

---

## FloatArray

All array data passes between C++ and Lua via the `FloatArray` userdata type. This is a **zero-copy** bridge — Lua reads and writes the same memory that the C++ engine uses.

**Indexing:** 0-based (matching C++ convention, not Lua's typical 1-based).

```lua
-- Read
local val = xt[i]

-- Write (only on mutable arrays)
xt[i] = val

-- Length
local n = #xt
```

Read-only arrays (like `pred_cond` and `pred_uncond` in guidance) will raise an error if you try to write to them.

---

## Available Globals

### Solver globals

| Global | Type | Description |
|--------|------|-------------|
| `step_index` | int | Current step index |
| `batch_n` | int | Number of batch elements |
| `n_per` | int | Elements per batch element |
| `params` | table | Plugin parameters |

### Guidance globals

| Global | Type | Description |
|--------|------|-------------|
| `step_idx` | int | Current step index (0-based) |
| `total_steps` | int | Total denoising steps |
| `dt` | float | Timestep delta |
| `t_curr` | float | Current timestep |
| `params` | table | Plugin parameters |

---

## Sandbox

Each plugin runs in an isolated Lua 5.4 VM with:

**Available:** `math`, `string`, `table`, `print`, `type`, `pairs`, `ipairs`, `tonumber`, `tostring`, `require` (for companion data files)

**Blocked:** `os`, `io`, `debug`, `dofile`, `loadfile` — no filesystem access, no shell commands, no process control.

The `require()` function works for loading companion Lua data files (e.g., precomputed constants in a separate `.lua` file in the same directory), but cannot load C modules.

---

## Complete Examples

### Minimal Solver

```lua
-- my_solver.lua
solver = {
    name        = "my_solver",
    display     = "My Solver",
    description = "Simple Euler variant",
    nfe         = 1,
    order       = 1,
    needs_model = 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
```

### Scheduler with Custom Curve

```lua
-- my_schedule.lua
scheduler = {
    name        = "my_schedule",
    display     = "Quadratic",
    description = "Quadratic timestep spacing",
    params      = {
        { key = "power", type = "slider", label = "Power",
          default = 2.0, min = 1.0, max = 4.0, step = 0.1 },
    },
}

function schedule(output, num_steps, shift)
    local p = (params and params.power) or 2.0
    for i = 0, num_steps - 1 do
        local frac = i / num_steps
        output[i] = (1.0 - frac) ^ p
    end
    -- Apply shift warp
    if shift ~= 1.0 then
        for i = 0, num_steps - 1 do
            local t = output[i]
            output[i] = shift * t / (1.0 + (shift - 1.0) * t)
        end
    end
end
```

### Guidance with APG + Custom Logic

```lua
-- my_guidance.lua
guidance = {
    name        = "my_guidance",
    display     = "My Guidance",
    description = "Warm-up guidance with linear ramp",
    params      = {
        { key = "warmup_steps", type = "slider", label = "Warm-Up Steps",
          default = 3, min = 0, max = 10, step = 1 },
    },
}

function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
    local warmup = (params and params.warmup_steps) or 3
    local progress = math.min((step_idx or 0) / math.max(warmup, 1), 1.0)
    local effective_scale = 1.0 + (guidance_scale - 1.0) * progress

    apg(pred_cond, pred_uncond, effective_scale, result, Oc, T, norm_threshold)
end
```

---

## Tips

- **Test with Euler + Linear first.** The simplest solver/scheduler combination isolates your plugin's behaviour.
- **Use `print()` for debugging.** Output goes to the terminal panel in the app.
- **Lua tables for scratch space.** If you need temporary arrays, use Lua tables: `local tmp = {}; for i = 0, n-1 do tmp[i] = 0 end`. They're slower than FloatArrays but work for intermediate calculations.
- **Be careful with the loop range.** FloatArrays are 0-indexed: `for i = 0, n - 1 do ... end`.
- **Guidance plugins: always use `apg()`.** Raw math without APG produces audio artifacts. Apply your custom logic as a delta on top.
- **Stateful plugins** can use file-level `local` variables to carry state across steps (e.g., previous velocity buffers, error accumulators). These reset when the plugin is reloaded.