| |
| """DPM-Solver++(2M) with denoise-to-zero, as Autoware runs it around the decoder (multi_step mode). |
| |
| Port of ``PKG/src/inference/solver/dpm_solver.cpp`` (guidance disabled, Autoware's default and decision D11) and of |
| the prefix constraint of ``PKG/src/inference/multi_step_inference.cpp:300-340``. The scalar schedule is float32 math |
| evaluated with the C library's ``expf`` / ``logf`` / ``expm1f`` / ``sqrtf`` (what ``std::exp(float)`` etc. call in |
| the node) in the C++ operation order, so the timesteps and update coefficients are the node's bit for bit; numpy's |
| own float32 transcendentals differ from glibc in the last bit for many arguments. Without a loadable libm the module |
| falls back to numpy (``SCALAR_MATH`` says which). |
| |
| For ``steps = 10``: 11 decoder evaluations at ``timesteps[0..9]`` and ``1/N = 0.001`` (denoise-to-zero); the 11 |
| published iterates (``denoising_steps``) are the corrected ``x`` after each update and after the final evaluation. |
| The update formulas, with ``x`` and the model outputs float32 arrays and every scalar a float32:: |
| |
| first (step 1): x = (sigma_t / sigma_s) * x - (alpha_t * phi_1) * m |
| second (steps 2..): d = (m0 - m1) / r0 |
| x = (sigma_t / sigma_0) * x - (alpha_t * phi_1) * m0 - (0.5 * (alpha_t * phi_1)) * d |
| |
| :class:`SolverPlan` lists these scalars per update (the RT-dev coefficient tables of the on-device loop). |
| """ |
| from __future__ import annotations |
|
|
| import ctypes |
| import ctypes.util |
| from dataclasses import dataclass, field |
| from typing import Callable, List, Optional, Tuple |
|
|
| import numpy as np |
|
|
| from ..reference import config as C |
|
|
| __all__ = ["SCALAR_MATH", "marginal_log_mean_coeff", "marginal_alpha", "marginal_std", "marginal_lambda", |
| "inverse_lambda", "log_snr_timesteps", "SolverUpdate", "SolverPlan", "solver_plan", "first_update", |
| "second_update", "dpm_solver_sample", "apply_prefix_constraint", "SampleResult"] |
|
|
| f32 = np.float32 |
| _T, _N = C.NOISE_SCHEDULE_T, C.NOISE_SCHEDULE_TOTAL_N |
| _B0, _B1 = C.NOISE_SCHEDULE_BETA0, C.NOISE_SCHEDULE_BETA1 |
|
|
|
|
| class _ScalarMath: |
| """float32 ``exp`` / ``log`` / ``expm1`` / ``sqrt`` of the C library, loaded on first use (numpy fallback).""" |
|
|
| def __init__(self) -> None: |
| self._fns = None |
| self.source = "unloaded" |
|
|
| def _load(self) -> None: |
| fns = None |
| try: |
| name = ctypes.util.find_library("m") or "libm.so.6" |
| libm = ctypes.CDLL(name) |
| fns = {} |
| for fn in ("expf", "logf", "expm1f", "sqrtf"): |
| f = getattr(libm, fn) |
| f.restype, f.argtypes = ctypes.c_float, [ctypes.c_float] |
| fns[fn] = f |
| self.source = f"libm ({name})" |
| except (OSError, AttributeError): |
| fns = None |
| self.source = "numpy (no C math library found)" |
| self._fns = fns or {} |
|
|
| def _call(self, fn: str, x: float, fallback: Callable) -> np.float32: |
| if self._fns is None: |
| self._load() |
| if fn in self._fns: |
| return f32(self._fns[fn](float(f32(x)))) |
| return f32(fallback(f32(x))) |
|
|
| def exp(self, x): |
| return self._call("expf", x, np.exp) |
|
|
| def log(self, x): |
| return self._call("logf", x, np.log) |
|
|
| def expm1(self, x): |
| return self._call("expm1f", x, np.expm1) |
|
|
| def sqrt(self, x): |
| return self._call("sqrtf", x, np.sqrt) |
|
|
|
|
| SCALAR_MATH = _ScalarMath() |
| _m = SCALAR_MATH |
|
|
|
|
| |
| def marginal_log_mean_coeff(t) -> np.float32: |
| """``-0.25f * t * t * (beta1 - beta0) - 0.5f * t * beta0``.""" |
| t = f32(t) |
| return f32(f32(f32(f32(-0.25) * t) * t) * f32(_B1 - _B0)) - f32(f32(f32(0.5) * t) * _B0) |
|
|
|
|
| def marginal_alpha(t) -> np.float32: |
| return _m.exp(marginal_log_mean_coeff(t)) |
|
|
|
|
| def marginal_std(t) -> np.float32: |
| return _m.sqrt(f32(f32(1.0) - _m.exp(f32(f32(2.0) * marginal_log_mean_coeff(t))))) |
|
|
|
|
| def marginal_lambda(t) -> np.float32: |
| lmc = marginal_log_mean_coeff(t) |
| log_std = f32(f32(0.5) * _m.log(f32(f32(1.0) - _m.exp(f32(f32(2.0) * lmc))))) |
| return f32(lmc - log_std) |
|
|
|
|
| def _log_add_exp(a, b) -> np.float32: |
| m = max(f32(a), f32(b)) |
| return f32(m + _m.log(f32(_m.exp(f32(f32(a) - m)) + _m.exp(f32(f32(b) - m))))) |
|
|
|
|
| def inverse_lambda(lam) -> np.float32: |
| beta_delta = f32(_B1 - _B0) |
| tmp = f32(f32(f32(2.0) * beta_delta) * _log_add_exp(f32(f32(-2.0) * f32(lam)), f32(0.0))) |
| delta = f32(f32(_B0 * _B0) + tmp) |
| return f32(f32(tmp / f32(_m.sqrt(delta) + _B0)) / beta_delta) |
|
|
|
|
| def log_snr_timesteps(steps: int) -> List[np.float32]: |
| """``steps + 1`` times uniform in log-SNR between t = 1 and t = 1/N (``dpm_solver.cpp:80-94``).""" |
| t0 = f32(f32(1.0) / _N) |
| lam_t, lam_0 = marginal_lambda(_T), marginal_lambda(t0) |
| out = [] |
| for i in range(steps + 1): |
| ratio = f32(f32(i) / f32(steps)) |
| out.append(inverse_lambda(f32(lam_t + f32(f32(lam_0 - lam_t) * ratio)))) |
| return out |
|
|
|
|
| |
| def first_update(x_s: np.ndarray, model_s: np.ndarray, s, t) -> np.ndarray: |
| h = f32(marginal_lambda(t) - marginal_lambda(s)) |
| sigma_s, sigma_t, alpha_t = marginal_std(s), marginal_std(t), marginal_alpha(t) |
| phi_1 = _m.expm1(f32(-h)) |
| a, b = f32(sigma_t / sigma_s), f32(alpha_t * phi_1) |
| return (a * x_s.astype(np.float32) - b * model_s.astype(np.float32)).astype(np.float32) |
|
|
|
|
| def second_update(x_s: np.ndarray, model_prev: Tuple[np.ndarray, np.ndarray], t_prev: Tuple[float, float], |
| t) -> np.ndarray: |
| m1, m0 = model_prev |
| t1, t0 = t_prev |
| lam1, lam0, lam_t = marginal_lambda(t1), marginal_lambda(t0), marginal_lambda(t) |
| sigma0, sigma_t, alpha_t = marginal_std(t0), marginal_std(t), marginal_alpha(t) |
| h0, h = f32(lam0 - lam1), f32(lam_t - lam0) |
| r0 = f32(h0 / h) |
| phi_1 = _m.expm1(f32(-h)) |
| a, b = f32(sigma_t / sigma0), f32(alpha_t * phi_1) |
| c = f32(f32(0.5) * b) |
| d1_0 = ((m0.astype(np.float32) - m1.astype(np.float32)) / r0).astype(np.float32) |
| return (a * x_s.astype(np.float32) - b * m0.astype(np.float32) - c * d1_0).astype(np.float32) |
|
|
|
|
| @dataclass(frozen=True) |
| class SolverUpdate: |
| """One solver update to time ``t``: ``x = a * x - b * m0 - c * (m0 - m1) / r0`` (``c = 0`` for the first-order |
| update, which has no ``m1``).""" |
|
|
| order: int |
| t: float |
| a: float |
| b: float |
| c: float |
| r0: float |
|
|
|
|
| @dataclass(frozen=True) |
| class SolverPlan: |
| """Everything the on-device loop needs for ``steps`` (the RT-dev tables of PLAN.md 2.12): the decoder evaluation |
| times (``eval_times[k]`` feeds evaluation ``k``; the t-embedding / adaLN tables are built for these values), the |
| update after each of the first ``steps`` evaluations, and the published iterate times.""" |
|
|
| steps: int |
| timesteps: Tuple[float, ...] |
| eval_times: Tuple[float, ...] |
| updates: Tuple[SolverUpdate, ...] |
| denoising_timesteps: Tuple[float, ...] |
| scalar_math: str = field(default="") |
|
|
|
|
| def solver_plan(steps: int = C.DPM_SOLVER_STEPS) -> SolverPlan: |
| if steps < C.DPM_SOLVER_ORDER: |
| raise ValueError("DpmSolver steps must be greater than or equal to solver order.") |
| ts = log_snr_timesteps(steps) |
| updates = [] |
| for step in range(1, steps + 1): |
| t = ts[step] |
| if step < C.DPM_SOLVER_ORDER: |
| s = _T if step == 1 else ts[step - 1] |
| h = f32(marginal_lambda(t) - marginal_lambda(s)) |
| a = f32(marginal_std(t) / marginal_std(s)) |
| b = f32(marginal_alpha(t) * _m.expm1(f32(-h))) |
| updates.append(SolverUpdate(1, float(t), float(a), float(b), 0.0, 1.0)) |
| else: |
| t1 = _T if step - 2 == 0 else ts[step - 2] |
| t0 = ts[step - 1] |
| lam1, lam0, lam_t = marginal_lambda(t1), marginal_lambda(t0), marginal_lambda(t) |
| h0, h = f32(lam0 - lam1), f32(lam_t - lam0) |
| b = f32(marginal_alpha(t) * _m.expm1(f32(-h))) |
| updates.append(SolverUpdate(2, float(t), float(f32(marginal_std(t) / marginal_std(t0))), float(b), |
| float(f32(f32(0.5) * b)), float(f32(h0 / h)))) |
| final_t = f32(f32(1.0) / _N) |
| return SolverPlan(steps, tuple(float(t) for t in ts), tuple(float(t) for t in ts[:steps]) + (float(final_t),), |
| tuple(updates), tuple(float(t) for t in ts[1:]) + (float(final_t),), SCALAR_MATH.source) |
|
|
|
|
| |
| @dataclass |
| class SampleResult: |
| final_x: np.ndarray |
| denoising_steps: List[np.ndarray] |
| denoising_timesteps: List[float] |
| eval_times: List[float] |
| nfe: int |
|
|
|
|
| def apply_prefix_constraint(x: np.ndarray, current_states: np.ndarray) -> np.ndarray: |
| """``x[agent, 0, :] = current_states[agent]`` in place (``multi_step_inference.cpp:327-340``); ``x`` is |
| ``[321, 81, 4]`` (or with a leading batch dim).""" |
| x[..., 0, :] = current_states |
| return x |
|
|
|
|
| def dpm_solver_sample(initial_x: np.ndarray, model_fn: Callable[[np.ndarray, np.float32], np.ndarray], |
| correcting_fn: Callable[[np.ndarray], None], steps: int = C.DPM_SOLVER_STEPS, |
| on_iterate: Optional[Callable[[int, np.ndarray], None]] = None) -> SampleResult: |
| """``DpmSolver::sample`` without guidance. ``model_fn(x, t)`` is one decoder evaluation (x0 prediction), |
| ``correcting_fn(x)`` the in-place prefix constraint.""" |
| if steps < C.DPM_SOLVER_ORDER: |
| raise ValueError("DpmSolver steps must be greater than or equal to solver order.") |
| x = np.array(initial_x, dtype=np.float32, copy=True) |
| correcting_fn(x) |
| ts = log_snr_timesteps(steps) |
| steps_out: List[np.ndarray] = [] |
| step_ts: List[float] = [] |
| eval_times: List[float] = [] |
|
|
| def evaluate(xx: np.ndarray, t) -> np.ndarray: |
| eval_times.append(float(t)) |
| return np.asarray(model_fn(xx, f32(t)), np.float32) |
|
|
| def record(xx: np.ndarray, t) -> None: |
| steps_out.append(xx.copy()) |
| step_ts.append(float(t)) |
| if on_iterate is not None: |
| on_iterate(len(steps_out) - 1, xx) |
|
|
| t_prev = [_T] |
| m_prev = [evaluate(x, ts[0])] |
| for step in range(1, C.DPM_SOLVER_ORDER): |
| t = ts[step] |
| x = first_update(x, m_prev[-1], t_prev[-1], t) |
| correcting_fn(x) |
| record(x, t) |
| t_prev.append(t) |
| m_prev.append(evaluate(x, t)) |
| for step in range(C.DPM_SOLVER_ORDER, steps + 1): |
| t = ts[step] |
| x = second_update(x, (m_prev[0], m_prev[1]), (t_prev[0], t_prev[1]), t) |
| correcting_fn(x) |
| record(x, t) |
| t_prev = [t_prev[1], t] |
| m_prev = [m_prev[1], None] |
| if step < steps: |
| m_prev[1] = evaluate(x, t) |
| final_t = f32(f32(1.0) / _N) |
| x = evaluate(x, final_t) |
| x = np.array(x, dtype=np.float32, copy=True) |
| correcting_fn(x) |
| record(x, final_t) |
| return SampleResult(x, steps_out, step_ts, eval_times, len(eval_times)) |
|
|
