Loading Flux Dev transformers from_single_file fails since support for Flux Fill was added

[dxqbYD]

Describe the bug

Flux models are sometimes distributed as transformer only, because all the other pipeline components are usually not changed.

Examples include:

• finetunes on CivitAI, to offer a smaller download size

• https://huggingface.co/lllyasviel/flux1-dev-bnb-nf4/tree/main

Since support for Flux Fill was added here 862a7d5, diffusers.FluxTransformer2DModel.from_single_file fails to load these files:

File "/.../src/diffusers/src/diffusers/loaders/single_file_utils.py", line 598, in infer_diffusers_model_type
    if checkpoint["img_in.weight"].shape[1] == 384:
       ~~~~~~~~~~^^^^^^^^^^^^^^^^^
KeyError: 'img_in.weight'

Reproduction

Download one of the example mentioned above, and run

from diffusers import FluxTransformer2DModel

transformer = FluxTransformer2DModel.from_single_file(your_file)

Logs

No response

System Info

HEAD

Who can help?

@DN6

[dxqbYD]

the bnb-nf4 checkpoint linked above is apparently not only the transformer - but it has the same problem anyway. Other affected transformer-only files can be found on CivitAI, or generated using this Comfy workflow: https://civitai.com/models/982277/flux-lora-merge

[DN6]

@dxqbYD I'm not sure the bnb checkpoint will load. We don't support bnb single checkpoints at the moment. But for the others, can you try loading with this branch?
#10545

[DN6]

@dxqbYD Do you have a link to a specific checkpoint that is failing to load that isn't a bnb NF4 checkpoint?

[dxqbYD]

@dxqbYD I'm not sure the bnb checkpoint will load. We don't support bnb single checkpoints at the moment. But for the others, can you try loading with this branch? #10545

Can confirm that this PR works on all my test cases, which are:

• a civitai checkpoint that offers downloads for 20 GB (full pipeline) and another file with 11 GB (transformer only I suppose). The 11 GB file can be loaded now
• a file that I've created myself using https://civitai.com/models/982277/flux-lora-merge
• https://huggingface.co/lllyasviel/flux1-dev-bnb-nf4/tree/main fails as you have expected, but now for a different reason

Thanks for your quick reaction!

[yiyixuxu]

closing as fixed in #10545