Skip to content

Fix FlowMatchEulerDiscreteScheduler double-shift in set_timesteps#13760

Open
Ricardo-M-L wants to merge 1 commit into
huggingface:mainfrom
Ricardo-M-L:fix-flow-match-double-shift
Open

Fix FlowMatchEulerDiscreteScheduler double-shift in set_timesteps#13760
Ricardo-M-L wants to merge 1 commit into
huggingface:mainfrom
Ricardo-M-L:fix-flow-match-double-shift

Conversation

@Ricardo-M-L
Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes #13243FlowMatchEulerDiscreteScheduler.set_timesteps applies the static timestep shift a second time when use_dynamic_shifting=False, so calling set_timesteps(num_train_timesteps) on a scheduler produces a different sigma schedule than the one __init__ cached.

Root cause

__init__ records sigma_min/sigma_max from the post-shift sigma range:

sigmas = timesteps / num_train_timesteps
if not use_dynamic_shifting:
    sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # ← shift applied
...
self.sigma_min = self.sigmas[-1].item()  # ← post-shift value
self.sigma_max = self.sigmas[0].item()   # ← post-shift value

set_timesteps then rebuilds sigmas via linspace(_sigma_to_t(sigma_max), _sigma_to_t(sigma_min), N) — already in shifted space — and applies the same shift again:

sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)

Reproduction

from diffusers import FlowMatchEulerDiscreteScheduler
import torch

s = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=False)
init_sigmas = s.sigmas.clone()
s.set_timesteps(num_inference_steps=1000)

print(init_sigmas[-3:])       # tensor([0.0089, 0.0060, 0.0030])
print(s.sigmas[-4:-1])        # tensor([0.0148, 0.0119, 0.0089])  ← different
torch.allclose(init_sigmas, s.sigmas[:-1])  # False

The same args produce two different schedules.

Fix

Record the pre-shift sigmas in __init__ and derive sigma_min/sigma_max from them. set_timesteps then regenerates the linear sigma range correctly and the static shift is applied exactly once. self.sigmas is unchanged (still post-shift); only the cached endpoints change.

Test

Added tests/schedulers/test_scheduler_flow_match_euler_discrete.py:

  • test_set_timesteps_matches_init_with_static_shift — fails on main, passes with the fix.
  • test_dynamic_shifting_is_unaffected — pins that the use_dynamic_shifting=True path is untouched (no static shift in __init__, so this code path never had the bug).

Before submitting

Who can review?

@yiyixuxu @DN6

`__init__` was storing `sigma_min`/`sigma_max` from the post-shift sigma
range when `use_dynamic_shifting=False`. `set_timesteps` then rebuilt
sigmas via `linspace(_sigma_to_t(sigma_max), _sigma_to_t(sigma_min), N)`
— already in shifted space — and applied the same static shift again:

    sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)

So calling `scheduler.set_timesteps(num_train_timesteps)` on a scheduler
constructed with the same args produced a different sigma schedule than
the one cached in `__init__`. For `shift=3, N=1000` the smallest sigma
became 0.00893 vs the init value of 0.00299.

Record the pre-shift sigmas in `__init__` and derive
`sigma_min`/`sigma_max` from those, so `set_timesteps` regenerates the
linear range and the shift is applied exactly once. `self.sigmas` is
unchanged (still post-shift). Added a regression test that constructs
the scheduler, snapshots `self.sigmas`, calls `set_timesteps` with the
same `num_inference_steps`, and asserts the two sigma schedules match.
The `use_dynamic_shifting=True` path is unaffected (no static shift in
`__init__`) and a second test pins that.

Fixes huggingface#13243
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] FlowMatchEulerDiscreteScheduler.__init__ computes sigma_min/sigma_max after shift, causing duplicate shift in set_timesteps

1 participant