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
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).

Timings on an Apple-silicon Mac (CPU, 8 threads): DiT ~3.9 s/step, text encoder 2.6 s, VAE 1.3 s — ~19 s per 512×512 image at 4 steps.

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.

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.