Add FLUX.2-klein-4B text-to-image and image-editing sample (chunked DiT on CompiledModel GPU)

[john-rocky]

FLUX.2 [klein] 4B (Black Forest Labs, Apache-2.0) generating images end to end on the LiteRT CompiledModel GPU delegate. Nothing runs on the CPU: the Qwen3-4B text encoder, the 4B rectified-flow transformer and the VAE decoder are all Accelerator.GPU graphs.

The upstream model card describes klein as running "on consumer GPUs, with as little as 13 GB VRAM". This sample runs the same weights on a Pixel 8a's Mali-G610.

Device	Pixel 8a, 306 s (encoder 35 s, ~70 s/step, VAE 7 s)
Quality	PSNR 36.8 dB / corr 0.9987 vs the fp32 diffusers pipeline
Graphs	12 x INTEGER-int8, 6.2 GB total, largest 912 MB
Weights	https://huggingface.co/litert-community/FLUX.2-klein-4B-LiteRT

What it demonstrates

Sequential residency. CompiledModel cannot load a graph over 2 GB, and the GPU budget is well under that anyway. The model is split into chunks resident one at a time, so peak footprint is a single ~912 MB graph. Chunks of roughly 500 MB - 1 GB are the sweet spot; a 1.8 GB chunk sends the shader compiler into a pathological blowup.

A distilled sampling loop. klein reports is_distilled, so the pipeline runs no classifier-free guidance: one transformer pass per step rather than two, and a plain flow-matching Euler update.

Host/GPU split. Tokenization, embed_tokens, the mask, both rotary tables, the scheduler and the two tail permutations are precomputed by conversion/gen_prep_klein.py. The tail (unpack by position id, batch-norm denorm, 2x2 unpatchify) is two pure permutations around one elementwise op, so the permutations ship as int32 gather maps recovered with an arange probe through the stock pipeline functions.

Two GPU-delegate constraints worth flagging

Neither is visible to a desktop op check: the graphs compile, report zero unsupported operations, and return wrong numbers.

1. BROADCAST_TO is rejected outright, and Tensor.expand lowers to it. Grouped-query attention's repeat_kv hits this and the 4D tensor limit in one expression, so it is rewritten as a CONCATENATION of per-KV-head slices — same head order, exact.

2. A broadcast ADD whose left operand is a BATCH_MATMUL result is silently miscomputed. That is softmax(q @ kᵀ * scale + mask[1,1,S,S]), i.e. every masked attention. There is no error and no NaN: the probabilities still sum to 1 and still honour the causal and padding masks, but the logits are wrong and the image comes out as structured garbage. The fix here is to hand the mask in already expanded to [1, num_heads, S, S] — zero extra ops, same formula.

   The diagnostic signature is that token 0 is bit-exact and every later token is wrong, since token 0 attends only to key 0. A broken RoPE looks identical, because RoPE is the identity at position 0. conversion/probe_enc_layer.py exports one encoder layer with a tap at every stage, which is how this was localized; materializing an intermediate as a graph output made the bug disappear, which is what identified the fused pattern.

I'm happy to file the second one as a separate delegate issue if that's useful — it affects any encoder-style transformer with a [1,1,S,S] mask.

Verification

conversion/gen_verify_klein.py runs the exact device loop on the host, driven only by the twelve exported .tflite graphs, and reports PSNR against the fp32 reference. The Kotlin is a transcription of that loop.

Note that the desktop int8 path is a pessimistic proxy: the same graphs score 36.4 dB through the host CPU int8 kernels and 44.1 dB on the GPU delegate.

Directory

Placed at compiled_model_api/text_to_image/flux2_klein_kotlin_gpu/ so it can sit alongside another text-to-image sample. Happy to rename if you'd prefer a different convention.