Bug flux2 klein: ValueError: Got [32, 32, 32, 32] but expected positional dim 64

[Juste-Leo2]

I used Qwe3-4B-IQ4_NL.gguf with flux-2-klein-4b-Q4_K_K_M.gguf

my configuration:

• I7 8760H
• GTX 1070 (8Go)
• 16Go RAM
• windows 11
• torch 2.8.0+cu128

the workflow used:

• flux2-4b.json

Here are the error logs

Prompt executed in 56.38 seconds
FETCH ComfyRegistry Data: 100/119
FETCH ComfyRegistry Data: 105/119
FETCH ComfyRegistry Data: 110/119
FETCH ComfyRegistry Data: 115/119
FETCH ComfyRegistry Data [DONE]
[ComfyUI-Manager] default cache updated: https://api.comfy.org/nodes
FETCH DATA from: https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/custom-node-list.json [DONE]
[ComfyUI-Manager] All startup tasks have been completed.
got prompt
Using pytorch attention in VAE
Using pytorch attention in VAE
VAE load device: cuda:0, offload device: cpu, dtype: torch.float32
gguf qtypes: F32 (145), Q6_K (1), IQ4_NL (216), Q5_K (36)
Dequantizing token_embd.weight to prevent runtime OOM.
CLIP/text encoder model load device: cuda:0, offload device: cpu, current: cpu, dtype: torch.float16
Requested to load ZImageTEModel_
loaded completely; 5742.68 MB usable, 2702.94 MB loaded, full load: True
gguf qtypes: F32 (60), Q5_K (30), Q4_K (50), BF16 (9)
!!! Exception during processing !!! Got [32, 32, 32, 32] but expected positional dim 64
Traceback (most recent call last):
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\execution.py", line 510, in execute
    output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs)
                                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\execution.py", line 324, in get_output_data
    return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, hidden_inputs=hidden_inputs)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\execution.py", line 298, in _async_map_node_over_list
    await process_inputs(input_dict, i)
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\execution.py", line 286, in process_inputs
    result = f(**inputs)
             ^^^^^^^^^^^
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-GGUF\nodes.py", line 176, in load_unet
    model = comfy.sd.load_diffusion_model_state_dict(
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\comfy\sd.py", line 1431, in load_diffusion_model_state_dict
    model = model_config.get_model(new_sd, "")
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\comfy\supported_models.py", line 767, in get_model
    out = model_base.Flux2(self, device=device)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\comfy\model_base.py", line 849, in __init__
    super().__init__(model_config, model_type, device=device, unet_model=unet_model)
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\comfy\model_base.py", line 140, in __init__
    self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\ComfyUI_windows_portable_nvidia_cu128\ComfyUI_windows_portable\ComfyUI\comfy\ldm\flux\model.py", line 63, in __init__
    raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
ValueError: Got [32, 32, 32, 32] but expected positional dim 64

[jprsyt5]

got prompt
gguf qtypes: F32 (145), Q6_K (1), IQ4_NL (216), Q5_K (36)
Dequantizing token_embd.weight to prevent runtime OOM.
CLIP/text encoder model load device: cuda:0, offload device: cpu, current: cpu, dtype: torch.float16
Requested to load Flux2TEModel_
loaded completely; 5598.55 MB usable, 2702.94 MB loaded, full load: True
gguf qtypes: F32 (60), Q5_K (30), Q4_K (50), BF16 (9)
model weight dtype torch.float16, manual cast: None
model_type FLUX
Requested to load Flux2
Unloaded partially: 1627.08 MB freed, 1075.86 MB remains loaded, 93.52 MB buffer reserved, lowvram patches: 0
loaded completely; 4143.52 MB usable, 2591.64 MB loaded, full load: True
100%|████████████████████████████████████████████████████████████████████████████████████| 8/8 [00:08<00:00,  1.01s/it]
Requested to load AutoencoderKL
Unloaded partially: 670.14 MB freed, 1921.54 MB remains loaded, 60.75 MB buffer reserved, lowvram patches: 0
loaded completely; 248.21 MB usable, 160.31 MB loaded, full load: True
Prompt executed in 24.48 seconds

Can’t reproduce it, it works fine for me using the wf you provided & same model. Are you sure you’ve already updated your Comfy?

Because in your logs it still detects Qwen3 as ZImageTEModel_, while both ZImage and Flux 2 Klein 4B use the same TE. There’s a different code path when you use Flux 2, in your logs it should show up as Flux2TEModel_

elif te_model == TEModel.QWEN3_4B:
            if clip_type == CLIPType.FLUX or clip_type == CLIPType.FLUX2:
                clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type="qwen3_4b")
                clip_target.tokenizer = comfy.text_encoders.flux.KleinTokenizer
            else:
                clip_target.clip = comfy.text_encoders.z_image.te(**llama_detect(clip_data))
                clip_target.tokenizer = comfy.text_encoders.z_image.ZImageTokenizer

People really need to understand that when a new model comes out and Comfy announces support for it, you still have to update your local Comfy, either by doing git fetch + git pull, using Comfy Manager, or running the .bat script provided by Comfy. It won’t magically updated in your local setup.

[Juste-Leo2]

got prompt
gguf qtypes: F32 (145), Q6_K (1), IQ4_NL (216), Q5_K (36)
Dequantizing token_embd.weight to prevent runtime OOM.
CLIP/text encoder model load device: cuda:0, offload device: cpu, current: cpu, dtype: torch.float16
Requested to load Flux2TEModel_
loaded completely; 5598.55 MB usable, 2702.94 MB loaded, full load: True
gguf qtypes: F32 (60), Q5_K (30), Q4_K (50), BF16 (9)
model weight dtype torch.float16, manual cast: None
model_type FLUX
Requested to load Flux2
Unloaded partially: 1627.08 MB freed, 1075.86 MB remains loaded, 93.52 MB buffer reserved, lowvram patches: 0
loaded completely; 4143.52 MB usable, 2591.64 MB loaded, full load: True
100%|████████████████████████████████████████████████████████████████████████████████████| 8/8 [00:08<00:00, 1.01s/it]
Requested to load AutoencoderKL
Unloaded partially: 670.14 MB freed, 1921.54 MB remains loaded, 60.75 MB buffer reserved, lowvram patches: 0
loaded completely; 248.21 MB usable, 160.31 MB loaded, full load: True
Prompt executed in 24.48 seconds
Can’t reproduce it, it works fine for me using the wf you provided & same model. Are you sure you’ve already updated your Comfy?

Because in your logs it still detects Qwen3 as ZImageTEModel_, while both ZImage and Flux 2 Klein 4B use the same TE. There’s a different code path when you use Flux 2, in your logs it should show up as Flux2TEModel_

elif te_model == TEModel.QWEN3_4B:
if clip_type == CLIPType.FLUX or clip_type == CLIPType.FLUX2:
clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type="qwen3_4b")
clip_target.tokenizer = comfy.text_encoders.flux.KleinTokenizer
else:
clip_target.clip = comfy.text_encoders.z_image.te(**llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.z_image.ZImageTokenizer
People really need to understand that when a new model comes out and Comfy announces support for it, you still have to update your local Comfy, either by doing git fetch + git pull, using Comfy Manager, or running the .bat script provided by Comfy. It won’t magically updated in your local setup.

Thanks for your reply. I updated comyUI GGUF but forgot about native comfyUI. It worked. Thank you for resolving the issue :)