litert-community/Z-Image-Turbo-LiteRT

Z-Image-Turbo — LiteRT (on-device text-to-image)

Alibaba Tongyi-MAI Z-Image-Turbo (6B, Apache-2.0) — a Single-Stream Diffusion Transformer (S3-DiT) — converted to LiteRT CompiledModel int8 graphs that generate an image fully on the phone GPU.

Prompt: "a red apple on a wooden table, studio lighting" (256 px, 8 steps) — generated end-to-end on a Pixel 8a Mali GPU (8 GB) through LiteRT. First 6B diffusion generator verified producing a real image on a commodity 8 GB phone, via chunked sequential residency. The on-device int8 output matches the fp32 reference at PSNR 40.4 dB (corr 0.9994) — visually indistinguishable.

Graphs in this repo

The 6 GB monolithic DiT exceeds LiteRT's file-load limit and a phone's GPU budget, so it is split into INTEGER-int8 graphs that load one at a time (peak footprint = a single sub-1 GB graph). Every graph below except the text encoder compiles fully on the GPU delegate; on a Galaxy S26 qwen_enc.tflite failed to compile on the GPU and on the NPU alike — see the Snapdragon NPU section:

File	Role	Size	I/O (256 px)
qwen_enc.tflite	text encoder (Qwen3-4B, penultimate hidden)	~3.5 GB	inputs_embeds[1,64,2560] → cap_feats[1,64,2560]
z_embx.tflite / z_refx.tflite	image patch embed + noise refiner	0.3 / 363 MB	img[1,256,64] → [1,256,3840]
z_embc.tflite / z_refc.tflite	caption embed + context refiner	10 / 355 MB	cap[1,32,2560] → [1,32,3840]
zc_main0..5.tflite	5 S3-DiT layers each (30 total)	866 MB ×6	hidden[1,288,3840] → [1,288,3840]
zc_final.tflite	final adaLN + projection	1.2 MB	[1,288,3840] → [1,288,64]
zvae.tflite	VAE decoder	50 MB	latent[1,16,32,32] → [1,3,256,256]

The refined image/context tokens meet as one unified [1,288,3840] hidden state passed between chunks; the composition is bit-exact to the monolithic DiT (corr 1.000000 desktop, 0.966 on-device int8/FP32-GPU). Tensors are raw float32, little-endian, row-major.

The text encoder (qwen_enc.tflite) is a standard INTEGER-int8 Qwen3-4B graph that emits the pipeline's conditioning cap_feats (the penultimate hidden state, hidden_states[-2]); tokenize the prompt with the Qwen2 BPE and embed_tokens on the host to form inputs_embeds, then slice the encoder output to the valid prompt length. The pad-token mask, x/c concat, classifier-free guidance and (un)patchify also run on the host — see the conversion scripts for the exact reference loop.

Usage (Kotlin — on-device, LiteRT CompiledModel GPU)

// One shared Environment across every graph (a null Environment leaks the OpenCL
// context and aborts after ~20 FP32 compiles).
val env = Environment.create()
fun gpu(name: String, inputs: List<FloatArray>): FloatArray {
    val opts = CompiledModel.Options(Accelerator.GPU).apply {
        // FP32 compute: the adaLN/attention path overflows fp16 to NaN.
        gpuOptions = CompiledModel.GpuOptions(precision = CompiledModel.GpuOptions.Precision.FP32)
    }
    val model = CompiledModel.create(File(dir, name).absolutePath, opts, env)
    val ins = model.createInputBuffers(); val outs = model.createOutputBuffers()
    inputs.forEachIndexed { i, a -> ins[i].writeFloat(a) }
    model.run(ins, outs)
    val out = outs[0].readFloat()
    ins.forEach { it.close() }; outs.forEach { it.close() }; model.close()
    return out
}
// Per step (Z-Image CFG): pos = DiT(cond), neg = DiT(uncond);
//   noise_pred = -(pos + guidance * (pos - neg)); latent += dsigma * noise_pred.
// chunked DiT per step: embx -> [host pad mask] -> refx ; embc -> refc ;
//   host concat -> zc_main0..5 -> zc_final -> [host unpatchify] -> VAE.

Usage (Python reference)

import numpy as np
from ai_edge_litert.compiled_model import CompiledModel

def tfl_run(path, *inputs):
    m = CompiledModel.from_file(path)
    sigs = m.get_signature_list(); key = list(sigs)[0]
    ind = m.get_input_tensor_details(key); outd = m.get_output_tensor_details(key)
    ib = m.create_input_buffers(0); ob = m.create_output_buffers(0)
    for name, buf, x in zip(sigs[key]["inputs"], ib, inputs):
        buf.write(np.ascontiguousarray(x, np.dtype(ind[name]["dtype"])))
    m.run_by_index(0, ib, ob)
    return [ob[i].read(int(np.prod(outd[n]["shape"])), np.dtype(outd[n]["dtype"]))
            .reshape(outd[n]["shape"]) for i, n in enumerate(sigs[key]["outputs"])]

# host-precompute cap_feats / RoPE / per-step adaln / sigmas / initial latent, then per
# step run the chunked DiT for cond + uncond, combine with the Z-Image CFG, Euler-update
# the latent, and decode the final latent with zvae.tflite. See the conversion scripts.

Notes

• Precision: int8 (INTEGER-compute) renders a faithful image — the on-device output matches the fp32 reference at PSNR 40.4 dB / corr 0.9994. int4 is garbage (PSNR 18).
• Two host-loop details that silently corrupt the image if missed (both cost a visible per-patch mesh): the cond and uncond prompts differ in length, so each branch needs its own context RoPE (cc/cs) and cap-pad mask — reusing the cond context for the uncond branch makes the uncond DiT wrong; and the latent must be denormalized before the VAE: latents / 0.3611 + 0.1159 (scaling_factor, shift_factor).
• GPU-delegate-only fixes (invisible to the desktop op-checker): move the pad-token MUL-after-FC to the host (bc coord for BATCH axis compile wall), force precision = FP32 (fp16 adaLN NaN), share one Environment (OpenCL context leak).
• Weights are not redistributed as the original checkpoint — the graphs are produced from the Apache-2.0 Z-Image-Turbo checkpoint with the conversion scripts.

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 5 warm-up runs then 20 timed runs, reported as the tool's mean.

This pipeline ships 13 graphs; the 5 below are the representative ones — the shards that dominate a step plus the small head and tail graphs. The repeated shards (zc_main1..5 and friends) are the same shape as the one measured here.

Graph	Graph on GPU	GPU (OpenCL)	CPU (XNNPACK, 4 threads)
z_embc.tflite	9 / 9	6 ms	2 ms
zc_main0.tflite	815 / 815	660 ms	3604 ms
zvae.tflite	641 / 641	did not run	3399 ms
zc_final.tflite	21 / 24	35 ms	14 ms
z_refc.tflite	298 / 298	100 ms	81 ms

These are the classic TFLite OpenCL delegate, not LiteRT's own accelerator. The Kotlin usage above runs through LiteRT CompiledModel, which is a different GPU implementation; any figure quoted elsewhere on this card came from that path and is not comparable to this table. Read this table as a reproducible floor anyone can re-measure with a public tool.

One graph does not run on this delegate at all. zvae fails the same way the FLUX VAE does — the OpenCL path cannot build an image of that size, and the OpenGL fallback rejects the batch shape.

The small head and tail graphs are faster on the CPU here — z_embc.tflite (2 ms on CPU against 6 ms on GPU), zc_final.tflite (14 ms on CPU against 35 ms on GPU), z_refc.tflite (81 ms on CPU against 100 ms on GPU) — so a host loop that sends every graph to the GPU is leaving time on the table.

Snapdragon NPU (Hexagon)

This repo publishes 13 graphs and the sweep measured each one separately. On the 11 where both accelerators ran, the NPU is faster on 1 and the GPU on 10. The NPU loads faster on every one of them. Per-file rows are below.

• qwen_enc.tflite — neither accelerator produced a usable row on the S26. Both ended the same way: LiteRtException: Failed to compile model.
• z_embx.tflite — the GPU runs it at 1.77 ms. The NPU does not — the request fell back to the CPU (XNNPACK) without saying so, which leaves no number attributable to the NPU.

file	backend	compiled	inference (median / min)	load
z_embc.tflite	NPU (Hexagon v81)	on-device JIT	1.09 ms / 1.01 ms	117 ms
z_embc.tflite	GPU (Adreno)	—	1.11 ms / 0.699 ms	316 ms
z_embx.tflite	GPU (Adreno)	—	1.77 ms / 1.68 ms	193 ms
z_refc.tflite	NPU (Hexagon v81)	on-device JIT	22.80 ms / 21.79 ms	723 ms
z_refc.tflite	GPU (Adreno)	—	16.97 ms / 16.62 ms	2536 ms
z_refx.tflite	NPU (Hexagon v81)	on-device JIT	127.2 ms / 121.3 ms	788 ms
z_refx.tflite	GPU (Adreno)	—	44.63 ms / 42.71 ms	3020 ms
zc_final.tflite	NPU (Hexagon v81)	on-device JIT	4.12 ms / 3.88 ms	103 ms
zc_final.tflite	GPU (Adreno)	—	0.951 ms / 0.699 ms	386 ms
zc_main0.tflite	NPU (Hexagon v81)	AOT (SM8850)	496.6 ms / 438.2 ms	2440 ms
zc_main0.tflite	GPU (Adreno)	—	136.7 ms / 134.3 ms	5847 ms
zc_main1.tflite	NPU (Hexagon v81)	AOT (SM8850)	367.4 ms / 354.7 ms	1769 ms
zc_main1.tflite	GPU (Adreno)	—	137.0 ms / 134.7 ms	5834 ms
zc_main2.tflite	NPU (Hexagon v81)	AOT (SM8850)	377.4 ms / 352.9 ms	1689 ms
zc_main2.tflite	GPU (Adreno)	—	136.6 ms / 133.4 ms	6739 ms
zc_main3.tflite	NPU (Hexagon v81)	AOT (SM8850)	374.0 ms / 356.3 ms	1696 ms
zc_main3.tflite	GPU (Adreno)	—	137.7 ms / 134.7 ms	5788 ms
zc_main4.tflite	NPU (Hexagon v81)	AOT (SM8850)	375.3 ms / 358.3 ms	1674 ms
zc_main4.tflite	GPU (Adreno)	—	137.0 ms / 134.8 ms	9221 ms
zc_main5.tflite	NPU (Hexagon v81)	AOT (SM8850)	378.7 ms / 357.4 ms	1678 ms
zc_main5.tflite	GPU (Adreno)	—	139.8 ms / 135.1 ms	6129 ms
zvae.tflite	NPU (Hexagon v81)	on-device JIT	2266.3 ms / 2008.8 ms	412 ms
zvae.tflite	GPU (Adreno)	—	232.4 ms / 218.0 ms	2498 ms

Measured on a Samsung Galaxy S26 (Snapdragon 8 Elite Gen 5 / SM8850, Hexagon v81, Android 16) with LiteRT CompiledModel 2.2.0, one accelerator per process, 5 warm-up runs then N=50 timed runs, median reported. Every run held thermal status NONE throughout. Headroom 0.54–0.84, where 1.0 is the throttling threshold.

The NPU rows marked on-device JIT ran the published file unchanged. LiteRT compiled it for the Hexagon on the device at first load. Those first compiles took 255 ms to 2.3 min here. The load column above is the cached load every later run pays. Recipe and the runtime libraries it needs: NPU guide.

The NPU rows marked AOT ran an artifact compiled ahead of time for SM8850 (ai-edge-litert 2.2.0 + QAIRT 2.47.0), not the published file. That artifact is not distributed here; the compile is one command in the NPU guide.

GPU wiring: GPU guide.

License

Apache-2.0, inherited from Z-Image-Turbo.