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

This reverts commit adc46a445f.
This commit is contained in:
dog
2026-07-09 18:03:01 +08:00
parent adc46a445f
commit 7c33c13a87

View File

@@ -24,7 +24,7 @@ from ai_mouse._postprocess import (
)
from ai_mouse.errors import GenerationError, ModelLoadError
_N_ODE_STEPS = 10
_N_EULER_STEPS = 10
class MouseModel:
@@ -90,16 +90,11 @@ class MouseModel:
)[None]
x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32)
# 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)
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
if not np.all(np.isfinite(x)):
raise GenerationError("Trajectory contains NaN/Inf after Euler integration")