Skip to content

Commit a6bf8b4

Browse files
committed
Fix double timestep shifting in FlowMatchEulerDiscreteScheduler.set_timesteps (#13243)
`__init__` records `self.sigma_max` / `self.sigma_min` *after* applying the `shift * sigmas / (1 + (shift - 1) * sigmas)` transform. `set_timesteps` then fed those already-shifted endpoints back into `np.linspace(...)` and applied the same shift again, producing a different schedule for the same `(num_train_timesteps, num_inference_steps, shift)` triple than `__init__`. Use the pre-shift `[num_train_timesteps, ..., 1]` range directly when no custom timesteps are provided so the shift is applied exactly once. After the fix, `__init__` and `set_timesteps(num_train_timesteps)` produce identical sigmas (to atol=1e-5) for every value of `shift` (verified for 1.0 and 3.0). `self.sigma_max`/`self.sigma_min` are left as-is to preserve their public attribute semantics for any downstream consumers.
1 parent 48f39c2 commit a6bf8b4

1 file changed

Lines changed: 6 additions & 2 deletions

File tree

src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -333,9 +333,13 @@ def set_timesteps(
333333

334334
if sigmas is None:
335335
if timesteps is None:
336+
# Use the raw, pre-shift `[num_train_timesteps, ..., 1]` range here:
337+
# `self.sigma_max`/`self.sigma_min` are recorded *after* the shift transform
338+
# in `__init__`, so feeding them back into the linspace and then re-applying
339+
# the shift below would shift the schedule twice (#13243).
336340
timesteps = np.linspace(
337-
self._sigma_to_t(self.sigma_max),
338-
self._sigma_to_t(self.sigma_min),
341+
self.config.num_train_timesteps,
342+
1,
339343
num_inference_steps,
340344
)
341345
sigmas = timesteps / self.config.num_train_timesteps

0 commit comments

Comments
 (0)