fix(flux2,z-image): run Diffusers schedulers on the node's sigma schedule

docs/src/content/docs/Users Guide/04.Image Generation/upscale.mdx

@@ -19,7 +19,7 @@ The panel is split into sections, top to bottom:
   - **Scale**: the output size is the source size multiplied by this value. The panel shows the resulting input and output size and warns when the output is large enough to need much more time and VRAM.
   - **Treatment preset**, or **Creativity** and **Structure** directly. Higher Creativity invents more detail and gives the prompt more influence. Higher Structure keeps the source composition more strictly; it is only shown when a Tile ControlNet is in use.
 - **Detail guidance**: the prompt. It is the project prompt shared with Generate. For upscaling, describe the medium, texture and style rather than new content.
-- **Generation**: the **Main model** for the detail pass (**SD1.5**, **SDXL**, **FLUX.1**, **FLUX.2**, **Z-Image**, **Anima**, **Krea-2** or **Qwen Image**), steps, CFG scale, scheduler, seed, and LoRAs that match the main model. FLUX.1, FLUX.2, Z-Image, Krea-2 and Qwen Image always run **Euler**, so they show no scheduler.
+- **Generation**: the **Main model** for the detail pass (**SD1.5**, **SDXL**, **FLUX.1**, **FLUX.2**, **Z-Image**, **Anima**, **Krea-2** or **Qwen Image**), steps, CFG scale, scheduler, seed, and LoRAs that match the main model. The scheduler list only holds schedulers that work for upscaling (see below).
 - **Advanced**: the components the main model needs (see below), tile size and tile overlap, VAE and VAE precision, and CLIP skip for SD1.5.
 
 ## What each model family needs
@@ -44,7 +44,7 @@ Z-Image, Anima, Krea-2 and Qwen Image count the start of the detail pass in step
 
 FLUX.2 starts where FLUX.1 does and then spreads all of your Steps over the rest of the pass, so 30 steps run 30 times. Picking a FLUX.2 model therefore sets its own step count: 4 for Klein 4B and 9B, 28 for Klein Base and [dev]. Switching to another family afterwards sets Steps back to 30.
 
-The panel offers only schedulers that start where Creativity puts the pass: **Euler** on FLUX.1, FLUX.2, Z-Image, Krea-2 and Qwen Image, and **Euler**, **DPM++ 2M**, **DPM++ 2M SDE** and **ER-SDE** on Anima. Heun and LCM build their own noise schedule, so they would treat the enlarged image as pure noise or leave noise in it.
+The panel offers only schedulers that start where Creativity puts the pass: **Euler** and **Heun** on FLUX.1 and Z-Image, and **Euler**, **DPM++ 2M**, **DPM++ 2M SDE** and **ER-SDE** on Anima. LCM builds its own noise schedule, so it would treat the enlarged image as pure noise. Anima's Heun starts below that point and leaves noise in the image. FLUX.2, Krea-2 and Qwen Image always run **Euler**, so they show no scheduler.
 
 The Z-Image Tile ControlNet is optional because the detail pass runs over the whole image at once: there are no tile seams for it to hold together. Adding it lets **Structure** pull the result back toward the enlarged image, at the cost of a second network's memory and time.
 

invokeai/app/invocations/flux2/flux2_denoise.py

@@ -31,7 +31,6 @@
 from invokeai.backend.flux2.extensions.regional_prompting_extension import Flux2RegionalPromptingExtension
 from invokeai.backend.flux2.ref_image_extension import Flux2RefImageExtension
 from invokeai.backend.flux2.sampling_utils import (
-    FLUX2_TXT2IMG_SCHEDULER_KWARGS,
     compute_empirical_mu,
     generate_img_ids_flux2,
     get_noise_flux2,
@@ -379,36 +378,24 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
         transformer_config = context.models.get_config(self.transformer.transformer)
         assert transformer_config.base == BaseModelType.Flux2 and transformer_config.type == ModelType.Main
 
-        # Calculate the timestep schedule using FLUX.2 specific schedule
-        # This matches diffusers' Flux2Pipeline implementation
-        # Note: Schedule shifting is handled by the scheduler via mu parameter
+        # The FLUX.2 schedule: diffusers' Flux2KleinPipeline linear sigmas with the exponential shift
+        # from mu, which its scheduler applies internally.
         image_seq_len = packed_h * packed_w
-        timesteps = get_schedule_flux2(
-            num_steps=self.num_steps,
-            image_seq_len=image_seq_len,
-        )
-        # Compute mu for dynamic schedule shifting (used by FlowMatchEulerDiscreteScheduler)
         mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=self.num_steps)
+        timesteps = time_shift_flux2(get_schedule_flux2(num_steps=self.num_steps, image_seq_len=image_seq_len), mu)
+
+        # Every path runs this shifted schedule, and it is shifted before clipping so the clip selects
+        # from the schedule the model is actually run on. That keeps denoising_start and denoising_end
+        # on one scale: a pass that ends at denoising_end = x stops at the sigma a later pass with
+        # denoising_start = x begins at. The clip works by sigma value, not by step count, so shifting
+        # first only changes how many steps fall below a given sigma.
+        timesteps = clip_timestep_schedule_fractional(timesteps, self.denoising_start, self.denoising_end)
 
-        # img2img and inpainting step this schedule manually (see the scheduler setup below), while
-        # txt2img hands it to the scheduler, which applies the exponential shift from mu itself. Apply
-        # the same shift here so both paths follow one schedule, and apply it before clipping so the
-        # clip selects from the schedule the model is actually run on. The clip works by sigma value,
-        # not by step count, so shifting first does not change what a given denoising_start means --
-        # it only changes how many steps fall below that sigma.
-        #
-        # Keyed on the same condition that decides whether a scheduler gets built below. Manual Euler
-        # also runs when the requested scheduler is unavailable (e.g. "lcm" when FlowMatchLCMScheduler
-        # cannot be imported); keying on denoise_mask/denoising_start alone would leave that run on the
-        # unshifted linear schedule -- the one this fix exists to avoid.
+        # Manual Euler steps img2img and inpainting, and also runs when the requested scheduler is
+        # unavailable (e.g. "lcm" when FlowMatchLCMScheduler cannot be imported).
         uses_manual_euler = (
             self.denoise_mask is not None or self.denoising_start > 1e-5 or self.scheduler not in FLUX_SCHEDULER_MAP
         )
-        if uses_manual_euler:
-            timesteps = time_shift_flux2(timesteps, mu)
-
-        # Clip the timesteps schedule based on denoising_start and denoising_end
-        timesteps = clip_timestep_schedule_fractional(timesteps, self.denoising_start, self.denoising_end)
 
         # Clipping keeps only the sigmas that fall inside the denoising window, so a narrow window
         # runs fewer steps than were requested. Under the shift that is severe: at 4 steps and
@@ -424,13 +411,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
             if self.add_noise:
                 assert noise is not None
                 # Noise the init latents using the first timestep from the clipped
-                # InvokeAI schedule.
-                #
-                # Known limitation: if a scheduler later uses a different first
-                # effective timestep/sigma than this precomputed schedule, the
-                # img2img preblend below may not match that scheduler exactly.
-                # This is an existing pipeline limitation and applies to both
-                # seed-generated noise and externally supplied noise.
+                # InvokeAI schedule, which manual Euler steps from.
                 t_0 = timesteps[0]
                 x = t_0 * noise + (1.0 - t_0) * init_latents
             else:
@@ -495,27 +476,12 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
         num_steps = len(timesteps) - 1
         cfg_scale_list = [self.cfg_scale] * num_steps
 
-        # Create scheduler with FLUX.2 Klein configuration
-        # For inpainting/img2img, use manual Euler stepping to preserve the exact
-        # clipped timestep schedule used for the initial latent/noise preblend.
-        # For txt2img, use the scheduler with dynamic shifting for optimal results.
-        #
-        # This split is intentional. Reusing a scheduler for img2img here can
-        # change the first effective timestep/sigma and break parity with the
-        # preblend computed above.
+        # txt2img uses the selected Diffusers scheduler; img2img and inpainting keep manual Euler on the
+        # respaced schedule above. The schedule is already shifted, so the scheduler must not shift it
+        # again: no dynamic shifting, and the default static shift of 1.0.
         scheduler = None
         if not uses_manual_euler:
-            # Only use scheduler for txt2img - use manual Euler for inpainting to preserve exact timesteps
-            scheduler_class = FLUX_SCHEDULER_MAP[self.scheduler]
-            # FlowMatchHeunDiscreteScheduler only supports num_train_timesteps and shift parameters
-            # FlowMatchEulerDiscreteScheduler and FlowMatchLCMScheduler support dynamic shifting
-            if self.scheduler == "heun":
-                scheduler = scheduler_class(
-                    num_train_timesteps=1000,
-                    shift=3.0,
-                )
-            else:
-                scheduler = scheduler_class(**FLUX2_TXT2IMG_SCHEDULER_KWARGS)
+            scheduler = FLUX_SCHEDULER_MAP[self.scheduler](num_train_timesteps=1000)
 
         # Prepare reference image extension for FLUX.2 Klein built-in editing
         ref_image_extension = None
@@ -638,7 +604,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                     neg_txt=neg_txt,
                     neg_txt_ids=neg_txt_ids,
                     scheduler=scheduler,
-                    mu=mu,
                     inpaint_extension=inpaint_extension,
                     img_cond_seq=img_cond_seq,
                     img_cond_seq_ids=img_cond_seq_ids,

invokeai/app/invocations/z_image/z_image_denoise.py

@@ -1,11 +1,11 @@
-import inspect
 import math
 from contextlib import ExitStack
 from typing import Callable, Iterator, Optional
 
 import einops
 import torch
 import torchvision.transforms as tv_transforms
+from diffusers import FlowMatchHeunDiscreteScheduler
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from PIL import Image
 from torchvision.transforms.functional import resize as tv_resize
@@ -27,7 +27,12 @@
 from invokeai.app.invocations.vae.z_image_image_to_latents import ZImageImageToLatentsInvocation
 from invokeai.app.invocations.z_image.z_image_control import ZImageControlField
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.flux.schedulers import ZIMAGE_SCHEDULER_LABELS, ZIMAGE_SCHEDULER_MAP, ZIMAGE_SCHEDULER_NAME_VALUES
+from invokeai.backend.flux.schedulers import (
+    ZIMAGE_SCHEDULER_LABELS,
+    ZIMAGE_SCHEDULER_MAP,
+    ZIMAGE_SCHEDULER_NAME_VALUES,
+    set_heun_timesteps_from_sigmas,
+)
 from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat
 from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec
 from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_TRANSFORMER_PREFIX
@@ -467,13 +472,8 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
             if self.add_noise:
                 assert noise is not None
                 # Noise the init latents using the first sigma from the clipped
-                # InvokeAI schedule.
-                #
-                # Known limitation: if the selected scheduler later starts from a
-                # different first effective sigma/timestep than sigmas[0], the
-                # img2img preblend below may not match that scheduler exactly.
-                # This is an existing pipeline limitation and affects both
-                # internally generated noise and externally supplied noise.
+                # InvokeAI schedule. Euler and Heun start from this sigma; LCM
+                # creates its own schedule, whose first sigma can differ.
                 s_0 = sigmas[0]
                 latents = s_0 * noise + (1.0 - s_0) * init_latents
             else:
@@ -506,23 +506,22 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
         # Initialize the diffusers scheduler if not using built-in Euler
         scheduler: SchedulerMixin | None = None
         use_scheduler = self.scheduler != "euler"
+        is_heun = False
 
         if use_scheduler:
             scheduler_class = ZIMAGE_SCHEDULER_MAP[self.scheduler]
             scheduler = scheduler_class(
                 num_train_timesteps=1000,
                 shift=1.0,
             )
-            # Set timesteps - LCM uses its own sigma schedule (num_inference_steps),
-            # while other schedulers can use custom sigmas if supported
-            is_lcm = self.scheduler == "lcm"
-            set_timesteps_sig = inspect.signature(scheduler.set_timesteps)
-            if not is_lcm and "sigmas" in set_timesteps_sig.parameters:
-                scheduler.set_timesteps(sigmas=sigmas, device=device)
+            is_heun = isinstance(scheduler, FlowMatchHeunDiscreteScheduler)
+            if is_heun:
+                # Heun runs the shifted, clipped schedule above, so it starts at the sigma the img2img
+                # preblend used and stops at the requested end.
+                set_heun_timesteps_from_sigmas(scheduler, sigmas[:-1], sigmas[-1], device)
             else:
-                # LCM or a scheduler without custom-sigma support computes its own
-                # schedule from num_inference_steps. That can diverge from sigmas[0]
-                # used in the img2img preblend above.
+                # LCM owns its schedule: clipping only sets its step count. Its first sigma can
+                # differ from sigmas[0] used in the img2img preblend above.
                 scheduler.set_timesteps(num_inference_steps=total_steps, device=device)
 
             # For Heun scheduler, the number of actual steps may differ
@@ -724,7 +723,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                     sigma_curr = sched_timestep.item() / scheduler.config.num_train_timesteps
 
                     # For Heun scheduler, track if we're in first or second order step
-                    is_heun = hasattr(scheduler, "state_in_first_order")
                     in_first_order = scheduler.state_in_first_order if is_heun else True
 
                     # Timestep tensor for Z-Image model
@@ -805,17 +803,19 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                     if inpaint_extension is not None:
                         latents = inpaint_extension.merge_intermediate_latents_with_init_latents(latents, sigma_prev)
 
-                    # For Heun, only increment user step after second-order step completes
+                    # For Heun, a user step completes with its second-order evaluation. The final
+                    # interval has none: Diffusers ends it with a first-order step to the terminal sigma.
                     if is_heun:
-                        if not in_first_order:
+                        if not in_first_order or step_index == num_scheduler_steps - 1:
                             user_step += 1
                             # Only call step_callback if we haven't exceeded total_steps
                             if user_step <= total_steps:
                                 pbar.update(1)
                                 step_callback(
                                     PipelineIntermediateState(
                                         step=user_step,
-                                        order=2,
+                                        # Progress over completed user steps, not the solver's order.
+                                        order=1,
                                         total_steps=total_steps,
                                         timestep=int(sigma_curr * 1000),
                                         latents=latents,

invokeai/backend/flux/denoise.py

@@ -13,7 +13,11 @@
 from invokeai.backend.flux.extensions.xlabs_controlnet_extension import XLabsControlNetExtension
 from invokeai.backend.flux.extensions.xlabs_ip_adapter_extension import XLabsIPAdapterExtension
 from invokeai.backend.flux.model import Flux
-from invokeai.backend.flux.schedulers import is_flow_match_lcm_scheduler, set_heun_timesteps_from_sigmas
+from invokeai.backend.flux.schedulers import (
+    is_flow_match_lcm_scheduler,
+    set_heun_timesteps_from_sigmas,
+    set_timesteps_from_sigmas,
+)
 from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.util.devices import TorchDevice
@@ -59,10 +63,7 @@ def denoise(
         set_heun_timesteps_from_sigmas(scheduler, timesteps[:-1], timesteps[-1], img.device)
     elif isinstance(scheduler, FlowMatchEulerDiscreteScheduler):
         # InvokeAI supplies model-evaluation sigmas plus its requested terminal endpoint.
-        scheduler.set_timesteps(sigmas=timesteps[:-1], device=img.device)
-        scheduler.sigmas[-1] = torch.as_tensor(
-            timesteps[-1], dtype=scheduler.sigmas.dtype, device=scheduler.sigmas.device
-        )
+        set_timesteps_from_sigmas(scheduler, timesteps[:-1], timesteps[-1], img.device)
     elif is_flow_match_lcm_scheduler(scheduler):
         # LCM owns its schedule; clipping controls the requested evaluation count,
         # while its exact start/end sigmas remain scheduler-defined.

invokeai/backend/flux/schedulers.py

@@ -68,6 +68,19 @@ def set_heun_timesteps_from_sigmas(
     scheduler.sample = None
 
 
+def set_timesteps_from_sigmas(
+    scheduler: SchedulerMixin,
+    sigmas: list[float],
+    terminal_sigma: float,
+    device: str | torch.device | None,
+) -> None:
+    """Set InvokeAI's shifted model sigmas on a scheduler whose set_timesteps() takes them (Euler, LCM)."""
+    # set_timesteps() appends a terminal sigma of 0, and step() reads sigmas[step_index + 1], so the
+    # requested terminal sigma has to replace it for a pass that ends before the schedule does.
+    scheduler.set_timesteps(sigmas=sigmas, device=device)
+    scheduler.sigmas[-1] = torch.as_tensor(terminal_sigma, dtype=scheduler.sigmas.dtype, device=scheduler.sigmas.device)
+
+
 def is_flow_match_lcm_scheduler(scheduler: SchedulerMixin) -> bool:
     """Check the optional LCM scheduler without requiring it in older Diffusers builds."""
     return _HAS_LCM and isinstance(scheduler, FlowMatchLCMScheduler)