[joshistoast]

Summary

Several flow-matching families compute their resolution-dependent schedule shift (mu) with the same two-anchor line fit:

• FLUX.1
• Krea-2 Raw
• Qwen-Image
• Qwen-Image-2.1
• Z-Image
• LTX-2

Each family had its own copy, in backend helpers, invocation methods or inline code, and the "hold the shift at the end of its fit" fix had to land in every copy separately.

invokeai/backend/rectified_flow/schedule_shift.py now owns:

• the fit (ShiftFit), including holding at the upper anchor and, for LTX-2, at both anchors;
• the linear time-shift form;
• loading a checkpoint's own scheduler config, with the stock config as fallback (load_scheduler). Keys a checkpoint omits take diffusers' defaults, as before.

Each family now declares only its anchors and form. The per-family helpers are removed. Qwen-Image's node gets a small _schedule seam that owns its config precedence, explicit-shift override and token count.

CogView4 (square-root form), FLUX.2 (fixed mu), Anima, Wan, SD3 and ERNIE are unchanged.

Commits:

1. Five families move with bit-identical results.
2. LTX-2 moves separately, because its old evaluation order rounded differently. See below.

QA Instructions

• uv run --no-sync pytest tests/backend/rectified_flow tests/backend/flux tests/backend/ltx2 tests/backend/qwen_image_2_1 tests/app/invocations/test_krea2_denoise.py tests/app/invocations/test_qwen_image_denoise.py tests/app/invocations/test_z_image_schedule_shift.py tests/app/invocations/test_denoise_noise_inputs.py passes.
• Full suite (-n 4): 15002 passed. The failures were the 112 model-identification and model-path tests, which need git-lfs fixtures that were not installed on the test machine, and fail the same way on main.

Equivalence check. These results are one-off evidence, not a committed test:

• Recording. Before the change, each family's mu (for Z-Image, the shift) and sigma schedule were recorded from the unmodified code by driving the real entry points on CPU without models. The Krea-2 Raw and Qwen-Image nodes ran with a real FlowMatchEulerDiscreteScheduler.
• Grid. 27 token counts below, at, between and past every family's anchors, from 1 to 65537. Schedules were recorded at 6 of them, each at 1, 2, 7 and 30 steps.
• Config paths. Checkpoint scheduler_config.json, both complete and missing keys. Qwen-Image explicit shifts 0.7 and 3.0. Z-Image explicit shifts 0, 0.5, 3 and 25.28.
• Five families. FLUX.1, Krea-2 Raw, Qwen-Image, Qwen-Image-2.1 and Z-Image are bit-identical (float.hex) after the refactor. A separate comparison against a reimplementation of the previous code confirmed mu for every token count from 1 to 70000, and FLUX schedules for steps up to 50.
• LTX-2. The shared fit adds the intercept as one term, as diffusers' calculate_shift does. LTX-2 used to evaluate x·slope + base − slope·base_seq left to right.
  • At and outside the anchors the values are identical.
  • Inside the line, 618 of the token counts from 1 to 19,999 move by exactly 1 ulp, and none move further.
  • Float32 sigmas are unchanged in all 1,854 schedules checked: each of the 618 moved counts at 8, 30 and 50 steps.
• Platform caveat. This was measured on Linux x86_64. Bit-exact torch schedules depend on the CPU kernel torch dispatches to, which is why the evidence is not committed as a test.

Review

Resolved:

• A committed bit-exact golden fixture was replaced with literal closed-form test cases. Torch schedule bits differ across CPU kernels: a 1-ulp FLUX difference was reproduced with ATEN_CPU_CAPABILITY=default.
• Qwen-Image's scheduler wiring had lost its test. It now goes through a tested seam that covers the checkpoint config, the stock config, token count and explicit shift.
• A node-level test again drives an explicit Qwen-Image shift through _run_diffusion to the scheduler, so dropping self.shift at the call site fails a test.

Remaining follow-up: the frontend Upscale feature keeps its own copies of the Z-Image, Krea-2 and Qwen-Image shift (qwenImageSigmas) for its start-index clipping.

Compatibility / Rollout

• No HTTP API or persisted-data change.
• Generation output is unchanged for the five families.
• LTX-2's mu can differ by 1 ulp between the anchors. Its float32 sigma schedules are unchanged in every case checked.
• Removed or changed backend helpers. None are exported from invocation_api, and remaining callers fail loudly, but custom nodes that import them need updating:
  • flux.sampling_utils.time_shift keeps its name, but its signature changes from (mu, sigma, t) to (shift, t).
  • get_schedule drops its base_shift/max_shift arguments.
  • Removed: get_lin_function, MU_FIT_*, Krea-2's calculate_shift and KREA2_* shift constants, LTX-2's calculate_shift and LTX2_{BASE,MAX}_*, Qwen-Image 2.1's calculate_mu.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

[Pfannkuchensack]

Approving. One request before merge: please restore one _run_diffusion-driven Qwen-Image explicit-shift test (see inline).

Old vs. new numerics

I compared the merge base (b3c6364) with this head in separate processes, comparing values as float.hex on CPU:

• FLUX.1: mu for 1–70,000 tokens. get_schedule at 32 token counts × 12 step counts, plus every token count from 1 to 20,000 at 30 steps.
• Z-Image: shift for 1–70,000 tokens and the _get_sigmas grid. Explicit shifts 0, 0.5, 1, 3, 25.28 and 103777. The real _run_diffusion from 256 to 4096 px, with the entry sigma read back from the img2img blend.
• Krea-2 Raw and Qwen-Image: the real _run_diffusion with the real FlowMatchEulerDiscreteScheduler. Five configs: no scheduler/ folder, an empty folder, a full checkpoint config, a custom config with shift_terminal, and a config with missing keys. 15 sizes from 16² to 4096², at two step counts. Krea-2 with shift None and 0; Qwen-Image with shift None, 0.7 and 3.0.
• Qwen-Image 2.1: mu for 1–70,000 tokens, and build_sigmas for Base, Turbo and None.
• LTX-2: mu for 1–70,000 tokens. build_sigmas for every token count from 1 to 20,000 at 8, 30 and 50 steps, and the refine schedules.

Everything is bit-identical except LTX-2's mu. There, 618 token counts strictly between 1024 and 4096 move by 1 ulp; nothing at or outside the anchors moves, and every float32 LTX-2 sigma schedule is unchanged. That matches the description. The holds from #9887, #9892 and #9902 are kept exactly.

The tests pin the numbers well. These changes all make the focused tests fail:

• changing any family's anchor or shift constant
• turning off hold_below_base or removing the upper hold
• making from_scheduler_config read the wrong key
• making load_scheduler ignore the checkpoint config
• making Krea-2 or Qwen-Image fall back to the stock fit

Locally, ruff check, ruff format --check and the focused tests pass.

Conflict in settings.ts

The conflict with #9934 is textual only: #9934 added the shiftedSigmas helper right above the comment you changed. Keep main's block as is and take your `z_image_denoise.Z_IMAGE_SHIFT_FIT` comment. With that resolution:

• the focused tests, test_upscale_index_clipped_starts.py and test_frontend_graph_node_types.py pass on the merged tree;
• nothing on main still references a removed helper;
• the Krea-2 and Qwen-Image comments #9934 added there still point to the right files.

Nits

1. Compatibility note. Please add a line under Compatibility / Rollout for the removed backend helpers. None of them are in invocation_api, and anything that still uses them fails loudly, but custom-node authors should know:
   • flux.sampling_utils.time_shift keeps its name but changes signature from (mu, sigma, t) to (shift, t).
   • get_schedule drops base_shift/max_shift.
   • Removed: get_lin_function, MU_FIT_*, Krea-2's calculate_shift and KREA2_* shift constants, LTX-2's calculate_shift and LTX2_{BASE,MAX}_*, and Qwen-Image 2.1's calculate_mu.
2. PR description. Since #9934, the follow-up ("the frontend Upscale feature keeps its own copies of the Z-Image and Krea-2 shift") also covers Qwen-Image's qwenImageSigmas.