You can deploy this model and start earning money today!

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

Name: Bonsai-Image-ternary-4B
Availability: InStock

Views

No views yet

LiteRT is Google's on-device runtime, the new name for TensorFlow Lite (Android: com.google.ai.edge.litert:litert), and litert-torch, the renamed ai-edge-torch, is its PyTorch converter: a PyTorch model converted unmodified with litert_torch.convert matched the original to 4e-7 on a Galaxy S26 (measured, LiteRT 2.2.0, Android 16, 2026-09-05).

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

bash

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.

Snapdragon NPU (Hexagon)

dit_gpu_int4b32.tflite — neither accelerator produced a usable row on the S26. NPU: the ahead-of-time compile for SM8850 failed on the host. GPU: the benchmark process was killed, most likely out of memory.
dit_int4b32.tflite — neither accelerator produced a usable row on the S26. NPU: the ahead-of-time compile for SM8850 failed on the host. GPU: the benchmark process was killed, most likely out of memory.
textenc_int4.tflite — neither accelerator produced a usable row on the S26. NPU: the ahead-of-time compile for SM8850 failed on the host. GPU: the benchmark process was killed, most likely out of memory.
textenc_int8_weightonly.tflite — neither accelerator produced a usable row on the S26. NPU: the ahead-of-time compile for SM8850 failed on the host. GPU: the benchmark process was killed, most likely out of memory.
vae_dec_fp32.tflite — neither accelerator produced a usable row on the S26. NPU: the graph compiles and then fails to run (LiteRtException: Failed to invoke the compiled model). GPU: LiteRtException: Failed to compile model.

Target: Samsung Galaxy S26 (Snapdragon 8 Elite Gen 5 / SM8850, Hexagon v81, Android 16).

GPU wiring: GPU guide. NPU recipe: NPU guide.

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:

kotlin

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

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