README.md

---
license: apache-2.0
base_model: black-forest-labs/FLUX.2-klein-4B
tags:
  - litert
  - text-to-image
  - on-device
  - diffusion-transformer
  - int8
pipeline_tag: text-to-image
library_name: litert
---

# FLUX.2-klein-4B — LiteRT (on-device text-to-image and image editing)

Black Forest Labs **[FLUX.2 [klein] 4B](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B)**
(Apache-2.0) converted to **LiteRT `CompiledModel`** int8 graphs, generating **and
editing** images fully on a phone GPU. The upstream model card says klein "runs on consumer GPUs, with
as little as 13 GB VRAM". These graphs run it on a Pixel 8a's Mali-G715, which has no
dedicated VRAM at all.

![generated on a Pixel 8a Mali GPU](pixel8a_generated.png)

Prompt: *"a red apple on a wooden table, studio lighting"*. 4 steps, 256×256, generated
end-to-end on a Pixel 8a Mali GPU in 306 s. Matches the fp32 `diffusers` pipeline at
**PSNR 36.8 dB / corr 0.9987**.

![edited on a Pixel 8a Mali GPU](pixel8a_edited.png)

Editing: source, the fp32 `diffusers` edit, and the same edit on a Pixel 8a. Prompt:
*"turn the apple into a green apple"*. 4 steps, 256×256, **PSNR 44.3 dB / SSIM 0.9998**
against the fp32 pipeline, in 328–369 s.

klein is **step-wise distilled**, so the sampling loop is unusually plain: 4 steps, no
classifier-free guidance (one DiT pass per step, not two), no sign flip — just a
flow-matching Euler update `latents += dsigma[step] * noise_pred`.

## What's here

The pipeline (Qwen3-4B text encoder → rectified-flow DiT → VAE) is exported as
**INTEGER-int8** LiteRT graphs — int×int compute, the path the GPU delegate actually
runs; weight-only-FLOAT quantization hangs the GPU compile. The 4B DiT and the 4B
encoder each exceed both LiteRT's 2 GB flatbuffer load limit and a phone's GPU budget,
so they are **split into chunks that are resident one at a time**: peak footprint is a
single ~912 MB graph rather than the 6.2 GB total.

| Graph | Role | int8 size | I/O (256 px) |
|-------|------|-----------|--------------|
| `ke_enc0` / `ke_enc1` / `ke_enc2` | Qwen3-4B layers 1-9 / 10-18 / 19-27 | 912 MB each | `[1,512,2560]` → `[1,512,2560]` |
| `kc_prep` | image + context embedders, 3 modulation FCs | 166 MB | `img[1,256,128]`, `txt[1,512,7680]`, `temb[1,3072]` → hidden + 3 modulations |
| `kc_double0` / `kc_double1` | 3 + 2 double-stream blocks | 739 / 492 MB | `img[1,256,3072]`, `txt[1,512,3072]` → same |
| `kc_single0..3` | 5 single-stream blocks each (20 total) | 615 MB each | `joint[1,768,3072]` → `[1,768,3072]` |
| `kc_final` | adaLN-continuous norm + projection | 19 MB | `[1,768,3072]` → `[1,256,128]` |
| `kv_vae` | VAE decoder | 50 MB | `latent[1,32,32,32]` → `img[1,3,256,256]` |
| `kce_prep` / `kce_double*` / `kce_single*` / `kce_final` | the same eight graphs re-exported for editing | 166 / 739 / 492 / 615×4 / 19 MB | image tokens `[1,512,128]`, joint `[1,1024,3072]` |
| `kv_vae_enc` | VAE encoder (editing only) | 35 MB | `img[1,3,256,256]` → `latent[1,32,32,32]` |

The text encoder is included because klein conditions on Qwen3-4B **hidden states from
layers 9 / 18 / 27**, interleaved to 7680 channels — not on a pooled embedding, so there
is no smaller drop-in replacement.

Tensors are raw `float32`, little-endian, row-major. Tokenization, `embed_tokens`, the
causal+padding mask, both rotary tables, the scheduler and the two tail permutations run
on the host.

## Typing the prompt

The prompt is not baked in — only two of the staged tensors depend on the words
(`inputs_embeds` and `enc_mask`), so the app encodes a typed prompt itself. Stage the
`tokenizer/` folder here (`qwen_vocab.txt`, `qwen_merges.txt`, `qwen_special.txt` and the
778 MB `qwen_embed_fp16.bin`) and the app shows editable prompt fields.

It carries a faithful `Qwen2Tokenizer` port, fixture-tested byte-for-byte against Python
(`tokenizer_fixture.txt`), and looks token rows up in the memory-mapped fp16 embedding
table — a `GATHER` over 151936 rows is not a GPU op, and the gathered row is the graph's
input anyway. The fp16 table is within 3e-8 of the fp32 weights because the checkpoint is
bf16. Tokenizing and embedding a prompt takes about 1 s; a typed "a blue ceramic teapot on
a marble counter, morning light" generates the teapot, not the baked apple.


## Image editing

`Flux2KleinPipeline.__call__` takes `image` as its **first** argument: klein is natively
an editing model and text-to-image is the `image=None` case. Editing VAE-encodes the
reference and appends its latent tokens to the noise tokens before every step,

```python
latent_model_input = torch.cat([latents, image_latents], dim=1)
latent_image_ids   = torch.cat([latent_ids, image_latent_ids], dim=1)
```

with the reference separated from the noise on the rotary `T` axis (noise `T = 0`, the
i-th reference `T = 10 + 10·i`). Afterwards only the noise half is kept:
`noise_pred[:, :latents.size(1)]`.

At 256×256 that grows the image sequence 256 → 512 tokens, so the joint sequence goes
768 → 1024. **The weights do not change**, so `kce_*` are the same tensors re-exported at
the longer shape and every chunk is byte-identical in size to its `kc_*` twin. Measured
on a Pixel 8a: peak RSS **+2 %**, per-step GPU time **1.6×**, shader compile **1.0–1.2×**
— and since compilation dominates, editing costs about **+7 %** end-to-end.

The VAE encoder needed one rewrite: `DiagonalGaussianDistribution.mode()` chunks the
64-channel moments in half, which lowers to the banned `SPLIT`. The mode *is* the mean,
so slicing the first 32 channels directly is bit-exact and GPU-clean.

**Caveats.** 256×256, one reference image, and the prompt is baked into the staged host
tensors. More references or a different resolution means a re-export, not a redesign.
Report SSIM next to PSNR: the residual error is *sparse*. On a moonlit-snow edit the
worst 0.5 % of pixels carry 76 % of the squared error (specular speckles flipping between
near-white and dark blue), which drags PSNR to 27.9 dB while SSIM stays at 0.989 and the
image is perceptually identical to the fp32 reference.

## Two things the graphs assume

Both come from the GPU delegate (ML Drift), and neither is visible to a desktop op check.

1. **The attention mask is pre-expanded across heads**: pass `[1, 32, 512, 512]`, not
   `[1, 1, 512, 512]`. A broadcast `ADD` whose left operand is a `BATCH_MATMUL` result is
   silently miscomputed — the probabilities still sum to 1 and still honour the causal and
   padding masks, but the logits are wrong.
2. **Compute must be FP32**: `GpuOptions(precision = FP32)`. The modulated (adaLN) blocks
   overflow fp16 and return NaN.

Also: create one `Environment` and share it across every `CompiledModel` (a null
environment leaks the OpenCL context), and close every `TensorBuffer` after each run.

## Usage — Python (reproduces the exact device loop)

```python
import numpy as np
from ai_edge_litert.compiled_model import CompiledModel


def run(path, *inputs):
    """Runs one chunk, then releases it — sequential residency, as on device."""
    model = CompiledModel.from_file(path)
    signatures = model.get_signature_list()
    key = list(signatures)[0]
    input_details = model.get_input_tensor_details(key)
    output_details = model.get_output_tensor_details(key)
    input_buffers = model.create_input_buffers(0)
    output_buffers = model.create_output_buffers(0)
    for name, buffer, value in zip(signatures[key]["inputs"], input_buffers, inputs):
        buffer.write(np.ascontiguousarray(value, np.dtype(input_details[name]["dtype"])))
    model.run_by_index(0, input_buffers, output_buffers)
    outputs = []
    for name, buffer in zip(signatures[key]["outputs"], output_buffers):
        detail = output_details[name]
        flat = buffer.read(int(np.prod(detail["shape"])), np.dtype(detail["dtype"]))
        outputs.append(flat.reshape(detail["shape"]).copy())
    return outputs


# Host prep (tokenizer, embed_tokens, mask, rotary tables, sigmas) omitted — see below.
hidden, taps = inputs_embeds, []
for i in range(3):
    hidden = run(f"ke_enc{i}.tflite", hidden, mask, enc_cos, enc_sin)[0]
    taps.append(hidden)
prompt_embeds = np.stack(taps, 1).transpose(0, 2, 1, 3).reshape(1, 512, 7680)

for step in range(4):
    image, text, mod_img, mod_txt, mod_single = run(
        "kc_prep.tflite", latents, prompt_embeds, temb[step:step + 1])
    for i in range(2):
        image, text = run(f"kc_double{i}.tflite", image, text, cos, sin, mod_img, mod_txt)
    joint = np.concatenate([text, image], axis=1)
    for i in range(4):
        joint = run(f"kc_single{i}.tflite", joint, cos, sin, mod_single)[0]
    latents = latents + dsigma[step] * run("kc_final.tflite", joint, temb[step:step + 1])[0]

latent = unpatchify(unpack(latents) * bn_std + bn_mean)   # two pure permutations
image = run("kv_vae.tflite", latent)[0]                   # [1,3,256,256] in [-1,1]
```

## Usage — Kotlin (Android, LiteRT GPU)

```kotlin
val environment = Environment.create()                    // create once, share

fun gpu(name: String, inputs: List<FloatArray>): List<FloatArray> {
    val options = CompiledModel.Options(Accelerator.GPU)
    options.gpuOptions = CompiledModel.GpuOptions(
        precision = CompiledModel.GpuOptions.Precision.FP32)
    val model = CompiledModel.create(File(dir, name).absolutePath, options, environment)
    val inputBuffers = model.createInputBuffers()
    val outputBuffers = model.createOutputBuffers()
    inputs.forEachIndexed { index, values -> inputBuffers[index].writeFloat(values) }
    model.run(inputBuffers, outputBuffers)
    val outputs = outputBuffers.map { it.readFloat() }
    inputBuffers.forEach { it.close() }
    outputBuffers.forEach { it.close() }
    model.close()                                          // one graph resident at a time
    return outputs
}

var hidden = inputsEmbeds
val taps = (0 until 3).map { gpu("ke_enc$it.tflite", listOf(hidden, mask, encCos, encSin))[0]
    .also { output -> hidden = output } }
// interleave the three taps -> [1, 512, 7680], then the 4-step DiT loop, then kv_vae
```

## Conversion

Quantization is `litert_torch` `full_dynamic_recipe(weight_dtype=INT8, granularity=CHANNELWISE)`.
The conversion scripts (`build_klein_enc.py`, `chunked_export_klein.py`,
`vae_deploy_klein.py`) and the host-prep / verification reference
(`gen_prep_klein.py`, `gen_verify_klein.py`) ship alongside the LiteRT sample app for
this model. Three rewrites are required for a GPU-clean graph, all exact:

- **RoPE without `GATHER_ND`** — bake the even/odd de-interleave into the rows of
  `to_q` / `to_k` and the fused `to_qkv_mlp_proj`, turning it into a contiguous
  half-split rotation. `q · k` is invariant to a permutation applied to both.
- **GQA `repeat_kv` as a `CONCATENATION`** — the stock `expand` is rank-5 *and* lowers
  to `BROADCAST_TO`, which the GPU delegate rejects outright.
- **Safe RMSNorm / LayerNorm** (max-normalized) and `ManualGroupNormND` in the VAE.

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 device. Weights are never
redistributed here — the graphs are produced from the original Apache-2.0 checkpoint
with those scripts.

## License

Apache-2.0, inherited from `black-forest-labs/FLUX.2-klein-4B` (weights and text encoder).