refactor(sampling): share one schedule-shift fit across flow-matching families

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 (`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`).

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)

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)

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)