How to convert Flux DiT diffusers format safetensors into original format?

[Leo-A11A]

I have fine tuned a Flux Kontext model and want to convert it to its original format.

[maiconaalmeida]

Using Python Script (Recommended)You can create a conversion script that loads your Diffusers model and saves it as a single safetensors file. Here's the approach:

`"""
Convert Flux Kontext from Diffusers format to original single-file format
"""
import torch
from diffusers import FluxKontextPipeline, FluxTransformer2DModel
from safetensors.torch import save_file
import argparse

def convert_diffusers_to_original(
model_path,
output_path,
dtype=torch.bfloat16
):
"""
Convert a Flux Kontext model from Diffusers format to original single-file format

Args:
    model_path: Path to the Diffusers model folder
    output_path: Path where to save the output safetensors file
    dtype: Data type for the model (default: torch.bfloat16)
"""
print(f"Loading model from: {model_path}")

# Load the transformer model from the Diffusers format
# Only load the transformer part, as that's what gets fine-tuned
transformer = FluxTransformer2DModel.from_pretrained(
    model_path,
    subfolder="transformer",
    torch_dtype=dtype
)

print("Extracting state dict...")
state_dict = transformer.state_dict()

# Create metadata
metadata = {
    "format": "pt",
    "model_type": "flux_kontext"
}

print(f"Saving to: {output_path}")
save_file(state_dict, output_path, metadata=metadata)

print(f"✓ Conversion complete! Saved to {output_path}")
print(f"Model size: {len(state_dict)} tensors")

def convert_with_full_pipeline(
model_path,
output_path,
dtype=torch.bfloat16
):
"""
Alternative method: Load full pipeline and extract transformer
"""
print(f"Loading full pipeline from: {model_path}")

# Load the full pipeline
pipe = FluxKontextPipeline.from_pretrained(
    model_path,
    torch_dtype=dtype
)

print("Extracting transformer state dict...")
state_dict = pipe.transformer.state_dict()

# Create metadata
metadata = {
    "format": "pt",
    "model_type": "flux_kontext"
}

print(f"Saving to: {output_path}")
save_file(state_dict, output_path, metadata=metadata)

print(f"✓ Conversion complete! Saved to {output_path}")
print(f"Model size: {len(state_dict)} tensors")

if name == "main":
parser = argparse.ArgumentParser(description="Convert Flux Kontext from Diffusers to original format")
parser.add_argument(
"--model_path",
type=str,
required=True,
help="Path to the Diffusers format model folder"
)
parser.add_argument(
"--output_path",
type=str,
required=True,
help="Output path for the safetensors file (e.g., flux_kontext_finetuned.safetensors)"
)
parser.add_argument(
"--dtype",
type=str,
default="bfloat16",
choices=["bfloat16", "float16", "float32", "float8_e4m3fn"],
help="Data type for the model"
)
parser.add_argument(
"--full_pipeline",
action="store_true",
help="Use full pipeline method (loads entire pipeline)"
)

args = parser.parse_args()

# Convert dtype string to torch dtype
dtype_map = {
    "bfloat16": torch.bfloat16,
    "float16": torch.float16,
    "float32": torch.float32,
    "float8_e4m3fn": torch.float8_e4m3fn
}
dtype = dtype_map[args.dtype]

# Run conversion
if args.full_pipeline:
    convert_with_full_pipeline(args.model_path, args.output_path, dtype)
else:
    convert_diffusers_to_original(args.model_path, args.output_path, dtype)`

Usage Instructions
Basic Usage

python convert_flux_kontext.py
--model_path ./path/to/your/diffusers_model
--output_path ./flux_kontext_finetuned.safetensors

With Different Data Types

For FP16

python convert_flux_kontext.py
--model_path ./your_model
--output_path ./model.safetensors
--dtype float16

For FP8 (quantized)

python convert_flux_kontext.py
--model_path ./your_model
--output_path ./model_fp8.safetensors
--dtype float8_e4m3fn

OR WITH Quick Inline Conversion
If you prefer a quick one-off conversion, you can use this simpler approach:
`import torch
from diffusers import FluxTransformer2DModel
from safetensors.torch import save_file

Load your fine-tuned model

model_path = "./path/to/your/diffusers_model"
transformer = FluxTransformer2DModel.from_pretrained(
model_path,
subfolder="transformer",
torch_dtype=torch.bfloat16
)

Save as single file

state_dict = transformer.state_dict()
save_file(state_dict, "flux_kontext_finetuned.safetensors", metadata={"format": "pt"})`

Important Notes

What gets converted: Since Flux Kontext fine-tuning typically only trains the transformer/UNet component, the script extracts just that part. The VAE and text encoders remain unchanged.
Using the converted model: To use your converted single-file model:

`from diffusers import FluxKontextPipeline, FluxTransformer2DModel
import torch

Load the transformer from your converted file

transformer = FluxTransformer2DModel.from_single_file(
"./flux_kontext_finetuned.safetensors",
torch_dtype=torch.bfloat16
)

Load full pipeline with your custom transformer

pipe = FluxKontextPipeline.from_pretrained(
"black-forest-labs/FLUX.1-Kontext-dev",
transformer=transformer,
torch_dtype=torch.bfloat16
)`

Requirements: Make sure you have the latest diffusers installed:
pip install git+https://github.com/huggingface/diffusers.git pip install safetensors

compatible with tools like ComfyUI and other workflows.