kny/FLUX.2-klein-4B-onnx-web

FLUX.2 [klein] 4B for ONNX Runtime Web (WebGPU)

This is an ONNX conversion of black-forest-labs/FLUX.2-klein-4B, made for in-browser inference with onnxruntime-web on WebGPU. FLUX.2 [klein] 4B is by Black Forest Labs and is released under Apache-2.0. Its text encoder is a Qwen3-4B derivative by Alibaba Cloud / Qwen, also released under Apache-2.0. Kim Nyberg converted the model to ONNX. See NOTICE.

Graph	File(s)	Precision	Size
Text encoder (Qwen3-4B, layers 0–26)	text_encoder/text_encoder_q4.onnx + .onnx_data, .onnx_data_1	q4 MatMulNBits (block 32, symmetric, accuracy_level 4); token embedding fp16	2.48 GB
Transformer / denoiser	transformer/transformer_q4.onnx + .onnx_data, .onnx_data_1	q4 MatMulNBits (block 32, symmetric); AdaLN modulation / time-embedding Gemms, norms and biases stay fp32	3.00 GB
VAE decoder	vae_decoder/vae_decoder_fp16.onnx + .onnx_data	fp16 weights and compute, fp32 I/O	0.10 GB

Every file is under 2 GB (the largest is 1.80 GB). All graphs use opset 18 (ai.onnx) plus com.microsoft.MatMulNBits. tokenizer/ (Qwen2 tokenizer + chat template) and scheduler/scheduler_config.json are copied unchanged from the original repo.

I/O

text_encoder_q4.onnx

name	dir	dtype	shape
input_ids	in	int64	[batch, seq]
attention_mask	in	int64	[batch, seq] (1 = token, 0 = pad)
prompt_embeds	out	float32	[batch, seq, 7680]

prompt_embeds is concat(hidden_states[9], hidden_states[18], hidden_states[27]) along the last dimension, exactly as Flux2KleinPipeline._get_qwen3_prompt_embeds computes it. Only the first 28 decoder layers are kept, and layer 27 is pruned as dead code, so the graph runs 27 layers. To match the reference, build the prompt with the chat template ([{"role":"user","content": prompt}], add_generation_prompt=true, enable_thinking=false) and right-pad to 512 tokens (max_sequence_length = 512).

transformer_q4.onnx (one denoising step, no CFG; klein is guidance-distilled)

name	dir	dtype	shape
hidden_states	in	float32	[1, img_seq, 128] packed latent tokens (img_seq = (H/16)·(W/16))
encoder_hidden_states	in	float32	[1, txt_seq, 7680] (= prompt_embeds)
timestep	in	float32	[1] sigma-time / 1000 (that is, scheduler timestep / 1000)
img_ids	in	int64	[img_seq, 4] rows (0, h, w, 0), raster order
txt_ids	in	int64	[txt_seq, 4] rows (0, 0, 0, l)
noise_pred	out	float32	[1, img_seq, 128]

The token dimensions are dynamic, so one graph serves 512² (1024 image tokens), 1024² (4096 image tokens) and other sizes that are multiples of 16. The batch dimension is fixed at 1 by the export. RoPE is computed in-graph in fp32; the original uses fp64, which WebGPU lacks, and the difference is at most 2e-5.

vae_decoder_fp16.onnx

name	dir	dtype	shape
latents	in	float32	[batch, 128, H/16, W/16] (packed latent tokens reshaped to NCHW, still BN-normalised)
image	out	float32	[batch, 3, H, W] RGB in [0, 1]

The graph folds in the pipeline's BatchNorm de-normalisation (x·sqrt(var+eps)+mean), the 2×2 un-patchify and the (x/2+0.5).clamp(0,1) post-processing.

Sampling loop (what the JS side must do)

1. Tokenize → text encoder → prompt_embeds ([1, 512, 7680]).
2. latents = N(0, 1) with shape [1, 128, H/16, W/16]. Flatten to tokens [1, (H/16)(W/16), 128] (channel-last, raster order).
3. FlowMatch Euler, 4 steps: sigmas = linspace(1, 1/steps, steps). Apply the exponential time shift with mu = compute_empirical_mu(img_seq, steps) (see diffusers pipeline_flux2_klein.py): sigma' = e^mu / (e^mu + (1/sigma - 1)), t = sigma'·1000. Each step runs noise_pred = transformer(...), then latents += (sigma_next − sigma)·noise_pred, with sigma_next = 0 after the last step.
4. Reshape tokens back to [1, 128, H/16, W/16] → VAE decoder → image.

Loading in onnxruntime-web

import * as ort from 'onnxruntime-web/webgpu';
const base = 'https://huggingface.co/kny/FLUX.2-klein-4B-onnx-web/resolve/main/';
const sess = await ort.InferenceSession.create(base + 'transformer/transformer_q4.onnx', {
  executionProviders: ['webgpu'],
  externalData: [
    { path: 'transformer_q4.onnx_data',   data: base + 'transformer/transformer_q4.onnx_data' },
    { path: 'transformer_q4.onnx_data_1', data: base + 'transformer/transformer_q4.onnx_data_1' },
  ],
});

Each path must exactly match the location stored in the graph, which is the bare file name in the same folder. The weights are about 5.6 GB in total, and wasm32 has a 4 GB heap, so create the sessions one at a time: run the text encoder, release() it, then load the transformer, and so on. Alternatively, run the text encoder with the wasm EP in a separate worker. You need a WebGPU adapter with ≥ 4 GB maxBufferSize/memory; desktop Chrome/Edge works best.

WebGPU op notes

• Transformer ops: Add, Cast, Concat, Cos, Div, Expand, Gather, Gemm, LayerNormalization, MatMul, MatMulNBits, Mul, Neg, Pow, Reciprocal, ReduceMean, Reshape, Shape, Sigmoid, Sin, Slice, Softmax, Split, Sqrt, Squeeze, Transpose, Unsqueeze. The WebGPU EP supports all of them.
• Text encoder adds GatherND, Range, Where, LessOrEqual, And, Max (causal/padding mask construction on int64/bool tensors). These are small shape/mask ops, and ORT may place some int64 ones on the CPU, which is expected and cheap. If a WebGPU build rejects GatherND or And, a fallback is to run the text encoder with the wasm EP; it is called once per prompt.
• VAE: Conv, InstanceNormalization (GroupNorm decomposition), Resize (nearest), Softmax, MatMul, Clip. All are supported. At 1024², the mid-block attention is 16384×16384 per head, about 1 GB in fp16, which may exceed some GPUs' buffer limits. Decode at 512² or tile if needed.
• There are no fp64 tensors and no custom ops besides com.microsoft.MatMulNBits.

Validation (CPU, onnxruntime 1.31)

See validation.md for the full numbers. A 512×512, 4-step, seed-42 run of the ONNX pipeline (q4 TE + q4 transformer + fp16 VAE) gives samples/onnx_q4_512.png. The PyTorch bf16 reference is samples/ref_torch_512.png.

Licence

Apache-2.0. See LICENSE and NOTICE. Original model: https://huggingface.co/black-forest-labs/FLUX.2-klein-4B