hlhc/FLUX.2-klein-4B-onnx

FLUX.2 [klein] 4B — ONNX

black-forest-labs/FLUX.2-klein-4B exported to ONNX for onnxruntime, with the pipeline's pre- and post-processing folded into the graphs: text encoder MatMulNBits int4 block 128, transformer MatMulNBits int8 block 128, VAE float32; activations float32. Same licence as the original (Apache-2.0, LICENSE.md).

graph	inputs → output
text_encoder.onnx	input_ids[1,512], attention_mask[1,512] → prompt_embeds[1,512,7680] (Qwen3 hidden states of layers 9/18/27)
transformer.onnx	hidden_states[1,S,128], encoder_hidden_states[1,512,7680], timestep[1] (= sigma), img_ids[1,S,4], txt_ids[1,512,4] → noise_pred[1,S,128]
vae_encoder.onnx	image[1,3,H,W] in [-1, 1] → latents[1,128,H/16,W/16] (patchified, batch-norm normalised)
vae_decoder.onnx	latents[1,128,h,w] → image[1,3,16h,16w] in [-1, 1]

config.json records the chat-template prefix/suffix, the position-id conventions, the schedule (flow-match Euler with the empirical mu shift), and the reference-image convention for editing. MatMulNBits runs on onnxruntime's CPU, CUDA and DirectML providers (not Core ML).

Files

file	size
text_encoder.onnx	0.00 GB
text_encoder.onnx.data	3.00 GB
transformer.onnx	0.00 GB
transformer.onnx.data	4.50 GB
vae_decoder.onnx	0.20 GB
vae_encoder.onnx	0.14 GB

Verification (cosine vs the PyTorch modules on identical inputs)

graph	cosine
text_encoder	0.99151
transformer	0.99789
vae_encoder	1.00000
vae_decoder	1.00000

verification_onnx.png is a 4-step, 512 px text-to-image through onnxruntime; verification_pytorch.png, when present, the same prompt through diffusers. Exported with scripts/export_klein_onnx.py (torch 2.14.0+cu130, diffusers 0.40.0, onnxruntime 1.30.0).