From 02e3cf3c7b9ca633a7a17262c5fb0daee544426e Mon Sep 17 00:00:00 2001
From: Josh Corbett <joshwcorbett@icloud.com>
Date: Mon, 19 Oct 2026 00:03:01 -0600
Subject: [PATCH 1/3] refactor(sampling): share one schedule-shift fit across
 flow-matching families

FLUX.1, Krea-2 Raw, Qwen-Image, Qwen-Image-2.1 and Z-Image evaluate their resolution-dependent mu through ShiftFit; checkpoint scheduler configs load through load_scheduler.
Recorded mu and sigma schedules match the previous per-family code bit for bit.
---
 .../contributing/new-model-integration.mdx    |   3 +-
 .../app/invocations/krea2/krea2_denoise.py    |  38 +-----
 .../qwen_image/qwen_image_denoise.py          |  77 ++++++-----
 .../invocations/z_image/z_image_denoise.py    |  63 ++-------
 invokeai/backend/flux/sampling_utils.py       |  70 +++-------
 invokeai/backend/krea2/sampling_utils.py      |  43 +++----
 invokeai/backend/qwen_image_2_1/sampling.py   |  21 ++-
 .../backend/rectified_flow/schedule_shift.py  |  89 +++++++++++++
 .../src/features/upscale/core/graph.test.ts   |   2 +-
 .../src/features/upscale/core/settings.ts     |   2 +-
 .../invocations/test_denoise_noise_inputs.py  |   4 +-
 tests/app/invocations/test_krea2_denoise.py   | 121 +++++++-----------
 .../invocations/test_qwen_image_denoise.py    | 113 +++++-----------
 .../test_z_image_schedule_shift.py            |  35 ++---
 tests/backend/flux/test_sampling_utils.py     |  47 ++-----
 tests/backend/qwen_image_2_1/test_sampling.py |   4 +-
 .../rectified_flow/test_schedule_shift.py     | 101 +++++++++++++++
 17 files changed, 404 insertions(+), 429 deletions(-)
 create mode 100644 invokeai/backend/rectified_flow/schedule_shift.py
 create mode 100644 tests/backend/rectified_flow/test_schedule_shift.py

diff --git a/docs/src/content/docs/contributing/new-model-integration.mdx b/docs/src/content/docs/contributing/new-model-integration.mdx
index 04e4530b2d8..27b0cf2278a 100644
--- a/docs/src/content/docs/contributing/new-model-integration.mdx
+++ b/docs/src/content/docs/contributing/new-model-integration.mdx
@@ -1,7 +1,7 @@
 ---
 title: New Model Type Integration Checklist
 description: A step-by-step checklist for integrating a new model architecture into InvokeAI, from the model manager to the webv2 frontend and the tests that enforce completeness.
-lastUpdated: 2026-10-02
+lastUpdated: 2026-10-18
 ---
 
 import { Steps, FileTree } from '@astrojs/starlight/components';
@@ -626,6 +626,7 @@ Packages are discovered automatically. A directory of node modules without an `_
 
 - **Backend package (optional).** Put reusable math in `invokeai/backend/new_model/`: noise, packing, position IDs, schedules and text encoding. FLUX, FLUX.2 and Krea-2 have one; Qwen-Image and Z-Image keep the loop in the denoise node. Use a package once code is shared between nodes or needs its own tests.
 - **Schedulers.** Flow-matching architectures use diffusers' `FlowMatchEulerDiscreteScheduler`, or the shared maps in `invokeai/backend/flux/schedulers.py`: `FLUX_SCHEDULER_MAP`, `ZIMAGE_SCHEDULER_MAP` and `ERNIE_IMAGE_SCHEDULER_MAP`, each with name values and labels. Declare which set the UI offers in `FeaturesFacet.scheduler_set`. Set `scheduler_applies_to_graph=True` only if the denoise node really takes a scheduler field. Reproduce the reference sigma schedule and shift (`mu`, time-shift type, terminal shift) exactly, and test it against the reference scheduler.
+- **Resolution-dependent shift.** A `mu` interpolated from the image token count goes through `invokeai/backend/rectified_flow/schedule_shift.py`. Declare the two anchors as a `ShiftFit`, or read them with `ShiftFit.from_scheduler_config` from the scheduler `load_scheduler` returns: a checkpoint's own `scheduler/scheduler_config.json`, else the family's stock config for single-file and GGUF installs. Hand `fit.mu(...)` to `FlowMatchEulerDiscreteScheduler.set_timesteps(mu=...)` (exponential form), or `fit.shift(...)`, which is `e^mu`, to `time_shift` for a hand-built schedule (linear form). The fit is held at its upper anchor. Reference pipelines extrapolate past it, and at large frames that puts nearly every step at high noise, so an image-to-image pass replaces its input instead of refining it. Test the family's anchors and form; the module's own tests cover the line and the hold.
 - **External noise (optional).** Only if the denoise node accepts a noise tensor: extend `LatentNoiseType` and its shape and grid logic in `invokeai/app/invocations/latent_noise.py`, and extend `noise_type` in `invocations/noise.py`. Validate with `validate_noise_tensor_shape`.
 - **Timestep-dependent conditioning.** Most architectures turn the text embeddings into the transformer's context once, before the loop. If a module between encoder and transformer reads the timestep -- Anima-3.8B's semantic connector does -- the context has to be recomputed at every step, for the positive, the negative and every regional conditioning. Pass the timestep exactly as the reference does: the connector scales sigma by 1000 before embedding it, so it gets a float32 sigma, where bf16's rounding would move its high frequencies by radians. A mask that leaves a query no key (an empty prompt encoded as one masked token) needs a defined result; `masked_sdpa` makes it zero instead of leaving it to the SDPA backend.
 - **Inpainting.** Rectified-flow models use `RectifiedFlowInpaintExtension` (`invokeai/backend/rectified_flow/rectified_flow_inpaint_extension.py`).
diff --git a/invokeai/app/invocations/krea2/krea2_denoise.py b/invokeai/app/invocations/krea2/krea2_denoise.py
index 04aa13a51d5..5a6543d3e8c 100644
--- a/invokeai/app/invocations/krea2/krea2_denoise.py
+++ b/invokeai/app/invocations/krea2/krea2_denoise.py
@@ -39,13 +39,9 @@
     Krea2TextConditioning,
 )
 from invokeai.backend.krea2.sampling_utils import (
-    KREA2_BASE_IMAGE_SEQ_LEN,
-    KREA2_BASE_SHIFT,
     KREA2_DISTILLED_MU,
-    KREA2_MAX_IMAGE_SEQ_LEN,
-    KREA2_MAX_SHIFT,
+    KREA2_SCHEDULER_CONFIG,
     build_sigmas,
-    calculate_shift,
     pack_latents,
     prepare_position_ids,
     unpack_latents,
@@ -57,6 +53,7 @@
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes, requires_sidecar_patching
 from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
+from invokeai.backend.rectified_flow.schedule_shift import ShiftFit, load_scheduler
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Krea2ConditioningInfo
 from invokeai.backend.util import sdpa_scope
@@ -353,8 +350,6 @@ def _is_distilled(self, context: InvocationContext) -> bool:
         return True
 
     def _run_diffusion(self, context: InvocationContext):
-        from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-
         self._validate_inputs()
 
         device = TorchDevice.choose_torch_device()
@@ -377,21 +372,8 @@ def _run_diffusion(self, context: InvocationContext):
         )
         pos_prompt_embeds = pos_extension.regional_text_conditioning.prompt_embeds
 
-        # Scheduler: load from the model's scheduler/ dir if present, else construct with Krea-2 defaults.
         model_path = context.models.get_absolute_path(context.models.get_config(self.transformer.transformer))
-        scheduler_path = Path(model_path) / "scheduler"
-        if scheduler_path.is_dir() and (scheduler_path / "scheduler_config.json").exists():
-            scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(str(scheduler_path), local_files_only=True)
-        else:
-            scheduler = FlowMatchEulerDiscreteScheduler(
-                use_dynamic_shifting=True,
-                base_shift=0.5,
-                max_shift=1.15,
-                base_image_seq_len=256,
-                max_image_seq_len=6400,
-                num_train_timesteps=1000,
-                time_shift_type="exponential",
-            )
+        scheduler = load_scheduler(Path(model_path), KREA2_SCHEDULER_CONFIG)
 
         # A non-positive shift means "unset", not "no time shift". The node editor builds every float field
         # with `schemaObject.default ?? 0`, so an untouched optional shift arrives here as 0.0 rather than
@@ -402,17 +384,9 @@ def _run_diffusion(self, context: InvocationContext):
         elif self._is_distilled(context):
             mu = KREA2_DISTILLED_MU
         else:
-            # Resolution-aware shift, held at the end of its fit (see `calculate_shift`). Honor the loaded
-            # scheduler's config so a Raw checkpoint that ships a customized scheduler_config.json is sampled
-            # with its own shift params. A key that file omits takes the diffusers scheduler's own default,
-            # not Krea-2's; the Krea-2 defaults below apply only to the stock scheduler built above.
-            mu = calculate_shift(
-                image_seq_len,
-                base_image_seq_len=getattr(scheduler.config, "base_image_seq_len", KREA2_BASE_IMAGE_SEQ_LEN),
-                max_image_seq_len=getattr(scheduler.config, "max_image_seq_len", KREA2_MAX_IMAGE_SEQ_LEN),
-                base_shift=getattr(scheduler.config, "base_shift", KREA2_BASE_SHIFT),
-                max_shift=getattr(scheduler.config, "max_shift", KREA2_MAX_SHIFT),
-            )
+            # Resolution-aware shift, held at the end of its fit (see `KREA2_SCHEDULER_CONFIG`), from the loaded
+            # scheduler's config: a Raw checkpoint that ships its own scheduler_config.json samples with its fit.
+            mu = ShiftFit.from_scheduler_config(scheduler.config).mu(image_seq_len)
 
         init_sigmas = build_sigmas(self.steps)
         scheduler.set_timesteps(sigmas=init_sigmas, mu=mu, device=device)
diff --git a/invokeai/app/invocations/qwen_image/qwen_image_denoise.py b/invokeai/app/invocations/qwen_image/qwen_image_denoise.py
index 7f53e7a6cc5..37f4efa455a 100644
--- a/invokeai/app/invocations/qwen_image/qwen_image_denoise.py
+++ b/invokeai/app/invocations/qwen_image/qwen_image_denoise.py
@@ -1,6 +1,7 @@
 import math
 from contextlib import ExitStack
-from typing import Callable, ClassVar, Iterator, Optional
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, ClassVar, Iterator, Optional
 
 import torch
 import torchvision.transforms as tv_transforms
@@ -31,11 +32,31 @@
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes, requires_sidecar_patching
 from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
+from invokeai.backend.rectified_flow.schedule_shift import ShiftFit, load_scheduler
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import QwenImageConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.sage_attention import sage_attention_scope
 
+if TYPE_CHECKING:
+    from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+
+# The stock Qwen-Image scheduler, for GGUF and single-file transformers, which ship no scheduler/ folder. Its shift
+# fit runs from 256 tokens (mu 0.5) to 8192 tokens (mu 0.9, about 1448x1448 of area) and is held there. Past that end
+# the line would keep climbing -- mu 1.31 at 2048x2048, 3.79 at 4096x4096, where 28 of 40 steps sit above sigma 0.9
+# and an img2img pass at strength 0.5 starts from sigma 0.95, so it replaces the image instead of refining it.
+# diffusers' QwenImagePipeline extrapolates; ComfyUI samples Qwen-Image at a constant mu of 1.15.
+QWEN_IMAGE_SCHEDULER_CONFIG = {
+    "use_dynamic_shifting": True,
+    "base_shift": 0.5,
+    "max_shift": 0.9,
+    "base_image_seq_len": 256,
+    "max_image_seq_len": 8192,
+    "shift_terminal": 0.02,
+    "num_train_timesteps": 1000,
+    "time_shift_type": "exponential",
+}
+
 
 @invocation(
     "qwen_image_denoise",
@@ -94,6 +115,21 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
         name = context.tensors.save(tensor=latents)
         return LatentsOutput.build(latents_name=name, latents=latents, seed=None)
 
+    @staticmethod
+    def _schedule(
+        model_path: Path, width: int, height: int, patch_size: int, shift: float | None
+    ) -> tuple["FlowMatchEulerDiscreteScheduler", float]:
+        """The scheduler for this transformer and the mu to set its timesteps with.
+
+        mu is the fixed `shift`'s log when one is set (Lightning LoRAs), else the loaded scheduler config's own
+        resolution-dependent fit over the target image's tokens, held at its end.
+        """
+        scheduler = load_scheduler(model_path, QWEN_IMAGE_SCHEDULER_CONFIG)
+        if shift is not None:
+            return scheduler, math.log(shift)
+        image_seq_len = (height // LATENT_SCALE_FACTOR) * (width // LATENT_SCALE_FACTOR) // patch_size**2
+        return scheduler, ShiftFit.from_scheduler_config(scheduler.config).mu(image_seq_len)
+
     def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) -> torch.Tensor | None:
         if self.denoise_mask is None:
             return None
@@ -288,50 +324,13 @@ def _run_diffusion(self, context: InvocationContext):
 
         latent_height = self.height // LATENT_SCALE_FACTOR
         latent_width = self.width // LATENT_SCALE_FACTOR
-        image_seq_len = (latent_height * latent_width) // (patch_size**2)
 
         # Use the actual FlowMatchEulerDiscreteScheduler to compute sigmas/timesteps,
         # exactly matching the diffusers pipeline.
-        import math
-
         import numpy as np
-        from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 
-        # Try to load the scheduler config from the model's directory (Diffusers models
-        # have a scheduler/ subdir). For GGUF models this path doesn't exist, so fall
-        # back to instantiating the scheduler with the known Qwen Image defaults.
         model_path = context.models.get_absolute_path(context.models.get_config(self.transformer.transformer))
-        scheduler_path = model_path / "scheduler"
-        if scheduler_path.is_dir() and (scheduler_path / "scheduler_config.json").exists():
-            scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(str(scheduler_path), local_files_only=True)
-        else:
-            scheduler = FlowMatchEulerDiscreteScheduler(
-                use_dynamic_shifting=True,
-                base_shift=0.5,
-                max_shift=0.9,
-                base_image_seq_len=256,
-                max_image_seq_len=8192,
-                shift_terminal=0.02,
-                num_train_timesteps=1000,
-                time_shift_type="exponential",
-            )
-
-        if self.shift is not None:
-            # Lightning LoRA: fixed shift
-            mu = math.log(self.shift)
-        else:
-            # Dynamic shift: a line through the scheduler config's two points, held at the last one. The stock
-            # config runs from 256 tokens (mu 0.5) to 8192 tokens (mu 0.9, about 1448x1448 of area). Past that
-            # end the line would keep climbing -- mu 1.31 at 2048x2048, 3.79 at 4096x4096, where 28 of 40 steps
-            # sit above sigma 0.9 and an img2img pass at strength 0.5 starts from sigma 0.95, so it replaces the
-            # image instead of refining it. diffusers' QwenImagePipeline extrapolates; ComfyUI samples Qwen-Image
-            # at a constant mu of 1.15. The scheduler registers every init argument, so a key that a checkpoint's
-            # scheduler_config.json omits takes FlowMatchEulerDiscreteScheduler's own default (max_shift 1.15,
-            # max_image_seq_len 4096), not the stock Qwen-Image value.
-            config = scheduler.config
-            m = (config.max_shift - config.base_shift) / (config.max_image_seq_len - config.base_image_seq_len)
-            b = config.base_shift - m * config.base_image_seq_len
-            mu = min(image_seq_len, config.max_image_seq_len) * m + b
+        scheduler, mu = self._schedule(model_path, self.width, self.height, patch_size, self.shift)
 
         init_sigmas = np.linspace(1.0, 1.0 / self.steps, self.steps).tolist()
         scheduler.set_timesteps(sigmas=init_sigmas, mu=mu, device=device)
diff --git a/invokeai/app/invocations/z_image/z_image_denoise.py b/invokeai/app/invocations/z_image/z_image_denoise.py
index 2293901d432..a96d0962b40 100644
--- a/invokeai/app/invocations/z_image/z_image_denoise.py
+++ b/invokeai/app/invocations/z_image/z_image_denoise.py
@@ -34,6 +34,7 @@
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
 from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes, requires_sidecar_patching
 from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
+from invokeai.backend.rectified_flow.schedule_shift import ShiftFit, time_shift
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ZImageConditioningInfo
 from invokeai.backend.util.attention import sdpa_score_matrix_bytes
@@ -62,11 +63,13 @@ def _padded(tokens: int) -> int:
     return tokens + (-tokens) % SEQ_MULTI_OF
 
 
-# The two points the shift fit is built on. 4096 image tokens is a 1024x1024 frame -- one
-# megapixel, and where the reference pipeline's fit ends. Naming them keeps the fit and the clamp
-# in `_calculate_shift` from drifting apart.
-Z_IMAGE_MU_FIT_MIN_SEQ_LEN = 256
-Z_IMAGE_MU_FIT_MAX_SEQ_LEN = 4096
+# The shift fit of diffusers' ZImagePipeline. 4096 image tokens is a 1024x1024 frame -- one megapixel, and where the
+# reference fit ends. Extrapolated past it the shift runs away fast -- 25 at 2048px, 810 at 3072px and 103_777 at
+# 4096px -- and `_get_sigmas` flattens with it. Because `denoising_start` is clipped by *index* here, the step count
+# survives but the entry point does not: a canvas img2img pass at strength 0.5 would start from sigma 0.88 at 1536px,
+# 0.96 at 2048px and 1.0000 at 4096px, i.e. throw the input image away and generate a new one. Below one megapixel
+# the hold changes nothing, and an explicit `shift` on the invocation still overrides it.
+Z_IMAGE_SHIFT_FIT = ShiftFit(base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.15)
 
 
 @invocation(
@@ -252,55 +255,13 @@ def _get_noise(
             generator=torch.Generator(device=rand_device).manual_seed(seed),
         ).to(device=device, dtype=dtype)
 
-    def _calculate_shift(
-        self,
-        image_seq_len: int,
-        base_image_seq_len: int = Z_IMAGE_MU_FIT_MIN_SEQ_LEN,
-        max_image_seq_len: int = Z_IMAGE_MU_FIT_MAX_SEQ_LEN,
-        base_shift: float = 0.5,
-        max_shift: float = 1.15,
-    ) -> float:
-        """Calculate timestep shift based on image sequence length.
-
-        Based on diffusers ZImagePipeline.calculate_shift method.
-        Returns a linear shift value (exp(mu) from the original formula).
-
-        The fit is a straight line through two points and is not evaluated past the second one.
-        Extrapolated it runs away fast -- shift reaches 25 at 2048px, 810 at 3072px and 103_777 at
-        4096px -- and `_get_sigmas` flattens with it. Because `denoising_start` is clipped by *index*
-        here, the step count survives but the entry point does not: a canvas img2img pass at
-        strength 0.5 would start from sigma 0.88 at 1536px, 0.96 at 2048px and 1.0000 at 4096px,
-        i.e. throw the input image away and generate a new one. Below one megapixel the clamp
-        changes nothing, and an explicit `shift` on the invocation still overrides it.
-        """
-        m = (max_shift - base_shift) / (max_image_seq_len - base_image_seq_len)
-        b = base_shift - m * base_image_seq_len
-        mu = min(image_seq_len, max_image_seq_len) * m + b
-        # Convert from exponential mu to linear shift value
-        return math.exp(mu)
-
     def _get_sigmas(self, shift: float, num_steps: int) -> list[float]:
-        """Generate sigma schedule with linear time shift.
-
-        Uses linear time shift: shift / (shift + (1/t - 1)).
-        The shift value is used directly as a multiplier.
-        Generates num_steps + 1 sigma values (including terminal 0.0).
-        """
-
-        def time_shift(shift: float, t: float) -> float:
-            """Apply linear time shift to a single timestep value."""
-            if t <= 0:
-                return 0.0
-            if t >= 1:
-                return 1.0
-            return shift / (shift + (1 / t - 1))
-
+        """`num_steps + 1` sigmas from 1 to 0, linearly time-shifted by `shift` (`time_shift`)."""
         sigmas = []
         for i in range(num_steps + 1):
             t = 1.0 - i / num_steps  # Goes from 1.0 to 0.0
-            sigma = time_shift(shift, t)
-            sigmas.append(sigma)
-
+            # The ends are fixed. Evaluated, t = 0 would divide by zero, and so would t = 1 at a shift of 0.
+            sigmas.append(1.0 if t >= 1 else 0.0 if t <= 0 else time_shift(shift, t))
         return sigmas
 
     @staticmethod
@@ -433,7 +394,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
         if self.shift is not None:
             shift = self.shift
         else:
-            shift = self._calculate_shift(img_seq_len)
+            shift = Z_IMAGE_SHIFT_FIT.shift(img_seq_len)
 
         # Generate sigma schedule with time shift
         sigmas = self._get_sigmas(shift, self.steps)
diff --git a/invokeai/backend/flux/sampling_utils.py b/invokeai/backend/flux/sampling_utils.py
index 5ccf439617b..1b60c7b1350 100644
--- a/invokeai/backend/flux/sampling_utils.py
+++ b/invokeai/backend/flux/sampling_utils.py
@@ -1,11 +1,11 @@
 # Initially pulled from https://github.com/black-forest-labs/flux
 
 import math
-from typing import Callable
 
 import torch
 from einops import rearrange, repeat
 
+from invokeai.backend.rectified_flow.schedule_shift import ShiftFit, time_shift
 from invokeai.backend.util.devices import TorchDevice
 
 
@@ -32,60 +32,30 @@ def get_noise(
     ).to(device=device, dtype=dtype)
 
 
-def time_shift(mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
-    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
-
-
-# The two points the mu fit below is built on. 4096 image tokens is a 1024x1024 frame -- one
-# megapixel, and the resolution the reference pipeline defaults to. The line says nothing beyond it.
-MU_FIT_MIN_SEQ_LEN = 256
-MU_FIT_MAX_SEQ_LEN = 4096
-
-
-def get_lin_function(
-    x1: float = MU_FIT_MIN_SEQ_LEN, y1: float = 0.5, x2: float = MU_FIT_MAX_SEQ_LEN, y2: float = 1.15
-) -> Callable[[float], float]:
-    """A line through two points. Callers must not evaluate it past `x2` -- see `get_schedule`."""
-    m = (y2 - y1) / (x2 - x1)
-    b = y1 - m * x1
-    return lambda x: m * x + b
-
-
-def get_schedule(
-    num_steps: int,
-    image_seq_len: int,
-    base_shift: float = 0.5,
-    max_shift: float = 1.15,
-    shift: bool = True,
-) -> list[float]:
+# FLUX.1's shift fit. 4096 image tokens is a 1024x1024 frame -- one megapixel, and the resolution the reference
+# pipeline defaults to.
+#
+# Extrapolated past it, mu reaches 3.23 at 2048px, 6.70 at 3072px and 11.55 at 4096px, and the schedule flattens
+# with it -- the smallest non-zero timestep rises to 0.47, 0.97 and 0.9997. `clip_timestep_schedule_fractional` then
+# has almost nothing left below `1 - denoising_start`. Steps surviving a 30-step request at `denoising_start=0.5`:
+# 8 at 1 MP, 6 at 1.6 MP, 4 at 2.4 MP, 2 at 4.2 MP, and one from ~2112px square. From ~3376px square any
+# `denoising_start` of 0.01 or more leaves one; only a full denoise from 0 escapes the clip. Asking for more steps
+# barely helps: at 2048px and 0.499, 100 keep 4 and 1000 keep 39.
+#
+# The effect is graded, not a cliff: between 1 MP and ~1.6 MP the unheld schedule is not degenerate, and holding mu
+# there is a judgement -- never evaluate the fit outside the range it was built on -- rather than a defect fix.
+# ComfyUI's default FLUX.1 sampling uses a fixed shift of 1.15 at every resolution. Below 256 tokens mu merely sags to
+# 0.457, which is harmless. The hold is on an *area*, so a 1024x1280 canvas is already above it.
+FLUX_SHIFT_FIT = ShiftFit(base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.15)
+
+
+def get_schedule(num_steps: int, image_seq_len: int, shift: bool = True) -> list[float]:
     # extra step for zero
     timesteps = torch.linspace(1, 0, num_steps + 1)
 
     # shifting the schedule to favor high timesteps for higher signal images
     if shift:
-        # The fit is a straight line through two points, so it is not evaluated past the second one.
-        # Extrapolated, mu reaches 3.23 at 2048px, 6.70 at 3072px and 11.55 at 4096px, and the
-        # schedule flattens with it -- the smallest non-zero timestep rises to 0.47, 0.97 and
-        # 0.9997. `clip_timestep_schedule_fractional` then has almost nothing left below
-        # `1 - denoising_start`. Steps surviving a 30-step request at `denoising_start=0.5`:
-        # 8 at 1 MP, 6 at 1.6 MP, 4 at 2.4 MP, 2 at 4.2 MP, and one from ~2112px square. From
-        # ~3376px square any `denoising_start` of 0.01 or more leaves one; only a full denoise from 0
-        # escapes the clip. Asking for more steps barely helps: at 2048px and 0.499, 100 keep 4 and
-        # 1000 keep 39.
-        #
-        # The effect is graded, not a cliff: between 1 MP and ~1.6 MP the unclamped schedule is not
-        # degenerate, and holding mu there is a judgement -- never evaluate the fit outside the
-        # range it was built on -- rather than a defect fix. `invokeai/backend/ltx2/sampling.py`
-        # holds its own mu at both anchors on the same reasoning; ComfyUI's default FLUX.1 sampling
-        # uses a fixed shift of 1.15 at every resolution.
-        #
-        # Only the upper anchor is held. Below 256 tokens mu merely sags to 0.457, which is
-        # harmless, and clamping there would cost small frames the interpolation they rely on.
-        #
-        # At or below 4096 tokens this is a no-op -- note that is an *area*, so a 1024x1280 canvas
-        # is already above it.
-        mu = get_lin_function(y1=base_shift, y2=max_shift)(min(image_seq_len, MU_FIT_MAX_SEQ_LEN))
-        timesteps = time_shift(mu, 1.0, timesteps)
+        timesteps = time_shift(FLUX_SHIFT_FIT.shift(image_seq_len), timesteps)
 
     return timesteps.tolist()
 
diff --git a/invokeai/backend/krea2/sampling_utils.py b/invokeai/backend/krea2/sampling_utils.py
index 4c137609fe3..d95db2d1be7 100644
--- a/invokeai/backend/krea2/sampling_utils.py
+++ b/invokeai/backend/krea2/sampling_utils.py
@@ -25,11 +25,21 @@
 KREA2_START_IDX = 34  # drop the system-prompt prefix tokens
 KREA2_NUM_SUFFIX_TOKENS = 5
 
-# Resolution-aware time-shift parameters (scheduler_config.json).
-KREA2_BASE_SHIFT = 0.5
-KREA2_MAX_SHIFT = 1.15
-KREA2_BASE_IMAGE_SEQ_LEN = 256
-KREA2_MAX_IMAGE_SEQ_LEN = 6400
+# The Raw checkpoint's scheduler_config.json, for checkpoints that ship none (see `load_scheduler`). Its shift fit ends
+# at 6400 tokens, 1280x1280 of area, and is held there: extrapolated, mu would be 2.2 at 2048x2048, which puts about
+# half of 28 steps above sigma 0.9, and 7.4 at 4096x4096, which puts all of them there, so an img2img start would have
+# no step near the sigma it asks for. Held, Raw ends at Turbo's fixed mu, which ComfyUI also uses for Krea-2 at every
+# size; diffusers' `Krea2Pipeline` extrapolates. The scheduler applies the exponential time shift to the mu it is
+# handed, so Raw's mu goes to `set_timesteps` as it is.
+KREA2_SCHEDULER_CONFIG = {
+    "use_dynamic_shifting": True,
+    "base_shift": 0.5,
+    "max_shift": 1.15,
+    "base_image_seq_len": 256,
+    "max_image_seq_len": 6400,
+    "num_train_timesteps": 1000,
+    "time_shift_type": "exponential",
+}
 # Fixed timestep shift for the distilled (Turbo) checkpoint.
 KREA2_DISTILLED_MU = 1.15
 
@@ -75,29 +85,6 @@ def prepare_position_ids(text_seq_len: int, grid_height: int, grid_width: int, d
     return torch.cat([text_ids, image_ids], dim=0)
 
 
-def calculate_shift(
-    image_seq_len: int,
-    base_image_seq_len: int = KREA2_BASE_IMAGE_SEQ_LEN,
-    max_image_seq_len: int = KREA2_MAX_IMAGE_SEQ_LEN,
-    base_shift: float = KREA2_BASE_SHIFT,
-    max_shift: float = KREA2_MAX_SHIFT,
-) -> float:
-    """Resolution-aware mu: a line through the two config points, held at the last one.
-
-    Past ``max_image_seq_len`` (6400 tokens, 1280x1280 of area) the line would keep climbing: mu 2.2 at
-    2048x2048 puts about half of 28 steps above sigma 0.9, and mu 7.4 at 4096x4096 puts all of them
-    there, so an img2img start has no step near the sigma it asks for. Held with the stock config, Raw
-    ends at Turbo's fixed mu, which ComfyUI also uses for Krea-2 at every size; diffusers'
-    ``Krea2Pipeline`` extrapolates.
-
-    NOTE: mu is fed straight into ``FlowMatchEulerDiscreteScheduler.set_timesteps(..., mu=mu)``;
-    the exponential time-shift happens inside the scheduler. Do NOT ``exp()`` this value.
-    """
-    m = (max_shift - base_shift) / (max_image_seq_len - base_image_seq_len)
-    b = base_shift - m * base_image_seq_len
-    return min(image_seq_len, max_image_seq_len) * m + b
-
-
 def build_sigmas(steps: int) -> List[float]:
     """Krea-2 sigma schedule: linspace(1.0, 1/steps, steps)."""
     return np.linspace(1.0, 1.0 / steps, steps).tolist()
diff --git a/invokeai/backend/qwen_image_2_1/sampling.py b/invokeai/backend/qwen_image_2_1/sampling.py
index 2a37cae1d8c..babbf2f95c7 100644
--- a/invokeai/backend/qwen_image_2_1/sampling.py
+++ b/invokeai/backend/qwen_image_2_1/sampling.py
@@ -5,6 +5,7 @@
 from diffusers import FlowMatchEulerDiscreteScheduler
 
 from invokeai.backend.model_manager.taxonomy import QwenImage21VariantType
+from invokeai.backend.rectified_flow.schedule_shift import ShiftFit
 
 # scheduler/scheduler_config.json of Qwen/Qwen-Image-2.1.
 BASE_SCHEDULER_CONFIG = {
@@ -30,17 +31,11 @@
 TURBO_SIGMAS = (1.0, 0.978453, 0.95418, 0.926626, 0.89508, 0.845148, 0.704534, 0.414568)
 
 
-def calculate_mu(image_seq_len: int) -> float:
-    """The pipeline's `calculate_shift` over target-image tokens only (condition images are not counted), held at
-    `max_shift` above `max_image_seq_len` (8192 tokens, about 1448x1448).
-
-    The pipeline keeps extrapolating the line, to 1.31 at 2048x2048 and 3.79 at 4096x4096, where image-to-image
-    at strength 0.5 over 40 steps would start at sigma 0.95 and replace the image; held, it starts at 0.70. ComfyUI
-    samples this model at a constant mu of 0.69.
-    """
-    c = BASE_SCHEDULER_CONFIG
-    m = (c["max_shift"] - c["base_shift"]) / (c["max_image_seq_len"] - c["base_image_seq_len"])
-    return min(image_seq_len, c["max_image_seq_len"]) * m + (c["base_shift"] - m * c["base_image_seq_len"])
+# The base model's shift, over target-image tokens only (condition images are not counted), held at `max_shift` above
+# `max_image_seq_len` (8192 tokens, about 1448x1448). The pipeline keeps extrapolating the line, to 1.31 at 2048x2048
+# and 3.79 at 4096x4096, where image-to-image at strength 0.5 over 40 steps would start at sigma 0.95 and replace the
+# image; held, it starts at 0.70. ComfyUI samples this model at a constant mu of 0.69.
+BASE_SHIFT_FIT = ShiftFit.from_scheduler_config(BASE_SCHEDULER_CONFIG)
 
 
 def turbo_sigmas(steps: int) -> list[float]:
@@ -57,7 +52,7 @@ def turbo_sigmas(steps: int) -> list[float]:
 
 def build_sigmas(variant: QwenImage21VariantType | None, steps: int, image_seq_len: int) -> torch.Tensor:
     """The full schedule, `steps + 1` sigmas ending in 0: `QwenImage21Pipeline`'s within the shift's fit, with the
-    base model's shift held above it (see `calculate_mu`).
+    base model's shift held above it (see `BASE_SHIFT_FIT`).
 
     `image_seq_len` is the target image's latent token count, `(height // 16) * (width // 16)`; only the base
     model's shift reads it.
@@ -73,6 +68,6 @@ def build_sigmas(variant: QwenImage21VariantType | None, steps: int, image_seq_l
     else:
         scheduler = FlowMatchEulerDiscreteScheduler.from_config(BASE_SCHEDULER_CONFIG)
         scheduler.set_timesteps(
-            sigmas=np.linspace(1.0, 1 / steps, steps).tolist(), mu=calculate_mu(image_seq_len), device="cpu"
+            sigmas=np.linspace(1.0, 1 / steps, steps).tolist(), mu=BASE_SHIFT_FIT.mu(image_seq_len), device="cpu"
         )
     return scheduler.sigmas.float()
diff --git a/invokeai/backend/rectified_flow/schedule_shift.py b/invokeai/backend/rectified_flow/schedule_shift.py
new file mode 100644
index 00000000000..7f04a7f7ad8
--- /dev/null
+++ b/invokeai/backend/rectified_flow/schedule_shift.py
@@ -0,0 +1,89 @@
+"""The resolution-dependent shift of a flow-matching sigma schedule.
+
+Larger frames are sampled with more of their steps at high noise. How much more is ``mu``, a straight line through
+two (image token count, mu) anchors -- diffusers' ``calculate_shift`` over a scheduler config's
+``base_image_seq_len``/``base_shift`` and ``max_image_seq_len``/``max_shift``. A family declares its anchors as a
+:class:`ShiftFit`, or reads them from its scheduler config, and takes the result in one of two forms:
+
+* ``mu``, for the exponential time shift ``e^mu / (e^mu + 1/t - 1)``. ``FlowMatchEulerDiscreteScheduler`` applies
+  it from ``set_timesteps(mu=...)`` with ``time_shift_type="exponential"``.
+* ``shift = e^mu``, for the linear form ``shift / (shift + 1/t - 1)`` (:func:`time_shift`). It is the same schedule,
+  written for a family that builds its sigmas by hand or takes a shift from the user.
+
+The line is not evaluated past its upper anchor. Extrapolated, it keeps climbing with the frame's area, the schedule
+spends almost every step near pure noise, and an image-to-image pass entered by strength starts so close to noise
+that it replaces its input. Each family's anchors record its own numbers. diffusers' pipelines extrapolate.
+"""
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Mapping, TypeVar
+
+import torch
+
+if TYPE_CHECKING:
+    from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+
+T = TypeVar("T", float, torch.Tensor)
+
+
+@dataclass(frozen=True)
+class ShiftFit:
+    """mu along the line through ``(base_seq_len, base_shift)`` and ``(max_seq_len, max_shift)``, held at its upper
+    end.
+
+    Below ``base_seq_len`` the line keeps running: mu only sags there, and small frames keep the interpolation they
+    were sampled with upstream.
+    """
+
+    base_seq_len: int
+    max_seq_len: int
+    base_shift: float
+    max_shift: float
+
+    @classmethod
+    def from_scheduler_config(cls, config: Mapping[str, Any]) -> "ShiftFit":
+        """The fit a ``FlowMatchEulerDiscreteScheduler`` config declares (see :func:`load_scheduler`)."""
+        return cls(
+            base_seq_len=config["base_image_seq_len"],
+            max_seq_len=config["max_image_seq_len"],
+            base_shift=config["base_shift"],
+            max_shift=config["max_shift"],
+        )
+
+    def mu(self, seq_len: int) -> float:
+        # diffusers' `calculate_shift`, in its order of operations, on the held sequence length.
+        m = (self.max_shift - self.base_shift) / (self.max_seq_len - self.base_seq_len)
+        b = self.base_shift - m * self.base_seq_len
+        return min(seq_len, self.max_seq_len) * m + b
+
+    def shift(self, seq_len: int) -> float:
+        """``e^mu``, the shift the linear form takes."""
+        return math.exp(self.mu(seq_len))
+
+
+def time_shift(shift: float, t: T) -> T:
+    """The linear time shift: ``t`` moved towards 1 by ``shift``, ``shift / (shift + 1/t - 1)``.
+
+    For a positive ``shift``, ``t = 1`` stays 1; a shift of 0 divides zero by zero there. A tensor ``t`` of 0 stays 0;
+    a float one divides by zero.
+    """
+    return shift / (shift + (1 / t - 1))
+
+
+def load_scheduler(model_path: Path, stock_config: Mapping[str, Any]) -> "FlowMatchEulerDiscreteScheduler":
+    """The checkpoint's own ``scheduler/scheduler_config.json``, else a scheduler built from ``stock_config``.
+
+    GGUF and single-file transformers have no ``scheduler/`` folder and get the stock config. A checkpoint's file is
+    read the way the reference pipeline reads it: the scheduler registers every init argument, so a key the file
+    omits takes ``FlowMatchEulerDiscreteScheduler``'s own default (``max_shift`` 1.15, ``max_image_seq_len`` 4096),
+    not the family's stock value.
+    """
+    # Imported on use: the families that only evaluate a fit do not need diffusers' schedulers.
+    from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+
+    scheduler_path = model_path / "scheduler"
+    if scheduler_path.is_dir() and (scheduler_path / "scheduler_config.json").exists():
+        return FlowMatchEulerDiscreteScheduler.from_pretrained(str(scheduler_path), local_files_only=True)
+    return FlowMatchEulerDiscreteScheduler(**stock_config)
diff --git a/invokeai/frontend/webv2/src/features/upscale/core/graph.test.ts b/invokeai/frontend/webv2/src/features/upscale/core/graph.test.ts
index 157dedbb9c3..b8d0cf057d2 100644
--- a/invokeai/frontend/webv2/src/features/upscale/core/graph.test.ts
+++ b/invokeai/frontend/webv2/src/features/upscale/core/graph.test.ts
@@ -478,7 +478,7 @@ describe('compileUpscaleGraph for Krea-2', () => {
 
 describe('Creativity on index-clipping architectures', () => {
   // Expected starts come from the backend's own schedules (`ZImageDenoiseInvocation._get_sigmas`
-  // with `_calculate_shift`, `AnimaDenoiseInvocation._get_sigmas`) and the nodes' clipping line,
+  // with `Z_IMAGE_SHIFT_FIT`, `AnimaDenoiseInvocation._get_sigmas`) and the nodes' clipping line,
   // `int(start * steps)`: each enters within one step of the sigma FLUX enters at, 1 - start.
   const startFor = (values: UpscaleWidgetValues, creativity: number, inputSide: number) =>
     compileUpscaleGraph(
diff --git a/invokeai/frontend/webv2/src/features/upscale/core/settings.ts b/invokeai/frontend/webv2/src/features/upscale/core/settings.ts
index a2505e81d82..1449ded0848 100644
--- a/invokeai/frontend/webv2/src/features/upscale/core/settings.ts
+++ b/invokeai/frontend/webv2/src/features/upscale/core/settings.ts
@@ -39,7 +39,7 @@ import type { SpandrelModelConfig, TileControlNetModelConfig, UpscaleWidgetValue
 // `indexClipping`; `getIndexClippedDenoisingStart` in graph.ts uses them to start those nodes where
 // FLUX would. They go away when those nodes clip by value.
 
-/** `z_image_denoise._calculate_shift`, held at the fit's 4096-token end since the clamp (#378). */
+/** `z_image_denoise.Z_IMAGE_SHIFT_FIT`, held at the fit's 4096-token end since the clamp (#378). */
 const Z_IMAGE_SHIFT_MIN_SEQ_LEN = 256;
 const Z_IMAGE_SHIFT_MAX_SEQ_LEN = 4096;
 const Z_IMAGE_BASE_SHIFT = 0.5;
diff --git a/tests/app/invocations/test_denoise_noise_inputs.py b/tests/app/invocations/test_denoise_noise_inputs.py
index 500d8b94600..e90e2f530ec 100644
--- a/tests/app/invocations/test_denoise_noise_inputs.py
+++ b/tests/app/invocations/test_denoise_noise_inputs.py
@@ -18,7 +18,7 @@
 )
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.invocations.sd3.sd3_denoise import SD3DenoiseInvocation
-from invokeai.app.invocations.z_image.z_image_denoise import ZImageDenoiseInvocation
+from invokeai.app.invocations.z_image.z_image_denoise import Z_IMAGE_SHIFT_FIT, ZImageDenoiseInvocation
 from invokeai.backend.flux.sampling_utils import clip_timestep_schedule_fractional, get_schedule
 from invokeai.backend.flux.schedulers import (
     ANIMA_SCHEDULER_MAP,
@@ -1058,7 +1058,7 @@ def forward(self, x: list[torch.Tensor], t: torch.Tensor, cap_feats: list[torch.
 def test_z_image_img2img_preblend_matches_scheduler_first_sigma(scheduler_name: str):
     invocation = ZImageDenoiseInvocation.model_construct(steps=8, width=1024, height=1024)
     img_seq_len = (invocation.height // 8 // 2) * (invocation.width // 8 // 2)
-    shift = invocation._calculate_shift(img_seq_len)
+    shift = Z_IMAGE_SHIFT_FIT.shift(img_seq_len)
     sigmas = invocation._get_sigmas(shift, invocation.steps)
     sigmas = sigmas[int(0.25 * (len(sigmas) - 1)) :]
     scheduler_class = ZIMAGE_SCHEDULER_MAP[scheduler_name]
diff --git a/tests/app/invocations/test_krea2_denoise.py b/tests/app/invocations/test_krea2_denoise.py
index d1cc6bbfa36..1e0198b7a3c 100644
--- a/tests/app/invocations/test_krea2_denoise.py
+++ b/tests/app/invocations/test_krea2_denoise.py
@@ -1,3 +1,4 @@
+import json
 import math
 from contextlib import contextmanager, nullcontext
 from types import SimpleNamespace
@@ -5,6 +6,7 @@
 
 import pytest
 import torch
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 
 from invokeai.app.invocations.fields import (
     DenoiseMaskField,
@@ -737,7 +739,7 @@ def __call__(self, *, hidden_states, encoder_hidden_states, position_ids, **_kwa
 
 
 def _mu_branch_invocation(shift: float | None = None) -> Krea2DenoiseInvocation:
-    # shift=None so the resolution-aware calculate_shift branch runs (not the fixed-mu override); width/height
+    # shift=None so the resolution-aware mu branch runs (not the fixed-mu override); width/height
     # 16 -> a 2x2 latent -> 1x1 grid -> image_seq_len == 1.
     return Krea2DenoiseInvocation.model_construct(
         transformer=TransformerField(transformer=_model_identifier(), loras=[]),
@@ -757,7 +759,7 @@ def _mu_branch_invocation(shift: float | None = None) -> Krea2DenoiseInvocation:
 
 
 def _mu_branch_context(tmp_path, transformer: _Transformer, variant=Krea2VariantType.Base):
-    # variant=Base -> _is_distilled() is False -> the resolution-aware mu (calculate_shift) branch is taken.
+    # variant=Base -> _is_distilled() is False -> the resolution-aware mu branch is taken.
     # Pass variant=Turbo to exercise the distilled fixed-mu branch instead.
     conditionings = {
         "positive": ConditioningFieldData(conditionings=[Krea2ConditioningInfo(prompt_embeds=torch.ones(1, 2, 12, 8))]),
@@ -775,16 +777,9 @@ def _mu_branch_context(tmp_path, transformer: _Transformer, variant=Krea2Variant
     )
 
 
-def _install_mu_capturing_scheduler(
-    monkeypatch, captured: list[float], config_kwargs: dict, stop: bool = False
-) -> None:
-    class _MuCapturingScheduler:
-        def __init__(self, **_kwargs) -> None:
-            # The real diffusers scheduler exposes shift params on .config (a FrozenDict, attribute-access);
-            # the denoise reads them via getattr(scheduler.config, <key>, <krea-2 default>). Omitting a key
-            # here (SimpleNamespace without the attr) exercises that fallback, which only the stock scheduler
-            # reaches: a real scheduler_config.json missing a key gets the diffusers scheduler's default.
-            self.config = SimpleNamespace(num_train_timesteps=1000, **config_kwargs)
+def _install_mu_capturing_scheduler(monkeypatch, captured: list[float], stop: bool = False) -> None:
+    class _MuCapturingScheduler(FlowMatchEulerDiscreteScheduler):
+        """The real scheduler, so a checkpoint's scheduler_config.json loads as it would in production."""
 
         def set_timesteps(self, *, sigmas, mu, device) -> None:
             captured.append(mu)
@@ -800,25 +795,6 @@ def set_timesteps(self, *, sigmas, mu, device) -> None:
     )
 
 
-def test_run_diffusion_mu_honors_scheduler_config_shift_params(monkeypatch, tmp_path) -> None:
-    # Regression: the undistilled (Raw) resolution-aware mu path must derive its shift params from the loaded
-    # scheduler's config so a checkpoint shipping a customized scheduler_config.json samples with its own
-    # shift, rather than the hardcoded Krea-2 defaults.
-    from invokeai.backend.krea2.sampling_utils import calculate_shift
-
-    _patch_runtime(monkeypatch)  # patches TorchDevice + LayerPatcher (and the default _Scheduler)...
-    captured: list[float] = []
-    custom = {"base_image_seq_len": 100, "max_image_seq_len": 1000, "base_shift": 0.3, "max_shift": 1.5}
-    _install_mu_capturing_scheduler(monkeypatch, captured, custom)  # ...then override the scheduler.
-
-    _mu_branch_invocation()._run_diffusion(_mu_branch_context(tmp_path, _Transformer()))
-
-    image_seq_len = 1  # width=height=16
-    assert captured == [pytest.approx(calculate_shift(image_seq_len, **custom))]
-    # Sanity: the custom config actually moves mu away from the stock-default result.
-    assert captured[0] != pytest.approx(calculate_shift(image_seq_len))
-
-
 class _StopAtTimesteps(Exception):
     """Ends a run once the scheduler has been handed its mu."""
 
@@ -831,57 +807,56 @@ class _StopAtTimesteps(Exception):
         (2048, 768, 1.1229166666666667),  # 6144 tokens: area decides, not a side
         (1280, 1280, 1.15),  # 6400 tokens: the line's end
         (2048, 2048, 1.15),  # 16384 tokens: would be 2.206 extrapolated
-        (4096, 4096, 1.15),  # 65536 tokens: would be 7.41
     ],
 )
-def test_raw_mu_follows_the_fit_up_to_its_end_and_holds_there(width: int, height: int, mu: float) -> None:
-    from invokeai.backend.krea2.sampling_utils import calculate_shift
-
-    # One token per 16x16 pixels, as `_run_diffusion` counts them.
-    assert calculate_shift((width // 16) * (height // 16)) == pytest.approx(mu)
-
-
-def test_raw_mu_holds_at_the_end_of_a_custom_scheduler_fit(monkeypatch, tmp_path) -> None:
-    # A checkpoint's own scheduler_config.json moves the end of the fit, and the hold moves with it.
+def test_raw_mu_follows_the_stock_fit_up_to_its_end_and_holds_there(monkeypatch, tmp_path, width, height, mu) -> None:
+    # The model folder has no scheduler/, as for a GGUF or single-file transformer, so the stock scheduler is used.
     _patch_runtime(monkeypatch)
     captured: list[float] = []
-    custom = {"base_image_seq_len": 256, "max_image_seq_len": 1024, "base_shift": 0.3, "max_shift": 1.5}
-    _install_mu_capturing_scheduler(monkeypatch, captured, custom, stop=True)
+    _install_mu_capturing_scheduler(monkeypatch, captured, stop=True)
     invocation = _mu_branch_invocation()
-    invocation.width = invocation.height = 2048
+    invocation.width, invocation.height = width, height
 
     with pytest.raises(_StopAtTimesteps):
         invocation._run_diffusion(_mu_branch_context(tmp_path, _Transformer()))
 
-    assert captured == [pytest.approx(1.5)]
+    assert captured == [pytest.approx(mu)]
 
 
-def test_run_diffusion_mu_falls_back_to_krea2_defaults_for_absent_config_keys(monkeypatch, tmp_path) -> None:
-    # A scheduler config that omits some shift params must fall back to the Krea-2 defaults for exactly the
-    # missing keys, honoring the ones it does provide.
-    from invokeai.backend.krea2.sampling_utils import (
-        KREA2_BASE_SHIFT,
-        KREA2_MAX_IMAGE_SEQ_LEN,
-        calculate_shift,
+@pytest.mark.parametrize(
+    ("size", "mu"),
+    [
+        (16, 0.168),  # 1 token: 0.3 + 1.2 * (1 - 100) / 900, below this fit's lower end
+        (2048, 1.5),  # 16384 tokens: held at this config's max_shift
+    ],
+)
+def test_raw_mu_follows_and_holds_at_the_end_of_a_checkpoints_own_scheduler_fit(
+    monkeypatch, tmp_path, size: int, mu: float
+) -> None:
+    (tmp_path / "scheduler").mkdir()
+    (tmp_path / "scheduler" / "scheduler_config.json").write_text(
+        json.dumps(
+            {
+                "_class_name": "FlowMatchEulerDiscreteScheduler",
+                "use_dynamic_shifting": True,
+                "base_image_seq_len": 100,
+                "max_image_seq_len": 1000,
+                "base_shift": 0.3,
+                "max_shift": 1.5,
+                "time_shift_type": "exponential",
+            }
+        )
     )
-
     _patch_runtime(monkeypatch)
     captured: list[float] = []
-    # Provide only base_image_seq_len + max_shift; base_shift and max_image_seq_len must default.
-    partial = {"base_image_seq_len": 100, "max_shift": 1.5}
-    _install_mu_capturing_scheduler(monkeypatch, captured, partial)
-
-    _mu_branch_invocation()._run_diffusion(_mu_branch_context(tmp_path, _Transformer()))
-
-    image_seq_len = 1
-    expected = calculate_shift(
-        image_seq_len,
-        base_image_seq_len=100,
-        max_image_seq_len=KREA2_MAX_IMAGE_SEQ_LEN,
-        base_shift=KREA2_BASE_SHIFT,
-        max_shift=1.5,
-    )
-    assert captured == [pytest.approx(expected)]
+    _install_mu_capturing_scheduler(monkeypatch, captured, stop=True)
+    invocation = _mu_branch_invocation()
+    invocation.width = invocation.height = size
+
+    with pytest.raises(_StopAtTimesteps):
+        invocation._run_diffusion(_mu_branch_context(tmp_path, _Transformer()))
+
+    assert captured == [pytest.approx(mu)]
 
 
 @pytest.mark.parametrize("shift", [0.0, -1.0])
@@ -896,7 +871,7 @@ def test_run_diffusion_treats_a_non_positive_shift_as_unset_on_turbo(monkeypatch
 
     _patch_runtime(monkeypatch)
     captured: list[float] = []
-    _install_mu_capturing_scheduler(monkeypatch, captured, {})
+    _install_mu_capturing_scheduler(monkeypatch, captured)
 
     _mu_branch_invocation(shift=shift)._run_diffusion(
         _mu_branch_context(tmp_path, _Transformer(), variant=Krea2VariantType.Turbo)
@@ -907,22 +882,20 @@ def test_run_diffusion_treats_a_non_positive_shift_as_unset_on_turbo(monkeypatch
 
 def test_run_diffusion_treats_a_zero_shift_as_unset_on_raw(monkeypatch, tmp_path) -> None:
     # The same bypass must fall through to the resolution-aware branch for an undistilled checkpoint.
-    from invokeai.backend.krea2.sampling_utils import calculate_shift
-
     _patch_runtime(monkeypatch)
     captured: list[float] = []
-    _install_mu_capturing_scheduler(monkeypatch, captured, {})
+    _install_mu_capturing_scheduler(monkeypatch, captured)
 
     _mu_branch_invocation(shift=0.0)._run_diffusion(_mu_branch_context(tmp_path, _Transformer()))
 
-    assert captured == [pytest.approx(calculate_shift(1))]
+    assert captured == [pytest.approx(0.5 + 0.65 * (1 - 256) / 6144)]  # 1 token on the stock fit
 
 
 def test_run_diffusion_still_honors_an_explicit_positive_shift(monkeypatch, tmp_path) -> None:
     # The override itself must keep working - only non-positive values are reinterpreted as "unset".
     _patch_runtime(monkeypatch)
     captured: list[float] = []
-    _install_mu_capturing_scheduler(monkeypatch, captured, {})
+    _install_mu_capturing_scheduler(monkeypatch, captured)
 
     _mu_branch_invocation(shift=1.8)._run_diffusion(
         _mu_branch_context(tmp_path, _Transformer(), variant=Krea2VariantType.Turbo)
diff --git a/tests/app/invocations/test_qwen_image_denoise.py b/tests/app/invocations/test_qwen_image_denoise.py
index 98a28366bb8..62a524cb3e5 100644
--- a/tests/app/invocations/test_qwen_image_denoise.py
+++ b/tests/app/invocations/test_qwen_image_denoise.py
@@ -2,18 +2,10 @@
 
 import json
 from pathlib import Path
-from types import SimpleNamespace
 
 import pytest
-import torch
-from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 
-from invokeai.app.invocations.qwen_image import qwen_image_denoise
 from invokeai.app.invocations.qwen_image.qwen_image_denoise import QwenImageDenoiseInvocation
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
-    ConditioningFieldData,
-    QwenImageConditioningInfo,
-)
 
 
 class TestPrepareCfgScale:
@@ -208,88 +200,29 @@ def test_edit_matches_diffusers_layout(self):
         assert result[0][1] == (1, 64, 48)
 
 
-class _StopAtTimesteps(Exception):
-    """Ends a run once the scheduler has been handed its mu."""
-
-
-class _TinyQwenImageTransformer:
-    """Stands in for `QwenImageTransformer2DModel`: the config the node reads before it sets the schedule."""
-
-    config = SimpleNamespace(patch_size=2, out_channels=16)
-
-
-def _run_to_scheduler(monkeypatch, model_path: Path, width: int, height: int, shift: float | None = None) -> float:
-    """Run `_run_diffusion` with the real scheduler up to `set_timesteps` and return the mu it was handed."""
-    monkeypatch.setattr(qwen_image_denoise, "QwenImageTransformer2DModel", _TinyQwenImageTransformer)
-    monkeypatch.setattr(qwen_image_denoise.TorchDevice, "choose_torch_device", lambda: torch.device("cpu"))
-    captured: list[float] = []
-
-    def capture_mu(self, *, sigmas, mu, device) -> None:
-        captured.append(mu)
-        # A frame of real size need not be denoised to know the mu it was handed.
-        raise _StopAtTimesteps
-
-    monkeypatch.setattr(FlowMatchEulerDiscreteScheduler, "set_timesteps", capture_mu)
-    conditioning = ConditioningFieldData(conditionings=[QwenImageConditioningInfo(prompt_embeds=torch.zeros(1, 4, 8))])
-    context = SimpleNamespace(
-        models=SimpleNamespace(
-            load=lambda _identifier: SimpleNamespace(model=_TinyQwenImageTransformer()),
-            get_config=lambda _identifier: SimpleNamespace(),
-            get_absolute_path=lambda _config: model_path,
-        ),
-        conditioning=SimpleNamespace(load=lambda _name: conditioning),
-    )
-    node = QwenImageDenoiseInvocation.model_construct(
-        latents=None,
-        reference_latents=None,
-        denoise_mask=None,
-        denoising_start=0.0,
-        denoising_end=1.0,
-        transformer=SimpleNamespace(transformer=SimpleNamespace(), loras=[]),
-        positive_conditioning=SimpleNamespace(conditioning_name="positive"),
-        negative_conditioning=None,
-        cfg_scale=1.0,
-        width=width,
-        height=height,
-        steps=40,
-        seed=0,
-        shift=shift,
-    )
-
-    with pytest.raises(_StopAtTimesteps):
-        node._run_diffusion(context)
-
-    assert len(captured) == 1
-    return captured[0]
-
-
 @pytest.mark.parametrize(
     ("width", "height", "mu"),
     [
-        # One token per 16x16 pixels. The stock fit is mu = 0.5 + 0.4 * (tokens - 256) / 7936.
+        # One token per 16x16 pixels (8x VAE, 2x2 patches). The stock fit is mu = 0.5 + 0.4 * (tokens - 256) / 7936.
+        (16, 16, 0.5 - 0.4 * 255 / 7936),  # 1 token: below the lower anchor the line keeps running
         (1024, 1024, 0.6935483870967742),  # 4096 tokens: 0.5 + 0.4 * 3840 / 7936
         (2048, 1008, 0.8935483870967742),  # 8064 tokens: 0.5 + 0.4 * 7808 / 7936; area decides, not a side
         (2048, 1024, 0.9),  # 8192 tokens: the fit's end
-        (2048, 2048, 0.9),  # 16384 tokens: 1.313 extrapolated
         (4096, 4096, 0.9),  # 65536 tokens: 3.790 extrapolated
     ],
 )
-def test_stock_mu_follows_the_fit_up_to_its_end_and_holds_there(monkeypatch, tmp_path, width, height, mu) -> None:
-    # A GGUF or single-file transformer has no scheduler/ folder, so the node builds the stock Qwen-Image scheduler.
-    assert _run_to_scheduler(monkeypatch, tmp_path / "qwen-image.gguf", width, height) == pytest.approx(mu)
+def test_without_a_scheduler_folder_mu_follows_the_stock_fit_and_holds_at_its_end(
+    tmp_path: Path, width: int, height: int, mu: float
+) -> None:
+    # A GGUF or single-file transformer has no scheduler/ folder.
+    _, actual = QwenImageDenoiseInvocation._schedule(tmp_path / "qwen-image.gguf", width, height, 2, None)
 
+    assert actual == pytest.approx(mu, abs=1e-12)
 
-@pytest.mark.parametrize(
-    ("size", "mu"),
-    [
-        (384, 0.8),  # 576 tokens: 0.3 + 1.2 * 320 / 768
-        (2048, 1.5),  # 16384 tokens: held at this config's max_shift
-    ],
-)
-def test_mu_follows_and_holds_at_the_end_of_a_checkpoints_own_scheduler_fit(monkeypatch, tmp_path, size, mu) -> None:
-    scheduler_dir = tmp_path / "scheduler"
-    scheduler_dir.mkdir()
-    (scheduler_dir / "scheduler_config.json").write_text(
+
+def _with_scheduler_config(model_path: Path) -> Path:
+    (model_path / "scheduler").mkdir(parents=True)
+    (model_path / "scheduler" / "scheduler_config.json").write_text(
         json.dumps(
             {
                 "_class_name": "FlowMatchEulerDiscreteScheduler",
@@ -303,11 +236,25 @@ def test_mu_follows_and_holds_at_the_end_of_a_checkpoints_own_scheduler_fit(monk
             }
         )
     )
+    return model_path
+
+
+@pytest.mark.parametrize(
+    ("size", "mu"),
+    [
+        (384, 0.8),  # 576 tokens: 0.3 + 1.2 * 320 / 768
+        (2048, 1.5),  # 16384 tokens: held at this config's max_shift
+    ],
+)
+def test_a_checkpoints_own_scheduler_config_sets_the_fit(tmp_path: Path, size: int, mu: float) -> None:
+    scheduler, actual = QwenImageDenoiseInvocation._schedule(_with_scheduler_config(tmp_path), size, size, 2, None)
 
-    assert _run_to_scheduler(monkeypatch, tmp_path, size, size) == pytest.approx(mu)
+    assert actual == pytest.approx(mu, abs=1e-12)
+    assert scheduler.config.max_shift == 1.5
 
 
-def test_an_explicit_shift_sets_mu_at_any_size(monkeypatch, tmp_path) -> None:
-    mu = _run_to_scheduler(monkeypatch, tmp_path / "qwen-image.gguf", 4096, 4096, shift=3.0)
+def test_an_explicit_shift_sets_mu_at_any_size_with_the_checkpoints_scheduler(tmp_path: Path) -> None:
+    scheduler, mu = QwenImageDenoiseInvocation._schedule(_with_scheduler_config(tmp_path), 4096, 4096, 2, 3.0)
 
-    assert mu == pytest.approx(1.0986122886681098)  # ln 3
+    assert mu == pytest.approx(1.0986122886681098, abs=1e-12)  # ln 3
+    assert scheduler.config.max_shift == 1.5
diff --git a/tests/app/invocations/test_z_image_schedule_shift.py b/tests/app/invocations/test_z_image_schedule_shift.py
index 4027b2bd211..b48356cb01f 100644
--- a/tests/app/invocations/test_z_image_schedule_shift.py
+++ b/tests/app/invocations/test_z_image_schedule_shift.py
@@ -5,7 +5,7 @@
 `denoising_start` is clipped by index here, the step count survives and the entry point moves
 instead, so a pass asked to refine an image starts from noise and generates a new one.
 
-The cells that matter drive the real `_run_diffusion`: the real `_calculate_shift`, the real
+The cells that matter drive the real `_run_diffusion`: the real shift fit, the real
 `_get_sigmas`, the real clipping and the real img2img blend all run, and the entry sigma is read
 back out of the blended tensor rather than recomputed here. A test that recomputed the arithmetic
 would agree with a broken node as readily as with a working one.
@@ -17,10 +17,7 @@
 import pytest
 import torch
 
-from invokeai.app.invocations.z_image.z_image_denoise import (
-    Z_IMAGE_MU_FIT_MAX_SEQ_LEN,
-    ZImageDenoiseInvocation,
-)
+from invokeai.app.invocations.z_image.z_image_denoise import Z_IMAGE_SHIFT_FIT, ZImageDenoiseInvocation
 
 # Canvas img2img is the path that reaches this node today. `canvasDenoisingStart` sends
 # `1 - strength` for Z-Image (`features/generation/core/canvas/compileCanvasGraph.ts`), so a
@@ -128,22 +125,26 @@ def capture(_context, latents):
 @pytest.mark.parametrize(
     ["px", "expected_shift"],
     [
+        (16, 1.579),  # 1 token: below the lower point the line keeps running, e^(0.5 - 0.65 * 255 / 3840)
         (512, 1.878),  # interpolation, below the fit's upper point
         (1024, 3.158),  # 4096 tokens: the upper point, where the fit ends
-        (1536, 3.158),  # beyond it, held -- extrapolation would give 7.5
-        (2048, 3.158),  # extrapolation would give 25.3
-        (3072, 3.158),  # extrapolation would give 809.7
-        (4096, 3.158),  # extrapolation would give 103_777
+        (2048, 3.158),  # beyond it, held -- extrapolation would give 25.3
     ],
 )
-def test_the_shift_is_not_extrapolated_past_the_fit(px: int, expected_shift: float):
-    assert invocation()._calculate_shift(seq_len(px)) == pytest.approx(expected_shift, abs=1e-3)
+def test_the_shift_follows_its_anchors_and_holds_past_the_fit(px: int, expected_shift: float):
+    assert Z_IMAGE_SHIFT_FIT.shift(seq_len(px)) == pytest.approx(expected_shift, abs=1e-3)
 
 
-def test_the_fit_is_still_live_below_its_upper_point():
-    """Holding the fit at its endpoint is not the same as replacing it with a constant."""
-    node = invocation()
-    assert node._calculate_shift(256) < node._calculate_shift(1024) < node._calculate_shift(Z_IMAGE_MU_FIT_MAX_SEQ_LEN)
+@pytest.mark.parametrize(
+    ["shift", "num_steps", "expected"],
+    [
+        (3.0, 1, [1.0, 0.0]),  # one step is the whole range
+        (3.0, 4, [1.0, 0.9, 0.75, 0.5, 0.0]),  # 3 / (3 + 1/t - 1) at t = 0.75, 0.5, 0.25
+        (0.0, 3, [1.0, 0.0, 0.0, 0.0]),  # a shift of 0 keeps the ends and flattens the rest
+    ],
+)
+def test_the_sigmas_keep_their_ends_at_any_shift(shift: float, num_steps: int, expected: list[float]):
+    assert invocation()._get_sigmas(shift, num_steps) == pytest.approx(expected, abs=1e-12)
 
 
 @pytest.mark.parametrize("px", [1024, 1536, 2048, 3072, 4096])
@@ -170,8 +171,8 @@ def test_the_entry_sigma_does_not_depend_on_the_step_count():
 def test_an_explicit_shift_overrides_the_clamp():
     """The node exposes `shift` so a user can ask for the pre-clamp behaviour back.
 
-    Both halves go through `_run_diffusion`, so deleting the override in favour of always calling
-    `_calculate_shift` fails here rather than passing on a recomputed schedule.
+    Both halves go through `_run_diffusion`, so deleting the override in favour of always using
+    the fit fails here rather than passing on a recomputed schedule.
     """
     auto_shift, auto_entry = drive_to_blend(2048, 30)
     forced_shift, forced_entry = drive_to_blend(2048, 30, shift=UNCLAMPED_SHIFT_AT_2048)
diff --git a/tests/backend/flux/test_sampling_utils.py b/tests/backend/flux/test_sampling_utils.py
index 30f38c4dd29..a8807dac4d7 100644
--- a/tests/backend/flux/test_sampling_utils.py
+++ b/tests/backend/flux/test_sampling_utils.py
@@ -4,11 +4,8 @@
 import torch
 
 from invokeai.backend.flux.sampling_utils import (
-    MU_FIT_MAX_SEQ_LEN,
-    MU_FIT_MIN_SEQ_LEN,
     clip_timestep_schedule,
     clip_timestep_schedule_fractional,
-    get_lin_function,
     get_schedule,
 )
 
@@ -88,10 +85,10 @@ def test_clip_timestep_schedule_fractional(
 def mu_of(image_seq_len: int) -> float:
     """Recover the shift the schedule was built with, through the public function.
 
-    `get_schedule(2, n)` is [1.0, x, 0.0] with x = exp(mu) / (exp(mu) + 1), because `time_shift`'s
-    `(1 / t - 1) ** sigma` term is 1 at t = 0.5 for every sigma. That makes the readback exact but
-    blind to the *shape* of the shift -- `test_the_shift_has_the_shape_this_readback_assumes` pins
-    that separately, and without it a changed sigma would slip past every cell below.
+    `get_schedule(2, n)` is [1.0, x, 0.0] with x = exp(mu) / (exp(mu) + 1), because `1 / t - 1` is 1 at
+    t = 0.5, so any power of it is too. That makes the readback exact but blind to the *shape* of the
+    shift -- `test_the_shift_has_the_shape_this_readback_assumes` pins that separately, and without it a
+    changed form would slip past every cell below.
     """
     x = get_schedule(num_steps=2, image_seq_len=image_seq_len, shift=True)[1]
     return math.log(x / (1.0 - x))
@@ -104,50 +101,30 @@ def test_the_shift_has_the_shape_this_readback_assumes():
     give 0.966 and 0.260 instead, while t = 0.5 would be unchanged.
     """
     assert float_lists_almost_equal(
-        get_schedule(num_steps=4, image_seq_len=MU_FIT_MAX_SEQ_LEN, shift=True),
+        get_schedule(num_steps=4, image_seq_len=4096, shift=True),
         [1.0, 0.904531, 0.759511, 0.512844, 0.0],
         tol=1e-5,
     )
 
 
-def test_the_anchors_are_where_the_clamp_assumes_they_are():
-    """Written as literals on purpose.
-
-    Asserting `line(MU_FIT_MAX_SEQ_LEN) == 1.15` would hold for any anchor value, because the
-    constants are also `get_lin_function`'s own defaults -- it would pass with the anchor moved to
-    8192, which is exactly the drift it is supposed to catch.
-    """
-    assert MU_FIT_MIN_SEQ_LEN == 256
-    assert MU_FIT_MAX_SEQ_LEN == 4096
-    assert get_lin_function()(4096) == pytest.approx(1.15)
-
-
 @pytest.mark.parametrize(
     ["image_seq_len", "expected_mu"],
     [
-        (MU_FIT_MIN_SEQ_LEN, 0.5),  # 256px: the lower point of the fit
-        (1024, 0.63),  # 512px: interpolation, unaffected
-        (MU_FIT_MAX_SEQ_LEN, 1.15),  # 1024px: the upper point, where the fit ends
-        (9216, 1.15),  # 1536px: beyond the fit, held
-        (16384, 1.15),  # 2048px: extrapolation would give 3.23
-        (36864, 1.15),  # 3072px: extrapolation would give 6.70
-        (65536, 1.15),  # 4096px: extrapolation would give 11.55
+        (1, 0.5 - 0.65 * 255 / 3840),  # 16px: below the lower point the line keeps running (0.457)
+        (256, 0.5),  # 256px: the lower point of the fit
+        (1024, 0.63),  # 512px: interpolation
+        (4096, 1.15),  # 1024px: the upper point, where the fit ends
+        (16384, 1.15),  # 2048px: held; extrapolation would give 3.23
     ],
 )
-def test_the_shift_is_not_extrapolated_past_the_fit(image_seq_len: int, expected_mu: float):
+def test_the_shift_follows_its_anchors_and_holds_past_the_fit(image_seq_len: int, expected_mu: float):
     assert mu_of(image_seq_len) == pytest.approx(expected_mu, abs=1e-3)
 
 
-def test_the_fit_is_still_live_below_its_upper_point():
-    """Holding the fit at its endpoint is not the same as replacing it with a constant, and this
-    names that choice: below one megapixel the schedule still adapts to the frame."""
-    assert mu_of(512) < mu_of(1024) < mu_of(2048) < mu_of(MU_FIT_MAX_SEQ_LEN)
-
-
 @pytest.mark.parametrize(
     ["image_seq_len", "why"],
     [
-        (MU_FIT_MAX_SEQ_LEN, "1024px: the anchor, a control -- unchanged by the clamp either way"),
+        (4096, "1024px: the anchor, a control -- unchanged by the clamp either way"),
         (16384, "2048px: 2 of 30 steps survived before the clamp"),
         (36864, "3072px: 1 of 30 survived before the clamp"),
     ],
diff --git a/tests/backend/qwen_image_2_1/test_sampling.py b/tests/backend/qwen_image_2_1/test_sampling.py
index b2478fc0e08..7f8d3fdfc6f 100644
--- a/tests/backend/qwen_image_2_1/test_sampling.py
+++ b/tests/backend/qwen_image_2_1/test_sampling.py
@@ -10,7 +10,7 @@
 import torch
 
 from invokeai.backend.model_manager.taxonomy import QwenImage21VariantType
-from invokeai.backend.qwen_image_2_1.sampling import TURBO_SIGMAS, build_sigmas, calculate_mu
+from invokeai.backend.qwen_image_2_1.sampling import BASE_SHIFT_FIT, TURBO_SIGMAS, build_sigmas
 
 BASE = QwenImage21VariantType.Base
 TURBO = QwenImage21VariantType.Turbo
@@ -55,7 +55,7 @@ def test_above_its_fit_the_base_schedule_holds_the_shift_at_the_fits_end(
 
 def test_just_below_the_fits_end_the_shift_still_follows_the_line() -> None:
     # 1440x1440 is 90x90 = 8100 tokens, 7844 of the line's 7936 tokens from 0.5 towards 0.9.
-    assert calculate_mu(8100) == pytest.approx(0.5 + 0.4 * 7844 / 7936, abs=1e-12)
+    assert BASE_SHIFT_FIT.mu(8100) == pytest.approx(0.5 + 0.4 * 7844 / 7936, abs=1e-12)
     sigmas = build_sigmas(BASE, 40, 8100)
     torch.testing.assert_close(sigmas[:3], torch.tensor([1.0, 0.989205, 0.978081]), atol=1e-6, rtol=0)
     torch.testing.assert_close(sigmas[-3:], torch.tensor([0.077366, 0.02, 0.0]), atol=1e-6, rtol=0)
diff --git a/tests/backend/rectified_flow/test_schedule_shift.py b/tests/backend/rectified_flow/test_schedule_shift.py
new file mode 100644
index 00000000000..c258e435b72
--- /dev/null
+++ b/tests/backend/rectified_flow/test_schedule_shift.py
@@ -0,0 +1,101 @@
+"""The shared schedule-shift fit: its line, its hold, its two forms, and where it reads its anchors from."""
+
+import json
+import math
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+
+from invokeai.backend.rectified_flow.schedule_shift import ShiftFit, load_scheduler, time_shift
+
+ONE_MEGAPIXEL_FIT = ShiftFit(base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.15)
+# A checkpoint's config may make the line fall. Holding the sequence length, not mu, keeps the hold at its end.
+FALLING_FIT = ShiftFit(base_seq_len=256, max_seq_len=1024, base_shift=1.5, max_shift=0.3)
+
+
+@pytest.mark.parametrize(
+    ("fit", "seq_len", "mu"),
+    [
+        (ONE_MEGAPIXEL_FIT, 1, 0.5 - 0.65 * 255 / 3840),  # below the lower anchor the line keeps running
+        (ONE_MEGAPIXEL_FIT, 256, 0.5),
+        (ONE_MEGAPIXEL_FIT, 1024, 0.63),
+        (ONE_MEGAPIXEL_FIT, 4095, 0.5 + 0.65 * 3839 / 3840),
+        (ONE_MEGAPIXEL_FIT, 4096, 1.15),
+        (ONE_MEGAPIXEL_FIT, 4097, 1.15),  # held from the upper anchor on
+        (ONE_MEGAPIXEL_FIT, 65536, 1.15),  # 11.55 extrapolated
+        (FALLING_FIT, 512, 1.5 - 1.2 * 256 / 768),
+        (FALLING_FIT, 4096, 0.3),
+    ],
+)
+def test_mu_follows_the_line_through_the_anchors_and_holds_at_its_upper_end(
+    fit: ShiftFit, seq_len: int, mu: float
+) -> None:
+    assert fit.mu(seq_len) == pytest.approx(mu, abs=1e-12)
+
+
+def test_the_linear_form_takes_e_to_the_mu() -> None:
+    assert ONE_MEGAPIXEL_FIT.shift(65536) == pytest.approx(3.158192909689768, abs=1e-12)  # e^1.15
+
+
+def test_the_time_shift_keeps_the_ends_and_moves_the_middle_towards_noise() -> None:
+    t = torch.tensor([1.0, 0.5, 0.25, 0.0], dtype=torch.float64)
+    torch.testing.assert_close(time_shift(3.0, t), torch.tensor([1.0, 0.75, 0.5, 0.0], dtype=torch.float64))
+    assert time_shift(3.0, 0.25) == 0.5
+
+
+def test_the_linear_form_of_e_to_the_mu_is_the_schedulers_exponential_shift() -> None:
+    mu = 1.15
+    sigmas = np.linspace(1.0, 1 / 8, 8).tolist()
+    scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True, time_shift_type="exponential")
+    scheduler.set_timesteps(sigmas=sigmas, mu=mu, device="cpu")
+
+    expected = time_shift(math.exp(mu), torch.tensor(sigmas, dtype=torch.float64)).float()
+    torch.testing.assert_close(scheduler.sigmas[:-1], expected)
+
+
+STOCK = {
+    "use_dynamic_shifting": True,
+    "base_image_seq_len": 256,
+    "max_image_seq_len": 8192,
+    "base_shift": 0.5,
+    "max_shift": 0.9,
+    "time_shift_type": "exponential",
+}
+
+
+STOCK_FIT = ShiftFit(base_seq_len=256, max_seq_len=8192, base_shift=0.5, max_shift=0.9)
+
+
+@pytest.mark.parametrize(
+    ("scheduler_folder", "checkpoint_config", "fit"),
+    [
+        # GGUF and single-file transformers ship no scheduler/ folder.
+        (False, None, STOCK_FIT),
+        (True, None, STOCK_FIT),
+        (
+            True,
+            {"base_image_seq_len": 100, "max_image_seq_len": 1000, "base_shift": 0.3, "max_shift": 1.5},
+            ShiftFit(base_seq_len=100, max_seq_len=1000, base_shift=0.3, max_shift=1.5),
+        ),
+        # A key the checkpoint omits takes diffusers' default, not the stock value: 4096 and 1.15, not 8192 and 0.9.
+        (
+            True,
+            {"base_image_seq_len": 100, "base_shift": 0.3},
+            ShiftFit(base_seq_len=100, max_seq_len=4096, base_shift=0.3, max_shift=1.15),
+        ),
+    ],
+    ids=["no-scheduler-folder", "folder-without-config", "checkpoint-config", "checkpoint-config-missing-keys"],
+)
+def test_the_fit_comes_from_the_checkpoints_scheduler_config_else_the_stock_one(
+    tmp_path: Path, scheduler_folder: bool, checkpoint_config: dict | None, fit: ShiftFit
+) -> None:
+    if scheduler_folder:
+        (tmp_path / "scheduler").mkdir()
+    if checkpoint_config is not None:
+        config = {"_class_name": "FlowMatchEulerDiscreteScheduler", "use_dynamic_shifting": True, **checkpoint_config}
+        (tmp_path / "scheduler" / "scheduler_config.json").write_text(json.dumps(config))
+
+    assert ShiftFit.from_scheduler_config(load_scheduler(tmp_path, STOCK).config) == fit

From 1e57dd188654994bd7a1a0dd6a87d9297327af5f Mon Sep 17 00:00:00 2001
From: Josh Corbett <joshwcorbett@icloud.com>
Date: Mon, 19 Oct 2026 00:03:01 -0600
Subject: [PATCH 2/3] refactor(ltx2): evaluate the schedule shift with the
 shared fit

LTX-2 holds its fit at both anchors through ShiftFit. mu between the anchors moves by at most one ulp; float32 sigmas are unchanged.
---
 .../contributing/new-model-integration.mdx    |  2 +-
 invokeai/backend/ltx2/constants.py            | 18 ++++++----
 invokeai/backend/ltx2/sampling.py             | 34 +++----------------
 .../backend/rectified_flow/schedule_shift.py  | 10 ++++--
 tests/backend/ltx2/test_ltx2_sampling.py      | 19 +++++------
 .../rectified_flow/test_schedule_shift.py     |  9 ++++-
 6 files changed, 41 insertions(+), 51 deletions(-)

diff --git a/docs/src/content/docs/contributing/new-model-integration.mdx b/docs/src/content/docs/contributing/new-model-integration.mdx
index 27b0cf2278a..c6325083423 100644
--- a/docs/src/content/docs/contributing/new-model-integration.mdx
+++ b/docs/src/content/docs/contributing/new-model-integration.mdx
@@ -626,7 +626,7 @@ Packages are discovered automatically. A directory of node modules without an `_
 
 - **Backend package (optional).** Put reusable math in `invokeai/backend/new_model/`: noise, packing, position IDs, schedules and text encoding. FLUX, FLUX.2 and Krea-2 have one; Qwen-Image and Z-Image keep the loop in the denoise node. Use a package once code is shared between nodes or needs its own tests.
 - **Schedulers.** Flow-matching architectures use diffusers' `FlowMatchEulerDiscreteScheduler`, or the shared maps in `invokeai/backend/flux/schedulers.py`: `FLUX_SCHEDULER_MAP`, `ZIMAGE_SCHEDULER_MAP` and `ERNIE_IMAGE_SCHEDULER_MAP`, each with name values and labels. Declare which set the UI offers in `FeaturesFacet.scheduler_set`. Set `scheduler_applies_to_graph=True` only if the denoise node really takes a scheduler field. Reproduce the reference sigma schedule and shift (`mu`, time-shift type, terminal shift) exactly, and test it against the reference scheduler.
-- **Resolution-dependent shift.** A `mu` interpolated from the image token count goes through `invokeai/backend/rectified_flow/schedule_shift.py`. Declare the two anchors as a `ShiftFit`, or read them with `ShiftFit.from_scheduler_config` from the scheduler `load_scheduler` returns: a checkpoint's own `scheduler/scheduler_config.json`, else the family's stock config for single-file and GGUF installs. Hand `fit.mu(...)` to `FlowMatchEulerDiscreteScheduler.set_timesteps(mu=...)` (exponential form), or `fit.shift(...)`, which is `e^mu`, to `time_shift` for a hand-built schedule (linear form). The fit is held at its upper anchor. Reference pipelines extrapolate past it, and at large frames that puts nearly every step at high noise, so an image-to-image pass replaces its input instead of refining it. Test the family's anchors and form; the module's own tests cover the line and the hold.
+- **Resolution-dependent shift.** A `mu` interpolated from the image token count goes through `invokeai/backend/rectified_flow/schedule_shift.py`. Declare the two anchors as a `ShiftFit`, or read them with `ShiftFit.from_scheduler_config` from the scheduler `load_scheduler` returns: a checkpoint's own `scheduler/scheduler_config.json`, else the family's stock config for single-file and GGUF installs. Hand `fit.mu(...)` to `FlowMatchEulerDiscreteScheduler.set_timesteps(mu=...)` (exponential form), or `fit.shift(...)`, which is `e^mu`, to `time_shift` for a hand-built schedule (linear form). The fit is held at its upper anchor (`hold_below_base` also holds it below the lower one, as LTX-2's is). Reference pipelines extrapolate past it, and at large frames that puts nearly every step at high noise, so an image-to-image pass replaces its input instead of refining it. Test the family's anchors and form; the module's own tests cover the line and the hold.
 - **External noise (optional).** Only if the denoise node accepts a noise tensor: extend `LatentNoiseType` and its shape and grid logic in `invokeai/app/invocations/latent_noise.py`, and extend `noise_type` in `invocations/noise.py`. Validate with `validate_noise_tensor_shape`.
 - **Timestep-dependent conditioning.** Most architectures turn the text embeddings into the transformer's context once, before the loop. If a module between encoder and transformer reads the timestep -- Anima-3.8B's semantic connector does -- the context has to be recomputed at every step, for the positive, the negative and every regional conditioning. Pass the timestep exactly as the reference does: the connector scales sigma by 1000 before embedding it, so it gets a float32 sigma, where bf16's rounding would move its high frequencies by radians. A mask that leaves a query no key (an empty prompt encoded as one masked token) needs a defined result; `masked_sdpa` makes it zero instead of leaving it to the SDPA backend.
 - **Inpainting.** Rectified-flow models use `RectifiedFlowInpaintExtension` (`invokeai/backend/rectified_flow/rectified_flow_inpaint_extension.py`).
diff --git a/invokeai/backend/ltx2/constants.py b/invokeai/backend/ltx2/constants.py
index 663dd4d62fd..1d9badf88aa 100644
--- a/invokeai/backend/ltx2/constants.py
+++ b/invokeai/backend/ltx2/constants.py
@@ -20,6 +20,8 @@
 
 from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
 
+from invokeai.backend.rectified_flow.schedule_shift import ShiftFit
+
 # --- Video geometry ------------------------------------------------------------------------------
 
 LTX2_LATENT_CHANNELS: Final = 128
@@ -69,13 +71,15 @@
 
 # --- Sampling ------------------------------------------------------------------------------------
 
-# Resolution-aware shift anchors: ``mu`` interpolates between ``base_shift`` at ``base_seq_len``
-# packed video tokens and ``max_shift`` at ``max_seq_len``, and shifts the sigma schedule toward
-# high noise for longer sequences.
-LTX2_BASE_SEQ_LEN: Final = 1024
-LTX2_MAX_SEQ_LEN: Final = 4096
-LTX2_BASE_SHIFT: Final = 0.95
-LTX2_MAX_SHIFT: Final = 2.05
+# The resolution-aware shift over packed video tokens, held at both anchors. The hold is the whole point. The line is
+# fitted between 1024 and 4096 tokens, and every canvas this architecture offers is past its top: a 1248x704x121 clip
+# is 13728 tokens, where the unheld line gives mu 5.5 and the resulting schedule spends 29 of 30 steps above sigma
+# 0.54 before covering 0.54 -> 0.1 in one. Held at 2.05, the schedule is the released pipeline's exactly -- ltx-core's
+# ``LTX2Scheduler`` is called with no latent, so it always takes the max-shift anchor, and the diffusers port's own
+# default operating point is inside the range.
+LTX2_SHIFT_FIT: Final = ShiftFit(
+    base_seq_len=1024, max_seq_len=4096, base_shift=0.95, max_shift=2.05, hold_below_base=True
+)
 
 # The guidance-distilled checkpoint's fixed schedule. The terminal 0 every schedule ends on is
 # appended by ``sampling.build_sigmas``, so this is the 8 sampled levels only.
diff --git a/invokeai/backend/ltx2/sampling.py b/invokeai/backend/ltx2/sampling.py
index 677287846fc..3c93e0ca18f 100644
--- a/invokeai/backend/ltx2/sampling.py
+++ b/invokeai/backend/ltx2/sampling.py
@@ -14,17 +14,10 @@
 add two divisions by a sigma that is already the step's own ratio.
 """
 
-import math
-
 import torch
 
-from invokeai.backend.ltx2.constants import (
-    LTX2_BASE_SEQ_LEN,
-    LTX2_BASE_SHIFT,
-    LTX2_DISTILLED_SIGMAS,
-    LTX2_MAX_SEQ_LEN,
-    LTX2_MAX_SHIFT,
-)
+from invokeai.backend.ltx2.constants import LTX2_DISTILLED_SIGMAS, LTX2_SHIFT_FIT
+from invokeai.backend.rectified_flow.schedule_shift import time_shift
 
 # The last sampled noise level every shifted schedule is stretched onto. Not cosmetic: at the
 # max-shift anchor the raw schedule's final sigma is 0.21, and without the stretch the run would
@@ -38,22 +31,6 @@
 _MIN_REFINE_ENTRY_FRACTION = 0.5
 
 
-def calculate_shift(video_seq_len: int) -> float:
-    """The schedule's ``mu``, interpolated from the token count and held at the anchors outside them.
-
-    The clamp is the whole point. The line is fitted between 1024 and 4096 tokens, and every canvas
-    this architecture offers is past its top: a 1248x704x121 clip is 13728 tokens, where the
-    unclamped line gives mu 5.5 and the resulting schedule spends 29 of 30 steps above sigma 0.54
-    before covering 0.54 -> 0.1 in one. Held at 2.05, the schedule is the released pipeline's
-    exactly -- ltx-core's ``LTX2Scheduler`` is called with no latent, so it always takes the
-    max-shift anchor, and the diffusers port's own default operating point is inside the range.
-    """
-    slope = (LTX2_MAX_SHIFT - LTX2_BASE_SHIFT) / (LTX2_MAX_SEQ_LEN - LTX2_BASE_SEQ_LEN)
-    shift = video_seq_len * slope + LTX2_BASE_SHIFT - slope * LTX2_BASE_SEQ_LEN
-
-    return min(max(shift, LTX2_BASE_SHIFT), LTX2_MAX_SHIFT)
-
-
 def build_sigmas(
     *, distilled: bool, num_steps: int, video_seq_len: int, start_sigma: float | None = None
 ) -> torch.Tensor:
@@ -81,11 +58,10 @@ def build_sigmas(
     if num_steps < 1:
         raise ValueError(f"LTX-2 needs at least one step; got {num_steps}.")
 
-    # Built in float64 and narrowed once at the end. The shift divides mu by itself at sigma 1,
-    # which in float32 is the exact mu over a rounded one -- the schedule would not start at 1.
+    # Built in float64 and narrowed once at the end. The shift divides e^mu by itself at sigma 1,
+    # which in float32 is the exact value over a rounded one -- the schedule would not start at 1.
     sigmas = torch.linspace(1.0, 1.0 / num_steps, num_steps, dtype=torch.float64)
-    mu = math.exp(calculate_shift(video_seq_len))
-    sigmas = mu / (mu + (1.0 / sigmas - 1.0))
+    sigmas = time_shift(LTX2_SHIFT_FIT.shift(video_seq_len), sigmas)
 
     # A single step is already the terminal one, and its `1 - sigma` is 0: there is nothing to
     # stretch onto and the scale factor would be a division by zero.
diff --git a/invokeai/backend/rectified_flow/schedule_shift.py b/invokeai/backend/rectified_flow/schedule_shift.py
index 7f04a7f7ad8..5829b17f6e2 100644
--- a/invokeai/backend/rectified_flow/schedule_shift.py
+++ b/invokeai/backend/rectified_flow/schedule_shift.py
@@ -33,14 +33,15 @@ class ShiftFit:
     """mu along the line through ``(base_seq_len, base_shift)`` and ``(max_seq_len, max_shift)``, held at its upper
     end.
 
-    Below ``base_seq_len`` the line keeps running: mu only sags there, and small frames keep the interpolation they
-    were sampled with upstream.
+    Below ``base_seq_len`` the line keeps running unless ``hold_below_base``: mu only sags there, and small frames keep
+    the interpolation they were sampled with upstream.
     """
 
     base_seq_len: int
     max_seq_len: int
     base_shift: float
     max_shift: float
+    hold_below_base: bool = False
 
     @classmethod
     def from_scheduler_config(cls, config: Mapping[str, Any]) -> "ShiftFit":
@@ -56,7 +57,10 @@ def mu(self, seq_len: int) -> float:
         # diffusers' `calculate_shift`, in its order of operations, on the held sequence length.
         m = (self.max_shift - self.base_shift) / (self.max_seq_len - self.base_seq_len)
         b = self.base_shift - m * self.base_seq_len
-        return min(seq_len, self.max_seq_len) * m + b
+        seq_len = min(seq_len, self.max_seq_len)
+        if self.hold_below_base:
+            seq_len = max(seq_len, self.base_seq_len)
+        return seq_len * m + b
 
     def shift(self, seq_len: int) -> float:
         """``e^mu``, the shift the linear form takes."""
diff --git a/tests/backend/ltx2/test_ltx2_sampling.py b/tests/backend/ltx2/test_ltx2_sampling.py
index c10ba1ac0d8..7de9ce44b58 100644
--- a/tests/backend/ltx2/test_ltx2_sampling.py
+++ b/tests/backend/ltx2/test_ltx2_sampling.py
@@ -8,7 +8,7 @@
 
 from invokeai.backend.ltx2.constants import (
     LTX2_DISTILLED_SIGMAS,
-    LTX2_MAX_SEQ_LEN,
+    LTX2_SHIFT_FIT,
     LTX2_STAGE_2_NOISE_SCALE,
 )
 from invokeai.backend.ltx2.packing import video_sequence_length
@@ -16,7 +16,6 @@
     LTX2_SHIFT_TERMINAL,
     build_refine_sigmas,
     build_sigmas,
-    calculate_shift,
     flow_step,
 )
 
@@ -96,12 +95,12 @@ def test_the_shift_interpolates_between_the_anchors_and_holds_outside_them() ->
     """Independently: mu is 0.95 at 1024 tokens and 2.05 at 4096, linear in between, flat outside.
 
     Extrapolating instead is what makes the schedule degenerate at real canvases -- see below."""
-    assert calculate_shift(1024) == pytest.approx(0.95)
-    assert calculate_shift(4096) == pytest.approx(2.05)
-    assert calculate_shift(2560) == pytest.approx(1.5)
-    assert calculate_shift(320) == pytest.approx(0.95)
-    assert calculate_shift(7168) == pytest.approx(2.05)
-    assert calculate_shift(13728) == pytest.approx(2.05)
+    assert LTX2_SHIFT_FIT.mu(1024) == pytest.approx(0.95)
+    assert LTX2_SHIFT_FIT.mu(4096) == pytest.approx(2.05)
+    assert LTX2_SHIFT_FIT.mu(2560) == pytest.approx(1.5)
+    assert LTX2_SHIFT_FIT.mu(320) == pytest.approx(0.95)
+    assert LTX2_SHIFT_FIT.mu(7168) == pytest.approx(2.05)
+    assert LTX2_SHIFT_FIT.mu(13728) == pytest.approx(2.05)
 
 
 def test_the_schedule_keeps_spending_steps_where_the_clip_is_formed() -> None:
@@ -120,7 +119,7 @@ def test_the_schedule_keeps_spending_steps_where_the_clip_is_formed() -> None:
 def test_every_canvas_the_panel_offers_samples_the_released_schedule() -> None:
     """All three presets are past the 4096-token anchor, so all three get the released pipeline's
     own schedule rather than three different extrapolations of it."""
-    reference = build_sigmas(distilled=False, num_steps=30, video_seq_len=LTX2_MAX_SEQ_LEN)
+    reference = build_sigmas(distilled=False, num_steps=30, video_seq_len=LTX2_SHIFT_FIT.max_seq_len)
 
     for width, height in ((928, 512), (1248, 704), (1376, 768)):
         sigmas = build_sigmas(distilled=False, num_steps=30, video_seq_len=video_sequence_length(121, height, width))
@@ -183,7 +182,7 @@ def test_a_zero_eta_ancestral_step_is_the_deterministic_one() -> None:
 def test_the_shifted_schedule_matches_the_reference_transformation() -> None:
     """Independently: shift sigma -> e^mu / (e^mu + 1/sigma - 1), then stretch onto the terminal."""
     num_steps, seq_len = 10, 2048
-    mu = math.exp(calculate_shift(seq_len))
+    mu = math.exp(0.95 + 1.1 * (seq_len - 1024) / 3072)
     raw = [mu / (mu + 1 / s - 1) for s in torch.linspace(1.0, 0.1, num_steps).tolist()]
     scale = (1 - raw[-1]) / (1 - LTX2_SHIFT_TERMINAL)
     expected = [1 - (1 - s) / scale for s in raw]
diff --git a/tests/backend/rectified_flow/test_schedule_shift.py b/tests/backend/rectified_flow/test_schedule_shift.py
index c258e435b72..bb4436ba118 100644
--- a/tests/backend/rectified_flow/test_schedule_shift.py
+++ b/tests/backend/rectified_flow/test_schedule_shift.py
@@ -14,6 +14,9 @@
 ONE_MEGAPIXEL_FIT = ShiftFit(base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.15)
 # A checkpoint's config may make the line fall. Holding the sequence length, not mu, keeps the hold at its end.
 FALLING_FIT = ShiftFit(base_seq_len=256, max_seq_len=1024, base_shift=1.5, max_shift=0.3)
+HELD_BOTH_ENDS_FIT = ShiftFit(
+    base_seq_len=1024, max_seq_len=4096, base_shift=0.95, max_shift=2.05, hold_below_base=True
+)
 
 
 @pytest.mark.parametrize(
@@ -28,9 +31,13 @@
         (ONE_MEGAPIXEL_FIT, 65536, 1.15),  # 11.55 extrapolated
         (FALLING_FIT, 512, 1.5 - 1.2 * 256 / 768),
         (FALLING_FIT, 4096, 0.3),
+        (HELD_BOTH_ENDS_FIT, 1, 0.95),  # held below the lower anchor too
+        (HELD_BOTH_ENDS_FIT, 1024, 0.95),
+        (HELD_BOTH_ENDS_FIT, 2560, 1.5),
+        (HELD_BOTH_ENDS_FIT, 13728, 2.05),
     ],
 )
-def test_mu_follows_the_line_through_the_anchors_and_holds_at_its_upper_end(
+def test_mu_follows_the_line_through_the_anchors_and_holds_where_it_ends(
     fit: ShiftFit, seq_len: int, mu: float
 ) -> None:
     assert fit.mu(seq_len) == pytest.approx(mu, abs=1e-12)

From 9f6fea56b539288915c61ddde7521630f48bea70 Mon Sep 17 00:00:00 2001
From: Josh Corbett <joshwcorbett@icloud.com>
Date: Mon, 19 Oct 2026 13:25:01 -0600
Subject: [PATCH 3/3] test(qwen-image): drive an explicit shift through the
 node to the scheduler

Restores the run-to-scheduler helper so a node-level test checks that the shift input reaches set_timesteps as ln(shift).
---
 .../invocations/test_qwen_image_denoise.py    | 69 +++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/tests/app/invocations/test_qwen_image_denoise.py b/tests/app/invocations/test_qwen_image_denoise.py
index 62a524cb3e5..158e745015b 100644
--- a/tests/app/invocations/test_qwen_image_denoise.py
+++ b/tests/app/invocations/test_qwen_image_denoise.py
@@ -2,10 +2,18 @@
 
 import json
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
+import torch
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 
+from invokeai.app.invocations.qwen_image import qwen_image_denoise
 from invokeai.app.invocations.qwen_image.qwen_image_denoise import QwenImageDenoiseInvocation
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
+    ConditioningFieldData,
+    QwenImageConditioningInfo,
+)
 
 
 class TestPrepareCfgScale:
@@ -258,3 +266,64 @@ def test_an_explicit_shift_sets_mu_at_any_size_with_the_checkpoints_scheduler(tm
 
     assert mu == pytest.approx(1.0986122886681098, abs=1e-12)  # ln 3
     assert scheduler.config.max_shift == 1.5
+
+
+class _StopAtTimesteps(Exception):
+    """Ends a run once the scheduler has been handed its mu."""
+
+
+class _TinyQwenImageTransformer:
+    """Stands in for `QwenImageTransformer2DModel`: the config the node reads before it sets the schedule."""
+
+    config = SimpleNamespace(patch_size=2, out_channels=16)
+
+
+def _run_to_scheduler(monkeypatch, model_path: Path, width: int, height: int, shift: float | None = None) -> float:
+    """Run `_run_diffusion` with the real scheduler up to `set_timesteps` and return the mu it was handed."""
+    monkeypatch.setattr(qwen_image_denoise, "QwenImageTransformer2DModel", _TinyQwenImageTransformer)
+    monkeypatch.setattr(qwen_image_denoise.TorchDevice, "choose_torch_device", lambda: torch.device("cpu"))
+    captured: list[float] = []
+
+    def capture_mu(self, *, sigmas, mu, device) -> None:
+        captured.append(mu)
+        # A frame of real size need not be denoised to know the mu it was handed.
+        raise _StopAtTimesteps
+
+    monkeypatch.setattr(FlowMatchEulerDiscreteScheduler, "set_timesteps", capture_mu)
+    conditioning = ConditioningFieldData(conditionings=[QwenImageConditioningInfo(prompt_embeds=torch.zeros(1, 4, 8))])
+    context = SimpleNamespace(
+        models=SimpleNamespace(
+            load=lambda _identifier: SimpleNamespace(model=_TinyQwenImageTransformer()),
+            get_config=lambda _identifier: SimpleNamespace(),
+            get_absolute_path=lambda _config: model_path,
+        ),
+        conditioning=SimpleNamespace(load=lambda _name: conditioning),
+    )
+    node = QwenImageDenoiseInvocation.model_construct(
+        latents=None,
+        reference_latents=None,
+        denoise_mask=None,
+        denoising_start=0.0,
+        denoising_end=1.0,
+        transformer=SimpleNamespace(transformer=SimpleNamespace(), loras=[]),
+        positive_conditioning=SimpleNamespace(conditioning_name="positive"),
+        negative_conditioning=None,
+        cfg_scale=1.0,
+        width=width,
+        height=height,
+        steps=40,
+        seed=0,
+        shift=shift,
+    )
+
+    with pytest.raises(_StopAtTimesteps):
+        node._run_diffusion(context)
+
+    assert len(captured) == 1
+    return captured[0]
+
+
+def test_an_explicit_shift_reaches_the_scheduler(monkeypatch, tmp_path) -> None:
+    mu = _run_to_scheduler(monkeypatch, tmp_path / "qwen-image.gguf", 4096, 4096, shift=3.0)
+
+    assert mu == pytest.approx(1.0986122886681098, abs=1e-12)  # ln 3