yellowdog/ai_mouse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: switch flow ODE sampling from Euler to Heun (10 steps, NFE 20)

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
dog
2026-07-09 17:55:45 +08:00
parent 441e6f3dfe
commit adc46a445f
1 changed files with 11 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/ai_mouse/mouse.py
View File

@@ -24,7 +24,7 @@ from ai_mouse._postprocess import (
) )
from ai_mouse.errors import GenerationError, ModelLoadError from ai_mouse.errors import GenerationError, ModelLoadError
_N_EULER_STEPS = 10 _N_ODE_STEPS = 10
class MouseModel: class MouseModel:
@@ -90,11 +90,16 @@ class MouseModel:
)[None] )[None]
x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32) x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32)
dt = 1.0 / _N_EULER_STEPS # Heun (2nd-order predictor-corrector): same step count as the old
for step in range(_N_EULER_STEPS): # Euler loop but far lower integration error for 2x NFE.
t = np.full((1,), step * dt, dtype=np.float32) dt = 1.0 / _N_ODE_STEPS
v = self._session.run(["v"], {"x_t": x, "t": t, "cond": cond})[0] for step in range(_N_ODE_STEPS):
x = x + v * dt 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)): if not np.all(np.isfinite(x)):
raise GenerationError("Trajectory contains NaN/Inf after Euler integration") raise GenerationError("Trajectory contains NaN/Inf after Euler integration")