huggingface / huggingface/diffusers

Flux.1 cannot load standard transformer in nf4

Open
#9,996 16 comments 0 reactions 1 assignee Claimed by @DN6 View on GitHub
bug wip
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### Describe the bug

loading different flux transformer models is fine except for nf4.
it works for 1% of fine-tunes provided on Huggingface, but it doesn't work for 99% standard fine-tunes available on CivitAI.

example of such model:

*note* i'm using `FluxTransformer2DModel` directly as its easiest for reproduction plus majority of flux fine-tunes are provided as transformer-only, not full models. but where full model does exist, its exactly the same problem using `FluxPipeline`

### Reproduction

```py
import torch
import bitsandbytes as bnb
import diffusers

print(f'torch=={torch.__version__} diffusers=={diffusers.__version__} bnb=={bnb.__version__}')
kwargs = { 'low_cpu_mem_usage': True, 'torch_dtype': torch.bfloat16, 'cache_dir': '/mnt/models/huggingface' }
files = [
'flux-c4pacitor_v2alpha-f1s-bf16.safetensors',
'flux-iniverse_v2-f1d-fp8.safetensors',
'flux-copax_timeless_xplus_mix2-nf4.safetensors',
]

for f in files:
print(f)
try:
transformer = diffusers.FluxTransformer2DModel.from_single_file(f, **kwargs)
print(transformer.__class__)
except Exception as e:
print(e)
transformer = None
torch.cuda.empty_cache()
```

### Logs

```shell
in `diffusers/loaders/single_file_utils.py:convert_flux_transformer_checkpoint_to_diffusers`

q, k, v, mlp = torch.split(checkpoint.pop(f"single_blocks.{i}.linear1.weight"), split_size, dim=0)

> RuntimeError: split_with_sizes expects split_sizes to sum exactly to 33030144 (input tensor's size at dimension 0), but got split_sizes=[3072, 3072, 3072, 12288]
```

### System Info

torch==2.5.1+cu124 diffusers==0.32.0.dev0 bnb==0.44.1

### Who can help?

@yiyixuxu @sayakpaul @DN6 @asomoza

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.