[bug]: Using Heun (2nd order) with a reference image in Flux 2 Klein 9B causes EinopsError

[AuralVertigo]

Is there an existing issue for this problem?

• I have searched the existing issues

Install method

Invoke's Launcher

Operating system

Windows

GPU vendor

Nvidia (CUDA)

GPU model

RTX 4090

GPU VRAM

24GB

Version number

v6.13.0.rc3

Browser

Vivaldi 8.0.4033.28 / Chromium 148.0.7778.183

System Information

No response

What happened

When trying to generate an image using the starter model FLUX.2 Klein 9B (GGUF Q8) in conjunction with a reference image and the Heun (2nd order) scheduler, the error below is being raised. The other schedulers, Euler and LCM, work as intended. Generations with Heun (2nd order) only fail when using a reference image, regular generation jobs are not affected. The error below was caused when using an input image with the size 1024x1280px and a desired output with the same dimensions.

2026-05-22 15:48:28,583]::[InvokeAI]::ERROR --> Error while invoking session 342ab238-c44f-4f34-af06-a7d8aee242ae, invocation 8aefefcf-da7e-4a15-b507-ec84c8383980 (flux2_denoise):  Error while processing rearrange-reduction pattern "b (h w) (c ph pw) -> b c (h ph) (w pw)".
 Input tensor shape: torch.Size([1, 10240, 128]). Additional info: {'h': 80, 'w': 64, 'ph': 2, 'pw': 2}.
 Shape mismatch, 10240 != 5120
[2026-05-22 15:48:28,583]::[InvokeAI]::ERROR --> Traceback (most recent call last):
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\invokeai\app\services\session_processor\session_processor_default.py", line 130, in run_node
    output = invocation.invoke_internal(context=context, services=self._services)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\invokeai\app\invocations\baseinvocation.py", line 244, in invoke_internal
    output = self.invoke(context)
             ^^^^^^^^^^^^^^^^^^^^
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\torch\utils\_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\invokeai\app\invocations\flux2_denoise.py", line 229, in invoke
    latents = self._run_diffusion(context)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\invokeai\app\invocations\flux2_denoise.py", line 470, in _run_diffusion
    x = denoise(
        ^^^^^^^^
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\invokeai\backend\flux2\denoise.py", line 191, in denoise
    step_callback(
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\invokeai\app\invocations\flux2_denoise.py", line 536, in step_callback
    state.latents = unpack_flux2(latents, self.height, self.width).squeeze()
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\invokeai\backend\flux2\sampling_utils.py", line 100, in unpack_flux2
    return rearrange(
           ^^^^^^^^^^
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\einops\einops.py", line 616, in rearrange
    return reduce(tensor, pattern, reduction="rearrange", **axes_lengths)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "H:\InvokeAI_Managed\.venv\Lib\site-packages\einops\einops.py", line 550, in reduce
    raise EinopsError(message + f"\n {e}") from None
einops.EinopsError:  Error while processing rearrange-reduction pattern "b (h w) (c ph pw) -> b c (h ph) (w pw)".
 Input tensor shape: torch.Size([1, 10240, 128]). Additional info: {'h': 80, 'w': 64, 'ph': 2, 'pw': 2}.
 Shape mismatch, 10240 != 5120

What you expected to happen

Generations using FLUX.2 Klein 9B (GGUF Q8) with the Heun (2nd order) scheduler and a reference image should work without issues in the same way as if using Euler or LCM.

How to reproduce the problem

1. Select model "FLUX.2 Klein 9B (GGUF Q8)" (or probably any other Flux 2 Klein variant), "FLUX.2 VAE" as the VAE and "FLUX.2 Klein Qwen3 8B Encoder" as the Qwen3 Encoder.
2. Add a prompt and a reference image. The size of this image doesn't seem to matter as the error is raised both with a matching and different size as the desired output image (e.g. 1024x1280px -> 1024x1280px or 768x1024px -> 1024x1024px)
3. Select "Heun (2nd order)" and hit "Invoke".

Additional context

No response

Discord username

No response