CoolFace
Modelpublic

yunfengwang/z-image-turbo-mnn

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
1likes
Model Card

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

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

Or in Python:

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

FileDescriptionQuant
unet.mnn + unet.mnn.weightDiT denoiser (12.3B params)int8 weights
text_encoder.mnn + .weightQwen3-4B text encoder (outputs hidden_states[-2])int8 weights
vae_decoder.mnnFlux VAE decoder (16ch latent → RGB)int8 weights
tokenizer/Qwen2 tokenizer files-

Inputs/outputs:

  • —textencoder: `inputids (1,512) int64, attentionmask (1,1,512,512) float` → `lasthidden_state (1,512,2560)`
  • —unet: sample (1,16,H/8,W/8), timestep (1,), encoder_hidden_states (1,256,2560) → out_sample
  • —vaedecoder: `latentsample (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)

Backend / precisionper-step8-step denoiseend-to-end
Metal fp32, int8 weights~13 s~104 s~118 s
Metal fp16, int8 weights*~6.2 s~50 s~61 s
CPU int8 (12 threads)~27 s~216 s~225 s

\* 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