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