From adc46a445f047a18e1df3e05ed2776d478a88198 Mon Sep 17 00:00:00 2001 From: dog Date: Thu, 9 Jul 2026 17:55:45 +0800 Subject: [PATCH] feat: switch flow ODE sampling from Euler to Heun (10 steps, NFE 20) Co-Authored-By: Claude Fable 5 --- src/ai_mouse/mouse.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/ai_mouse/mouse.py b/src/ai_mouse/mouse.py index f312e5d..ec26a67 100644 --- a/src/ai_mouse/mouse.py +++ b/src/ai_mouse/mouse.py @@ -24,7 +24,7 @@ from ai_mouse._postprocess import ( ) from ai_mouse.errors import GenerationError, ModelLoadError -_N_EULER_STEPS = 10 +_N_ODE_STEPS = 10 class MouseModel: @@ -90,11 +90,16 @@ class MouseModel: )[None] x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32) - dt = 1.0 / _N_EULER_STEPS - for step in range(_N_EULER_STEPS): - t = np.full((1,), step * dt, dtype=np.float32) - v = self._session.run(["v"], {"x_t": x, "t": t, "cond": cond})[0] - x = x + v * dt + # Heun (2nd-order predictor-corrector): same step count as the old + # Euler loop but far lower integration error for 2x NFE. + dt = 1.0 / _N_ODE_STEPS + for step in range(_N_ODE_STEPS): + t0 = np.full((1,), step * dt, dtype=np.float32) + v1 = self._session.run(["v"], {"x_t": x, "t": t0, "cond": cond})[0] + x_pred = (x + v1 * dt).astype(np.float32) + t1 = np.full((1,), (step + 1) * dt, dtype=np.float32) + v2 = self._session.run(["v"], {"x_t": x_pred, "t": t1, "cond": cond})[0] + x = x + (v1 + v2) * (dt / 2.0) if not np.all(np.isfinite(x)): raise GenerationError("Trajectory contains NaN/Inf after Euler integration")