fix(qwen-image-2-1): hold the base model's shift at the end of its fit

docs/src/content/docs/Users Guide/03.Models/Local Models/qwen-image-2-1.mdx

@@ -1,7 +1,7 @@
 ---
 title: Qwen-Image-2.1
 description: Generate images with Qwen-Image-2.1 and its 8-step Turbo variant, from the Diffusers folders or a GGUF transformer, including transparent (RGBA) output.
-lastUpdated: 2026-10-16
+lastUpdated: 2026-10-17
 sidebar:
   order: 8
 ---
@@ -67,5 +67,9 @@ Image-to-image, inpainting and outpainting read the colors of the input image an
 
 - **Steps** — 40 for the base model and 8 for Turbo. Turbo's schedule is tuned for 8 steps; other step counts are spread over the same curve and work, but 8 is what it was distilled for.
 - **CFG** — 1 by default, which runs one prompt per step. Above 1, the negative prompt is used and each step costs about twice as long.
-- There is no scheduler setting: both variants use the schedule they were trained with.
+- There is no scheduler setting: both variants use their own schedule.
 - **Denoising strength** (image-to-image, inpainting) — on the base model a strength changes the image about as much as on Qwen-Image or Z-Image. Turbo's 8-step schedule spends most of its steps at high noise, so there a strength sets the starting noise level directly (0.5 starts at noise level 0.5): the same value changes the image less on Turbo than on the base model.
+
+:::note[Image size and the noise schedule]
+The base model puts more of its steps at high noise the larger the image, up to the area of about 1448×1448 (2.1 megapixels). Larger images keep that schedule. Without the limit, an image-to-image pass at strength 0.5 would start at 77% noise on a 2048×2048 image and at 95%, almost pure noise, on a 4096×4096 one, so it would replace the image instead of refining it; with it, both start at 70%. Text-to-image above that size uses the held schedule too, so it differs from diffusers' `QwenImage21Pipeline` there, and base-model images made at those sizes with an earlier Invoke version don't reproduce exactly with the same seed. Turbo's fixed schedule does not depend on the image size.
+:::

invokeai/app/invocations/qwen_image_2_1/qwen_image_2_1_denoise.py

@@ -80,7 +80,7 @@ def unpack_latents(latents: torch.Tensor, height: int, width: int) -> torch.Tens
     title="Denoise - Qwen-Image-2.1",
     tags=["image", "qwen_image_2_1", "qwen-image-2.1"],
     category="image",
-    version="1.0.0",
+    version="1.0.1",
     classification=Classification.Prototype,
 )
 class QwenImage21DenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):

invokeai/backend/qwen_image_2_1/sampling.py

@@ -31,10 +31,16 @@
 
 
 def calculate_mu(image_seq_len: int) -> float:
-    """The pipeline's `calculate_shift` over target-image tokens only (condition images are not counted)."""
+    """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 image_seq_len * m + (c["base_shift"] - m * 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"])
 
 
 def turbo_sigmas(steps: int) -> list[float]:
@@ -50,7 +56,8 @@ 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, exactly as `QwenImage21Pipeline` samples it.
+    """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`).
 
     `image_seq_len` is the target image's latent token count, `(height // 16) * (width // 16)`; only the base
     model's shift reads it.

invokeai/frontend/api/openapi.json

@@ -95557,7 +95557,7 @@
         "tags": ["image", "qwen_image_2_1", "qwen-image-2.1"],
         "title": "Denoise - Qwen-Image-2.1",
         "type": "object",
-        "version": "1.0.0",
+        "version": "1.0.1",
         "output": {
           "$ref": "#/components/schemas/LatentsOutput"
         }

tests/backend/qwen_image_2_1/test_sampling.py

@@ -1,14 +1,16 @@
 """Qwen-Image-2.1's schedules against what `QwenImage21Pipeline` samples.
 
 The expected sigmas were recorded from the pipeline's own path -- `retrieve_timesteps` with the scheduler configs
-shipped in Qwen/Qwen-Image-2.1 and Qwen/Qwen-Image-2.1-Turbo -- not derived from the code under test.
+shipped in Qwen/Qwen-Image-2.1 and Qwen/Qwen-Image-2.1-Turbo -- not derived from the code under test. Above the end
+of the shift's fit (8192 tokens) the pipeline keeps extrapolating and Invoke holds the shift; those expectations come
+from diffusers' `FlowMatchEulerDiscreteScheduler` with the base config at the stated mu.
 """
 
 import pytest
 import torch
 
 from invokeai.backend.model_manager.taxonomy import QwenImage21VariantType
-from invokeai.backend.qwen_image_2_1.sampling import TURBO_SIGMAS, build_sigmas
+from invokeai.backend.qwen_image_2_1.sampling import TURBO_SIGMAS, build_sigmas, calculate_mu
 
 BASE = QwenImage21VariantType.Base
 TURBO = QwenImage21VariantType.Turbo
@@ -17,19 +19,48 @@
 @pytest.mark.parametrize(
     ("steps", "tokens", "head", "tail"),
     [
-        # 1024x1024 and 2048x2048: the shift grows with the image, which keeps more steps at high noise.
+        # 1024x1024, within the shift's fit.
         (40, 4096, [1.0, 0.986964, 0.973593], [0.067881, 0.02, 0.0]),
-        (40, 16384, [1.0, 0.992646, 0.985013], [0.10223, 0.02, 0.0]),
         (8, 4096, [1.0, 0.916024, 0.820046], [0.244054, 0.02, 0.0]),
     ],
 )
-def test_the_base_schedule_is_the_pipelines(steps: int, tokens: int, head: list[float], tail: list[float]) -> None:
+def test_within_its_fit_the_base_schedule_is_the_pipelines(
+    steps: int, tokens: int, head: list[float], tail: list[float]
+) -> None:
     sigmas = build_sigmas(BASE, steps, tokens)
     assert sigmas.shape == (steps + 1,)
     torch.testing.assert_close(sigmas[:3], torch.tensor(head), atol=1e-6, rtol=0)
     torch.testing.assert_close(sigmas[-3:], torch.tensor(tail), atol=1e-6, rtol=0)
 
 
+@pytest.mark.parametrize(
+    ("steps", "tokens", "head", "tail"),
+    [
+        # mu 0.9 at 2048x1024 (the fit's end), 2048x2048 and 4096x4096. Extrapolated, 4096x4096 would reach mu
+        # 3.79 and an image-to-image pass at strength 0.5 would start at sigma 0.95 instead of 0.70.
+        (40, 8192, [1.0, 0.989251, 0.978174], [0.077602, 0.02, 0.0]),
+        (40, 16384, [1.0, 0.989251, 0.978174], [0.077602, 0.02, 0.0]),
+        (40, 65536, [1.0, 0.989251, 0.978174], [0.077602, 0.02, 0.0]),
+        (8, 65536, [1.0, 0.927302, 0.841941], [0.272285, 0.02, 0.0]),
+    ],
+)
+def test_above_its_fit_the_base_schedule_holds_the_shift_at_the_fits_end(
+    steps: int, tokens: int, head: list[float], tail: list[float]
+) -> None:
+    sigmas = build_sigmas(BASE, steps, tokens)
+    torch.testing.assert_close(sigmas[:3], torch.tensor(head), atol=1e-6, rtol=0)
+    torch.testing.assert_close(sigmas[-3:], torch.tensor(tail), atol=1e-6, rtol=0)
+    torch.testing.assert_close(sigmas, build_sigmas(BASE, steps, 8192), atol=0, rtol=0)
+
+
+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)
+    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)
+
+
 @pytest.mark.parametrize("tokens", [4096, 16384])
 def test_turbo_samples_its_table_unshifted_at_every_size(tokens: int) -> None:
     sigmas = build_sigmas(TURBO, 8, tokens)