From 438eea115900dc102fd061486576c63d5e3a577a Mon Sep 17 00:00:00 2001
From: Alexander Eichhorn <alex@eichhorn.dev>
Date: Sun, 1 Feb 2026 22:25:39 +0100
Subject: [PATCH] fix(flux2): support Heun scheduler for FLUX.2 Klein models
 (#8794)

* fix(flux2): support Heun scheduler for FLUX.2 Klein models

FlowMatchHeunDiscreteScheduler does not support dynamic shifting parameters
(use_dynamic_shifting, base_shift, max_shift, etc.) or sigmas/mu in set_timesteps.
This caused FLUX.2 Klein to fail when using Heun scheduler.

- Create Heun scheduler with only num_train_timesteps and shift parameters
- Use num_inference_steps instead of sigmas for Heun's set_timesteps call
- Euler and LCM schedulers continue to use full dynamic shifting support

* fix(flux2): fix Heun scheduler detection using inspect.signature

The previous hasattr check for state_in_first_order failed because
the attribute doesn't exist before set_timesteps() is called. Now
using inspect.signature to check for sigmas parameter support,
matching the FLUX1 implementation.

---------

Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com>
---
 invokeai/app/invocations/flux2_denoise.py | 28 +++++++++++++++--------
 invokeai/backend/flux2/denoise.py         | 14 +++++++++---
 2 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py
index 6a85e58f74..f221c1a3c4 100644
--- a/invokeai/app/invocations/flux2_denoise.py
+++ b/invokeai/app/invocations/flux2_denoise.py
@@ -366,16 +366,24 @@ class Flux2DenoiseInvocation(BaseInvocation):
         if self.scheduler in FLUX_SCHEDULER_MAP and not is_inpainting:
             # Only use scheduler for txt2img - use manual Euler for inpainting to preserve exact timesteps
             scheduler_class = FLUX_SCHEDULER_MAP[self.scheduler]
-            scheduler = scheduler_class(
-                num_train_timesteps=1000,
-                shift=3.0,
-                use_dynamic_shifting=True,
-                base_shift=0.5,
-                max_shift=1.15,
-                base_image_seq_len=256,
-                max_image_seq_len=4096,
-                time_shift_type="exponential",
-            )
+            # 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(
+                    num_train_timesteps=1000,
+                    shift=3.0,
+                    use_dynamic_shifting=True,
+                    base_shift=0.5,
+                    max_shift=1.15,
+                    base_image_seq_len=256,
+                    max_image_seq_len=4096,
+                    time_shift_type="exponential",
+                )
 
         # Prepare reference image extension for FLUX.2 Klein built-in editing
         ref_image_extension = None
diff --git a/invokeai/backend/flux2/denoise.py b/invokeai/backend/flux2/denoise.py
index 6ac20be5a9..7b5bd6194e 100644
--- a/invokeai/backend/flux2/denoise.py
+++ b/invokeai/backend/flux2/denoise.py
@@ -4,6 +4,7 @@ This module provides the denoising function for FLUX.2 Klein models,
 which use Qwen3 as the text encoder instead of CLIP+T5.
 """
 
+import inspect
 import math
 from typing import Any, Callable
 
@@ -87,11 +88,18 @@ def denoise(
         # The scheduler will apply dynamic shifting internally using mu (if enabled in scheduler config)
         sigmas = np.array(timesteps[:-1], dtype=np.float32)  # Exclude final 0.0
 
-        # Pass mu if provided - it will only be used if scheduler has use_dynamic_shifting=True
-        if mu is not None:
+        # Check if scheduler supports sigmas parameter using inspect.signature
+        # FlowMatchHeunDiscreteScheduler and FlowMatchLCMScheduler don't support sigmas
+        set_timesteps_sig = inspect.signature(scheduler.set_timesteps)
+        supports_sigmas = "sigmas" in set_timesteps_sig.parameters
+        if supports_sigmas and mu is not None:
+            # Pass mu if provided - it will only be used if scheduler has use_dynamic_shifting=True
             scheduler.set_timesteps(sigmas=sigmas.tolist(), mu=mu, device=img.device)
-        else:
+        elif supports_sigmas:
             scheduler.set_timesteps(sigmas=sigmas.tolist(), device=img.device)
+        else:
+            # Scheduler doesn't support sigmas (e.g., Heun, LCM) - use num_inference_steps
+            scheduler.set_timesteps(num_inference_steps=len(sigmas), device=img.device)
         num_scheduler_steps = len(scheduler.timesteps)
         is_heun = hasattr(scheduler, "state_in_first_order")
         user_step = 0