fix(qwen-image): hold the dynamic shift at the end of its fit

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

@@ -1,7 +1,7 @@
 ---
 title: Qwen Image
 description: Generate and edit images with the Qwen Image 2512 text-to-image model and the Qwen Image Edit 2511 reference-image editing model, including GGUF installs and the Lightning LoRAs.
-lastUpdated: 2026-09-30
+lastUpdated: 2026-10-16
 sidebar:
   order: 8
 ---
@@ -91,6 +91,16 @@ Selecting a Qwen Image model applies these defaults:
 There is no scheduler choice — Qwen Image always uses its own flow-matching schedule, so the
 **Scheduler** control is hidden.
 
+:::note[Image size and the noise schedule]
+Qwen Image 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. To sample with a different schedule, set **Shift** on the
+**Denoise - Qwen Image** node in the workflow editor (Generate and Canvas don't send it). On this node
+Shift is the exponential shift e^mu, not mu as on Krea-2: about 3.72 reproduces the unlimited schedule at
+2048×2048, and larger images need larger values, about 44 at 4096×4096.
+:::
+
 The **negative prompt** is used only when CFG Scale is above 1; at CFG 1 it is ignored.
 
 ## Lightning LoRAs (4 and 8 steps)

invokeai/app/invocations/qwen_image/qwen_image_denoise.py

@@ -42,7 +42,7 @@
     title="Denoise - Qwen Image",
     tags=["image", "qwen_image"],
     category="image",
-    version="1.0.0",
+    version="1.0.1",
     classification=Classification.Prototype,
 )
 class QwenImageDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
@@ -82,7 +82,8 @@ class QwenImageDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         default=None,
         description="Override the sigma schedule shift. "
         "When set, uses a fixed shift (e.g. 3.0 for Lightning LoRAs) instead of the default dynamic shifting. "
-        "Leave unset for the base model's default schedule.",
+        "Leave unset for the base model's default schedule, which is resolution-dependent up to the area of the "
+        "scheduler fit's end (about 1448x1448, 2.1 megapixels, for the stock config) and held above it.",
     )
 
     @torch.no_grad()
@@ -319,15 +320,18 @@ def _run_diffusion(self, context: InvocationContext):
             # Lightning LoRA: fixed shift
             mu = math.log(self.shift)
         else:
-            # Default dynamic shifting
-            # Linear interpolation matching diffusers' calculate_shift
-            base_shift = scheduler.config.get("base_shift", 0.5)
-            max_shift = scheduler.config.get("max_shift", 0.9)
-            base_seq = scheduler.config.get("base_image_seq_len", 256)
-            max_seq = scheduler.config.get("max_image_seq_len", 4096)
-            m = (max_shift - base_shift) / (max_seq - base_seq)
-            b = base_shift - m * base_seq
-            mu = image_seq_len * m + b
+            # 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
 
         init_sigmas = np.linspace(1.0, 1.0 / self.steps, self.steps).tolist()
         scheduler.set_timesteps(sigmas=init_sigmas, mu=mu, device=device)

invokeai/frontend/api/openapi.json

@@ -96336,7 +96336,7 @@
               }
             ],
             "default": null,
-            "description": "Override the sigma schedule shift. When set, uses a fixed shift (e.g. 3.0 for Lightning LoRAs) instead of the default dynamic shifting. Leave unset for the base model's default schedule.",
+            "description": "Override the sigma schedule shift. When set, uses a fixed shift (e.g. 3.0 for Lightning LoRAs) instead of the default dynamic shifting. Leave unset for the base model's default schedule, which is resolution-dependent up to the area of the scheduler fit's end (about 1448x1448, 2.1 megapixels, for the stock config) and held above it.",
             "field_kind": "input",
             "input": "any",
             "orig_default": null,
@@ -96355,7 +96355,7 @@
         "tags": ["image", "qwen_image"],
         "title": "Denoise - Qwen Image",
         "type": "object",
-        "version": "1.0.0",
+        "version": "1.0.1",
         "output": {
           "$ref": "#/components/schemas/LatentsOutput"
         }

invokeai/frontend/api/schema.ts

@@ -41628,7 +41628,7 @@ export type components = {
             seed?: number;
             /**
              * Shift
-             * @description Override the sigma schedule shift. When set, uses a fixed shift (e.g. 3.0 for Lightning LoRAs) instead of the default dynamic shifting. Leave unset for the base model's default schedule.
+             * @description Override the sigma schedule shift. When set, uses a fixed shift (e.g. 3.0 for Lightning LoRAs) instead of the default dynamic shifting. Leave unset for the base model's default schedule, which is resolution-dependent up to the area of the scheduler fit's end (about 1448x1448, 2.1 megapixels, for the stock config) and held above it.
              * @default null
              */
             shift?: number | null;

tests/app/invocations/test_qwen_image_denoise.py

@@ -1,8 +1,19 @@
 """Tests for the Qwen Image denoise invocation."""
 
+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:
@@ -195,3 +206,108 @@ def test_edit_matches_diffusers_layout(self):
         assert len(result[0]) == 2
         assert result[0][0] == (1, 40, 56)
         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.
+        (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)
+
+
+@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(
+        json.dumps(
+            {
+                "_class_name": "FlowMatchEulerDiscreteScheduler",
+                "use_dynamic_shifting": True,
+                "base_image_seq_len": 256,
+                "max_image_seq_len": 1024,
+                "base_shift": 0.3,
+                "max_shift": 1.5,
+                "shift_terminal": 0.02,
+                "time_shift_type": "exponential",
+            }
+        )
+    )
+
+    assert _run_to_scheduler(monkeypatch, tmp_path, size, size) == pytest.approx(mu)
+
+
+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)
+
+    assert mu == pytest.approx(1.0986122886681098)  # ln 3