From 39f2909a75e465bd36fd57d5adcc6713c3c64555 Mon Sep 17 00:00:00 2001 From: Ricardo-M-L Date: Tue, 21 Apr 2026 10:22:41 +0800 Subject: [PATCH] Fix double-shift in `FlowMatchEulerDiscreteScheduler.__init__` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When `use_dynamic_shifting=False` and `shift != 1.0`, calling `set_timesteps(num_train_timesteps)` produced sigmas that were shifted twice, so the schedule visible to inference did not match the schedule `__init__` already stored on `self.sigmas`. Root cause: `__init__` applied the static shift to `sigmas` and then read `self.sigma_min` / `self.sigma_max` from the *post-shift* tensor. `set_timesteps` later re-derives a sigma schedule by linspace-ing between `_sigma_to_t(sigma_max)` and `_sigma_to_t(sigma_min)`, dividing by `num_train_timesteps` (recovering the stored, already- shifted sigmas), and then applying `shift` again at sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) Concretely with `shift=3.0`, `num_train_timesteps=1000`: * `__init__`: last sigma = 3 * (1/1000) / (1 + 2 * 1/1000) ≈ 0.00299 * `set_timesteps(1000)` before this fix: last sigma ≈ 0.00893 (≈3× off — the shift applied on top of itself) Fix: record `sigma_min` / `sigma_max` from the raw, pre-shift sigmas (which are `1/num_train_timesteps` and `1.0` by construction). The shift is then applied exactly once — inside `set_timesteps`, on the caller's chosen inference schedule — matching what the stored `self.sigmas` represented all along. No external caller reads these values post-init: the only in-repo writers are the Z-Image pipelines, which overwrite `scheduler.sigma_min = 0.0` before sampling. Fixes #13243 --- .../schedulers/scheduling_flow_match_euler_discrete.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py index 0e4f5c6a1f97..cdfef1bcbd2e 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py @@ -128,6 +128,13 @@ def __init__( timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) sigmas = timesteps / num_train_timesteps + # Store sigma_min / sigma_max from the *unshifted* sigmas. `set_timesteps` uses these + # to seed the per-call schedule and then re-applies `shift` itself — storing the + # post-shift values here would make `set_timesteps` shift a second time and produce + # sigmas that don't match `__init__`'s schedule (#13243). + self.sigma_min = sigmas[-1].item() + self.sigma_max = sigmas[0].item() + if not use_dynamic_shifting: # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) @@ -140,8 +147,6 @@ def __init__( self._shift = shift self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication - self.sigma_min = self.sigmas[-1].item() - self.sigma_max = self.sigmas[0].item() @property def shift(self):