Flux1-Dev inference with single file ComfyUI/SD-Forge Safetensors

[ddpasa]

Is it possible to run inference with diffusers using a single-file safetensors created for ComfyUI/SD-Forge?

It looks like FluxPipeline.from_single_file() might be intended for this purpose, but I'm getting the following errors:

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_single_file("./flux1-dev-fp8.safetensors", torch_dtype=torch.float8_e4m3fn, use_safetensors=True)

Traceback (most recent call last):
  File "/home/user/flux/imgen.py", line 9, in <module>
    pipe = FluxPipeline.from_single_file("./flux1-dev-fp8.safetensors", torch_dtype=torch.float8_e4m3fn, use_safetensors=True)
  File "/home/user/.local/lib/python3.13/site-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
    return fn(*args, **kwargs)
  File "/home/user/.local/lib/python3.13/site-packages/diffusers/loaders/single_file.py", line 509, in from_single_file
    loaded_sub_model = load_single_file_sub_model(
        library_name=library_name,
    ...<11 lines>...
        **kwargs,
    )
  File "/home/user/.local/lib/python3.13/site-packages/diffusers/loaders/single_file.py", line 127, in load_single_file_sub_model
    loaded_sub_model = create_diffusers_t5_model_from_checkpoint(
        class_obj,
    ...<4 lines>...
        local_files_only=local_files_only,
    )
  File "/home/user/.local/lib/python3.13/site-packages/diffusers/loaders/single_file_utils.py", line 2156, in create_diffusers_t5_model_from_checkpoint
    model.load_state_dict(diffusers_format_checkpoint)
    ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/.local/lib/python3.13/site-packages/torch/nn/modules/module.py", line 2641, in load_state_dict
    raise RuntimeError(
    ...<3 lines>...
    )
RuntimeError: Error(s) in loading state_dict for T5EncoderModel:
	Missing key(s) in state_dict: "encoder.embed_tokens.weight". 

I checked the safetensors file and the T5 encoder is present. However, it is named differently, which confuses diffusers.

[DN6]

@ddpasa It's possible to load the Flux transformer individually from the the file.

from diffusers import FluxTransformer2DModel
model = FluxTransformer2DModel.from_single_file("./flux1-dev-fp8.safetensors", torch_dtype=torch.float8_e4m3fn)

Just FYI, diffusers doesn't handle automatic casting during inference if the weights are loaded in FP8. To enable that behaviour, you should use layerwise casting

from diffusers import FluxTransformer2DModel
model = FluxTransformer2DModel.from_single_file("./flux1-dev-fp8.safetensors", torch_dtype=torch.bfloat16)
model.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)

Re: The text encoder, I'm afraid if the keys are different, the pipeline won't be able to load the T5 encoder. You could either convert the keys to match the HF T5 Encoder or use the model in the official repo.

import torch
from diffusers import FluxPipeline
from transformers import T5EncoderModel

ckpt_path = "<path to checkpoint file>"
pipe = FluxPipeline.from_single_file(
    ckpt_path,
    text_encoder_2=None,
    torch_dtype=torch.bfloat16
)
pipe.text_encoder_2 = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    torch_dtype=torch.bfloat16
)

[ddpasa]

thanks @DN6 !

This is not the original Flux model, but a finetune trained with sd-scripts. That is why I'm trying to load it this way.