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 each compile fully on the GPU delegate and load one at a time (peak footprint = a single sub-1 GB graph):

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.

License: Apache-2.0 (inherited from Z-Image-Turbo).