Error in loading flux fp8 model with local transformer_flux.py file

[DhavalTaunk08]

Describe the bug

Unable to use flux fp8 model from Kijai/flux-fp8 while having transformer_flux.py file in local. I have modified the scripts to remove any import error. I put some print statements in single_model_file.py to check why it is not loading the model.

Reproduction

The below code works fine.

single_model_file.py

def _get_single_file_loadable_mapping_class(cls):
    print(cls)
    diffusers_module = importlib.import_module(__name__.split(".")[0])

    for loadable_class_str in SINGLE_FILE_LOADABLE_CLASSES:
        loadable_class = getattr(diffusers_module, loadable_class_str)
        print(cls, loadable_class)
        print(issubclass(cls, loadable_class))
        if issubclass(cls, loadable_class):
            return loadable_class_str

    return None

from diffusers import FluxTransformer2DModel
transformer = FluxTransformer2DModel.from_single_file(
    "https://huggingface.co/Kijai/flux-fp8/blob/main/flux1-schnell-fp8-e4m3fn.safetensors", 
    torch_dtype=torch.bfloat16
)

I am getting the below output:

<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'>
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.unets.unet_stable_cascade.StableCascadeUNet'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.unets.unet_2d_condition.UNet2DConditionModel'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.controlnet.ControlNetModel'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.transformers.transformer_sd3.SD3Transformer2DModel'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.unets.unet_motion_model.MotionAdapter'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.controlnet_sparsectrl.SparseControlNetModel'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'>
True

But while using the class from my local code:

from transformer_flux import FluxTransformer2DModel
FluxTransformer2DModel.__module__ = 'diffusers.models.transformers.transformer_flux'
transformer = FluxTransformer2DModel.from_single_file(
    "https://huggingface.co/Kijai/flux-fp8/blob/main/flux1-schnell-fp8-e4m3fn.safetensors", 
    torch_dtype=torch.bfloat16
)

It is giving me following error:

<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'>
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.unets.unet_stable_cascade.StableCascadeUNet'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.unets.unet_2d_condition.UNet2DConditionModel'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.controlnet.ControlNetModel'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.transformers.transformer_sd3.SD3Transformer2DModel'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.unets.unet_motion_model.MotionAdapter'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.controlnet_sparsectrl.SparseControlNetModel'>
False
<class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'> <class 'diffusers.models.transformers.transformer_flux.FluxTransformer2DModel'>
False

Traceback (most recent call last):
  File "/workspace/GarmentTransferV2/test.py", line 441, in <module>
    main(args)
  File "/workspace/GarmentTransferV2/test.py", line 368, in main
    transformer_garment = FluxTransformerGarment2DModel.from_single_file(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/garment/lib/python3.11/site-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/workspace/garment/lib/python3.11/site-packages/diffusers/loaders/single_file_model.py", line 182, in from_single_file
    raise ValueError(
ValueError: FromOriginalModelMixin is currently only compatible with StableCascadeUNet, UNet2DConditionModel, AutoencoderKL, ControlNetModel, SD3Transformer2DModel, MotionAdapter, SparseControlNetModel, FluxTransformer2DModel

Any leads would be appreciated.

Logs

No response

System Info

• 🤗 Diffusers version: 0.30.3
• Platform: Linux-6.8.0-40-generic-x86_64-with-glibc2.35
• Running on Google Colab?: No
• Python version: 3.11.9
• PyTorch version (GPU?): 2.4.1+cu121 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.25.2
• Transformers version: 4.45.2
• Accelerate version: 1.0.1
• PEFT version: not installed
• Bitsandbytes version: not installed
• Safetensors version: 0.4.5
• xFormers version: not installed
• Accelerator: NVIDIA H100 80GB HBM3, 81559 MiB
• Using GPU in script?:

Who can help?

@DN6 @sayakpaul

[a-r-r-o-w]

You cannot pass URLs to from_single_file. Could you try the following instead?

from diffusers import FluxTransformer2DModel
from huggingface_hub import hf_hub_download

safetensors_file = hf_hub_download("Kijai/flux-fp8", filename="flux1-dev-fp8-e4m3fn.safetensors")
print(safetensors_file)

transformer = FluxTransformer2DModel.from_single_file(safetensors_file, subfolder="transformer")
print(transformer.config)

Edit: My bad, you can pass URLs to from_single_file, I got confused with something else. I think you just have to pass subfolder="transformer" when initializing to make your snippet work. This is because we try to fetch the init config from the original Flux-Dev repository (as Kijai/flux-fp8 is based on Flux-Dev).

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.