huggingface / huggingface/diffusers

Shape mismatch when loading a Flux2.Klein GGUF

Open
#13,001 4 comments 2 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### 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

```shell
Traceback (most recent call last):
File "test.py", line 8, in
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

Contributor guide

Open the contributing guide

Research direction

Reproduce the failure with Flux2Transformer2DModel.from_single_file and the provided GGUF example, then trace the shape handling through loaders/single_file_model.py, models/model_loading_utils.py, and quantizers/gguf/gguf_quantizer.py. Compare the inferred empty-state shape for Flux2.Klein with the loaded parameter shape; done means the Klein checkpoint loads without the reported mismatch.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.