z-image-turbo-mnn

Z-Image-Turbo MNN (int8, fp16-safe)

Z-Image-Turbo (6B single-stream DiT text-to-image model) converted to MNN format for fast on-device inference on Apple Silicon (Metal) and CPU.

Converted from scratch from the original weights (ONNX export → MNNConvert), with quantization and numerical transforms for correctness and speed.

Quick start

# one-liner: install + download model + generate
uvx zimage-mnn "a cute cat sitting on a windowsill" --size 1024 --steps 8

Or in Python:

from zimage_mnn import ZImageMNNPipeline

pipe = ZImageMNNPipeline.from_pretrained("<this-repo-id>")   # downloads & caches weights
image = pipe("a cute cat", size=1024, steps=8, seed=42)
image.save("cat.png")

Model files

Inputs/outputs:

• text_encoder: input_ids (1,512) int64, attention_mask (1,1,512,512) float → last_hidden_state (1,512,2560)
• unet: sample (1,16,H/8,W/8), timestep (1,), encoder_hidden_states (1,256,2560) → out_sample
• vae_decoder: latent_sample (1,16,H/8,W/8) → sample (1,3,H,W)

Conversion details

• Exported each component to ONNX from the original diffusers weights, then MNNConvert with --weightQuantBits 8 and --transformerFuse (fuses the 34 attention blocks into MNN Attention ops).
• fp16-safe weight rescaling (SmoothQuant): the original model has activation outliers up to ~1M that overflow fp16. Since the FFN w2 / attention to_out outputs feed directly into RMSNorm (scale-invariant), all 68 such layers are rescaled by 1/64, and each FFN applies SmoothQuant (w3 /= 64, w2 *= 64, an exact compensation). This keeps fp32 output nearly identical (cosine ≈ 0.99996) while making fp16 inference NaN-free across all 8 denoising steps.
• The text encoder is exported with a precomputed 4D causal+padding attention mask to avoid ONNX ops MNN does not support (GatherND, IsNaN).
• Attention is exported in the exact pattern MNN's FuseAttention template matches, so all 34 attention blocks are fused.

Performance (Apple M5 Pro, 48GB, 1024×1024, 8 steps)

* fp16 requires the MNN build with the non-causal flash-attention / fp16 fixes; the PyPI wheel runs fp32 by default for broad correctness.

Quality: single-step transformer output vs PyTorch fp32 — cosine ≈ 0.9996 (int8) / 0.9999 (fp16 path). Final images are visually indistinguishable from the original pipeline.

Scheduler

FlowMatch Euler, shift=3.0, sigmas linspace(1, 1/N, N) shifted by sigma' = shift*sigma/(1+(shift-1)*sigma), timestep passed to the unet as (1000 - sigma*1000)/1000, unet output negated before the Euler step. VAE decode: latent/0.3611 + 0.1159.

Credits

• Original model: Tongyi-MAI/Z-Image-Turbo
• Inference engine: alibaba/MNN