litert-community/Bonsai-Image-ternary-4B

Bonsai Image 4B — LiteRT (.tflite)

PrismML Bonsai Image 4B — a ternary-weight diffusion transformer built on the FLUX.2-klein-4B architecture — converted to LiteRT for on-device text-to-image. The whole pipeline runs in three fixed-shape .tflite graphs; the host does only tokenization, the FlowMatch-Euler loop, and latent unpatchify (generate.py in this repo, ~150 lines, no torch and no diffusers).

graph	file	recipe	size
DiT (Flux2Transformer2DModel, 3.88 B)	dit_int4b32.tflite	int4 block-32	2.11 GiB
DiT — GPU-shaped export (Apple GPU / Metal)	dit_gpu_int4b32.tflite	int4 block-32, same weights	2.11 GiB
text encoder (Qwen3-4B, top 9 layers pruned)	textenc_int4.tflite	int4 block-128 DRQ	1.68 GiB
text encoder — higher-fidelity variant	textenc_int8_weightonly.tflite	int8 weight-only	2.91 GiB
VAE decoder (AutoencoderKLFlux2)	vae_dec_fp32.tflite	fp32	0.19 GiB

Smallest working set: 3.97 GiB (int4 text encoder). Output is fixed at 512×512, 256 prompt tokens, 4 sampling steps by default (the model is step-distilled; more steps also work).

The ternary weights survive int4 exactly

The vendor ships the DiT with ternary weights (every block linear uses values {−scale, 0, +scale} per 128-group). In the int4 block-32 container those land as exactly {−7, 0, +7} — zero rounding decisions, verified in the converted artifact. The int4 file is a lossless container for this model's weights, at 2.11 GiB versus 14.4 GiB fp32.

Quality

Converting without quantizing anything is essentially exact (76.8 dB vs the PyTorch pipeline). With quantization, low PSNR against the PyTorch sample is sampler drift, not degradation: a small weight perturbation amplified through 4 sampling steps lands on a different draw from the same distribution. Across a 6-prompt stress set (legible text, a face, macro texture, a smooth gradient, a cluttered scene), reference-free sharpness stays flat across every recipe (Laplacian variance 1086 → 1069 at the lowest-PSNR variant), text stays legible, and faces stay correct in all columns. Pick a text encoder by size, not by fear: int4 if you want the smallest set, int8 weight-only if you want to track the reference sample more closely.

Run it

pip install ai-edge-litert numpy pillow transformers
python generate.py --model-dir . --prompt "a red fox sitting in fresh snow at sunrise" --seed 42 --out fox.png

The same sample is maintained in Google's litert-samples repository (models/bonsai/bonsai_image_4b).

Performance

All rows are for one 512×512 image at the default 4 sampling steps with a 256-token prompt. The pipeline is three graphs — text encoder once, DiT once per step, VAE decoder once — so the DiT column dominates.

Device	Backend	Text encoder	DiT (per step)	VAE	Total per image	Peak memory
Apple-silicon Mac	CPU, 8 threads	2.6 s	3.9 s	1.3 s	~19 s	—
Apple-silicon Mac	GPU (Metal, fp32) — DiT only	on CPU	0.74 s	on CPU	~6 s steady state	~22 GB
iPhone 17 Pro	CPU/XNNPACK, 6 threads	1.8 s	13 s	3.1 s	~64 s	~2.9 GiB
Pixel 8a (8 GB)	CPU	—	—	—	~7–8 min	—

Conditions worth knowing before you read those numbers:

• The Metal row uses dit_gpu_int4b32.tflite and needs fp32 GPU precision (gpu_options TOML precision = 2) — the default fp16 overflows this DiT's activation range. It also pays a one-time ~40 s Metal compile per launch, which the steady-state figure excludes, and its fp32-resident weights are what put memory near 22 GB, so a 32 GB+ Mac is the practical target. The CPU-shaped dit_int4b32.tflite does not run on the Metal accelerator at all (rank-5 rotary tensors).
• The iPhone DiT loads and prepares in ~4.5 s before the first step. Attach the XNNPACK delegate explicitly, with your thread count, when using the C API — without it the runtime falls back to reference kernels, which are orders of magnitude slower on this model.
• Device output is bit-exact against the Mac run: every DiT step matches to the last float, and the final PNG scores 51.2 dB PSNR.
• The full-app measurements land in the same place: ~62 s per image on an iPhone 17 Pro. Treat 8 GB of RAM as the floor and 12 GB+ as the practical target on Android.

On-device (iPhone 17 Pro, CPU via XNNPACK, 6 threads): the 2.11 GiB DiT loads and prepares in 4.5 s; 13 s per DiT step, text encoder 1.8 s, VAE 3.1 s — **64 s per 512×512 image** at 4 steps, ~2.9 GiB peak memory. Device output is bit-exact against the Mac run (every DiT step matches to the last float; 51.2 dB PSNR on the final PNG). One integration note: attach the XNNPACK delegate explicitly (with your thread count) when using the C API — without it the runtime falls back to reference kernels, which are orders of magnitude slower on this model.

Run it on-device (Android / iOS)

Complete single-screen apps for both platforms — including Swift/Kotlin ports of the Qwen3 tokenizer (token-exact against the Python tokenizer), the FlowMatch-Euler loop, and the latent unpatchify — are at hf-to-litertlm/bonsai_image_work/device (BonsaiApp for iOS, BonsaiAppAndroid for Android). Measured: iPhone 17 Pro ~62 s per image; Pixel 8a (8 GB RAM) completes at ~7–8 min per image, so treat 8 GB as the floor and 12 GB+ as the practical target.

Run it on the Apple GPU (macOS)

dit_gpu_int4b32.tflite is the same DiT re-exported GPU-clean — rope tables precomputed with the pipeline's constant position ids and the interleaved rotation rewritten rank-4, same weights, verified numerically identical at export. On an Apple-Silicon Mac the LiteRT Metal accelerator runs it at ~0.74 s/DiT-step, ~6 s per 512×512 image steady-state (one-time ~40 s Metal compile per launch). Force fp32 GPU precision (gpu_options TOML precision = 2) — default fp16 overflows this DiT's activation range. A complete macOS SwiftUI app (DiT on Metal, text encoder + VAE on CPU) is BonsaiAppMac in the same device directory linked above; the fp32-resident GPU weights put steady-state memory around 22 GB, so a 32 GB+ Mac is recommended. The CPU-shaped dit_int4b32.tflite does not run on the Metal accelerator (rank-5 rotary tensors), so pick the file that matches your target.

The LiteRT (Kotlin) essentials — the rest of the pipeline is plain host code:

// build.gradle: implementation("com.google.ai.edge.litert:litert:2.1.3")
import org.tensorflow.lite.Interpreter
import java.io.File

fun graph(path: String, threads: Int = 6) =
    Interpreter(File(path), Interpreter.Options().apply {   // File path => the 2.11 GiB DiT is mmap'd, not copied
        numThreads = threads
        setUseXNNPACK(true)   // required: reference kernels are both slow and wrong for blockwise int4
    })

// Stage the three graphs SEQUENTIALLY (close() each before opening the next) so
// peak memory stays ~DiT-sized (~2.9 GiB) instead of the ~4 GiB sum:
//   1. textenc:  ids (1,256) i32, mask (1,256) i32          -> embeds (1,256,7680) f32
//   2. DiT x N:  lat (1,1024,128), embeds, sigma (1,), img_ids (1024,4), txt_ids (256,4) -> velocity
//                lat += (sigma[k+1] - sigma[k]) * velocity
//   3. VAE:      z (1,32,64,64) f32                         -> image (1,3,512,512) f32, clamp(y/2+0.5)
// Map inputs by tensor NAME (serving_default_args_<n>), never by shape — two inputs share (1,256).

Conversion

Converted with [litert-torch] via the recipe published at hf-to-litertlm (export scripts, quantization, and the two conversion gotchas: Flux2's float64 RoPE frequency table must be forced to float32 to legalize tfl.pow, and blockwise zero-scales in all-zero blocks need patching for XNNPACK). The text encoder is exported as a prompt embedder — the pipeline reads hidden states from layers (9, 18, 27) only, so the top 9 of 36 layers and the LM head are pruned for free.

License

Apache-2.0, following the upstream prism-ml/bonsai-image-ternary-4B release.