fix(flux2): support Heun scheduler for FLUX.2 Klein models#8794
Merged
lstein merged 3 commits intoinvoke-ai:mainfrom Feb 1, 2026
Merged
fix(flux2): support Heun scheduler for FLUX.2 Klein models#8794lstein merged 3 commits intoinvoke-ai:mainfrom
lstein merged 3 commits intoinvoke-ai:mainfrom
Conversation
FlowMatchHeunDiscreteScheduler does not support dynamic shifting parameters (use_dynamic_shifting, base_shift, max_shift, etc.) or sigmas/mu in set_timesteps. This caused FLUX.2 Klein to fail when using Heun scheduler. - Create Heun scheduler with only num_train_timesteps and shift parameters - Use num_inference_steps instead of sigmas for Heun's set_timesteps call - Euler and LCM schedulers continue to use full dynamic shifting support
github-actions
bot
added
python
The previous hasattr check for state_in_first_order failed because the attribute doesn't exist before set_timesteps() is called. Now using inspect.signature to check for sigmas parameter support, matching the FLUX1 implementation.
Pfannkuchensack
requested review from
JPPhoto,
blessedcoolant,
dunkeroni and
lstein
as code owners
Collaborator
Author
Heun works fine now.
|
lstein
added
v6.12.0
JPPhoto
approved these changes
Feb 1, 2026
dunkeroni
pushed a commit
to dunkeroni/InvokeAI
that referenced
this pull request
Feb 2, 2026
…#8794) * fix(flux2): support Heun scheduler for FLUX.2 Klein models FlowMatchHeunDiscreteScheduler does not support dynamic shifting parameters (use_dynamic_shifting, base_shift, max_shift, etc.) or sigmas/mu in set_timesteps. This caused FLUX.2 Klein to fail when using Heun scheduler. - Create Heun scheduler with only num_train_timesteps and shift parameters - Use num_inference_steps instead of sigmas for Heun's set_timesteps call - Euler and LCM schedulers continue to use full dynamic shifting support * fix(flux2): fix Heun scheduler detection using inspect.signature The previous hasattr check for state_in_first_order failed because the attribute doesn't exist before set_timesteps() is called. Now using inspect.signature to check for sigmas parameter support, matching the FLUX1 implementation. --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com>
lstein
pushed a commit
that referenced
this pull request
Feb 6, 2026
* fix(flux2): support Heun scheduler for FLUX.2 Klein models FlowMatchHeunDiscreteScheduler does not support dynamic shifting parameters (use_dynamic_shifting, base_shift, max_shift, etc.) or sigmas/mu in set_timesteps. This caused FLUX.2 Klein to fail when using Heun scheduler. - Create Heun scheduler with only num_train_timesteps and shift parameters - Use num_inference_steps instead of sigmas for Heun's set_timesteps call - Euler and LCM schedulers continue to use full dynamic shifting support * fix(flux2): fix Heun scheduler detection using inspect.signature The previous hasattr check for state_in_first_order failed because the attribute doesn't exist before set_timesteps() is called. Now using inspect.signature to check for sigmas parameter support, matching the FLUX1 implementation. --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
FlowMatchHeunDiscreteScheduler does not support dynamic shifting parameters (use_dynamic_shifting, base_shift, max_shift, etc.) or sigmas/mu in set_timesteps. This caused FLUX.2 Klein to fail when using Heun scheduler.
Related Issues / Discussions
None
QA Instructions
Merge Plan
None - straightforward fix, can be merged directly.
Checklist
What's Newcopy (if doing a release after this PR)