field to override the timestep shift for sampling - #1692
Open
dxqb wants to merge 1 commit into
Open
Conversation
Flow-matching models shift the sampling schedule, either derived per sample from the token count (Flux 1, Flux 2, Qwen, Krea 2) or fixed by the scheduler config (Z-Image, Chroma, Anima, Ernie, HunyuanVideo). This adds an optional override_shift to SampleConfig and an "override shift:" entry in the sample frame, shown for the model types listed in ModelType.supports_sample_shift_override(). Empty keeps the model's own shift, so nothing changes unless the field is filled in. The models that derive a shift take the override in place of calculate_timestep_shift. Flux 2 is the odd one out: it works in mu directly via compute_empirical_mu, so the override goes in as log(shift) to keep the field meaning one thing across models. The fixed-shift models stamp it on their private scheduler copy with set_shift(), which set_timesteps() honours for both the step-count and the explicit-sigmas call shapes. The two mechanisms build different sigma grids, so the derived-shift samplers keep passing mu rather than switching to a static shift. The entry goes in the row below the steps entry, which moves the inpainting switch and the inpainting image paths down one row each. Ideogram is deliberately left out. Its sampler computes a resolution-aware mu and feeds it to a logit-normal sigma generator rather than to the scheduler's shift, so "shift" would not mean the same thing there. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.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
1 participant
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.
Flow-matching models shift the sampling schedule. This PR lets you override this automated calculation.
This is especially intended for video models where the linear interpolation defined in the scheduler config breaks down, because the number of tokens are far beyond the calibrated max shift and the calculation becomes an exponential extrapolation, resulting in huge shifts.
However, it can be useful for other models and SwarmUI also lets you override the shift, so this PR adds it first for existing models
Summary
Test plan
pre-commit run --all-filespassesAI assistance