Shape mismatch when loading a Flux2.Klein GGUF

[dxqb]

Describe the bug

When trying to load a Flux2.Klein GGUF:

ValueError: double_stream_modulation_img.linear.weight has an expected quantized shape of: (18432, 3072), but received shape: torch.Size([18432, 6144])

Reproducer and full error low below.

The inferred quantized shape of (18432, 3072) is the correct shape of this key.

(18432, 6144) is the shape of this key in Flux2.Dev, but not in Flux2.Klein. Is the empty shape dict initialized as Flux2.Dev even though Flux2.Klein is loaded?

Reproduction

import torch

from diffusers import FluxPipeline, Flux2Transformer2DModel, GGUFQuantizationConfig

ckpt_path = (
    "https://huggingface.co/unsloth/FLUX.2-klein-base-4B-GGUF/blob/main/flux-2-klein-base-4b-Q2_K.gguf"
)
transformer = Flux2Transformer2DModel.from_single_file(
    ckpt_path,
    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
    torch_dtype=torch.bfloat16,
)

Logs

Traceback (most recent call last):
  File "test.py", line 8, in <module>
    transformer = Flux2Transformer2DModel.from_single_file(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.12/site-packages/diffusers/loaders/single_file_model.py", line 491, in from_single_file
    load_model_dict_into_meta(
  File "venv/lib/python3.12/site-packages/diffusers/models/model_loading_utils.py", line 291, in load_model_dict_into_meta
    hf_quantizer.check_quantized_param_shape(param_name, empty_state_dict[param_name], param)
  File "venv/lib/python3.12/site-packages/diffusers/quantizers/gguf/gguf_quantizer.py", line 85, in check_quantized_param_shape
    raise ValueError(
ValueError: double_stream_modulation_img.linear.weight has an expected quantized shape of: (18432, 3072), but received shape: torch.Size([18432, 6144])

System Info

diffusers HEAD

Who can help?

@DN6 @yiyixuxu @sayakpaul