--- license: apache-2.0 base_model: Tongyi-MAI/Z-Image-Turbo base_model_relation: quantized pipeline_tag: text-to-image language: - en tags: - mlx - apple-silicon - quantized - 4-bit --- # Z-Image-Turbo, 4-bit MLX A 4-bit MLX conversion of [`Tongyi-MAI/Z-Image-Turbo`](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo). **5.91 GB instead of 32.84 GB.** The model, the architecture and the outputs are Tongyi-MAI's; the only thing done here is a change of numeric format. The original publishes its 6.15 B-parameter diffusion transformer in float32. That is 24.62 GB of the 32.84 GB download for weights that no Apple-silicon runtime evaluates at float32. At 50 Mbit/s the original is about 1 h 27 m of downloading; this is about 16 minutes. On a phone that is the difference between a model you can offer and one you cannot. ## What was changed Stated plainly, as Apache-2.0 §4(b) requires. 1. **The transformer's 275 weight matrices were quantized to 4 bits**, group size 64, affine mode — `mlx.core.quantize(w, group_size=64, bits=4)`. That is 6,152,880,128 of its 6,154,908,736 parameters. 24.62 GB becomes 3.47 GB. 2. **The text encoder's 253 weight matrices were quantized identically**, including the tied vocabulary embedding. That is 4,022,272,000 of its 4,022,468,096 parameters. 8.04 GB becomes 2.26 GB. 3. **Everything not quantized was written as bfloat16.** In the transformer that is a float32 → bfloat16 conversion; in the text encoder and the VAE the originals are already bfloat16 and the bytes are unchanged. 4. **The VAE was copied byte for byte.** No quantization, no cast. 5. **Nothing was dropped, renamed, retrained, fine-tuned, merged or pruned.** Every tensor in the original has a counterpart here, under its original name. The configs, the scheduler and the tokenizer are the original files, with one addition noted below. Nothing else. This is not a different model and it is not tuned; it is the same weights at a smaller numeric width. ### Why bfloat16 and never float16 Measured on these weights on an Apple M2, not assumed: - The transformer's first forward returns **65536 of 65536 non-finite values at float16** and a clean `[-6.75, 6.97]` at bfloat16. The caption state it is conditioned on reaches 1.31e4, and the squares an RMSNorm takes of that are 1.7e8 against float16's 65504 ceiling. The float16 run decodes to a black frame. - The text encoder is about 4x more accurate in bfloat16 — 0.0043 against 0.0176 relative L2 versus transformers on the same 73-token prompt. - The VAE returns **NaN for every pixel at float16** and finite output at bfloat16. Its `vae/config.json` says `force_upcast: true`, and it is right to. So: bfloat16 for every unquantized tensor, and the scales and biases beside each quantized one are bfloat16 too. ### What was deliberately left unquantized | Left alone | Size | Reason | |---|---|---| | The whole VAE (244 tensors) | 167.7 MB | Almost entirely `Conv2d`, which is not `Quantizable` in mlx-swift at all. 0.17 GB is not worth a custom path. | | `all_final_layer.2-1.adaLN_modulation.1.weight` | 983,040 params, 1.4 MB at bf16 | It is index `.1` of a two-element `Sequential(SiLU, Linear)`. mlx-swift's `Module.update(modules:)` is handed `[.none, QuantizedLinear]` for that array, finds neither a value nor a dictionary at `values.first`, and throws `unexpectedStructure` from inside a `try!` — a crash, not a wrong image. A transformer *block*'s `adaLN_modulation` is a one-element array and quantizes normally; only the final layer has the hole. | | RMSNorm gains, all biases, `cap_pad_token` | 1,045,568 params | Rank-1 or tiny; `quantize` does not touch them and there is nothing to gain. | The text encoder's last decoder layer (`layers.35`) and its final `norm` are **kept**, though a Z-Image pipeline stops at layer 34 and never evaluates them. They are split across the published shards rather than sitting in a tail, so dropping them saves about 100 MB out of 5.91 GB and costs the ability to load this repository with every key accounted for. Not a good trade. ## Naming: this repository keeps the original tensor names There are two ways to publish an MLX quantization and they are not compatible. - **Emit the framework's post-quantization module tree.** That is what [`deepsweet/Z-Image-Turbo-6B-MLX-Q4`](https://huggingface.co/deepsweet/Z-Image-Turbo-6B-MLX-Q4) does. It is convenient for exactly one port — the one whose module names the file was written from — and unreadable to every other, including ours. - **Keep the checkpoint's own tensor names.** Any loader that already knows how to read `Tongyi-MAI/Z-Image-Turbo` keeps its key mapping and only has to learn the packing. That is what this repository does. The cost of the second choice is that diffusers has no convention for a quantized checkpoint under those names, so the layout below is an invention. It is written out in full so it can be implemented from this page alone. ### The packing convention For every quantized module at prefix `P`, the original `P.weight` of shape `[out, in]` is replaced by **three** tensors under the same prefix: | Tensor | dtype | Shape | Meaning | |---|---|---|---| | `P.weight` | `uint32` | `[out, in / 8]` | The 4-bit codes, eight per word | | `P.scales` | `bfloat16` | `[out, in / 64]` | One scale per group of 64 | | `P.biases` | `bfloat16` | `[out, in / 64]` | One zero point per group | `P.bias`, where the module has one, is untouched and stays a `bfloat16` vector. Element `j` of a word occupies bits `[4j, 4j+4)`, least significant nibble first, and the code is an unsigned integer in `0…15`. Groups run along the input axis. To reconstruct: ``` code = (weight[o, i / 8] >> (4 * (i % 8))) & 0xF w[o, i] = code * scales[o, i / 64] + biases[o, i / 64] ``` which is `mlx.core.dequantize(weight, scales, biases, group_size=64, bits=4, mode="affine")`, and in mlx-swift `dequantized(weight, scales: scales, biases: biases, groupSize: 64, bits: 4)`. That equivalence was checked by hand-decoding a tensor with the expression above and diffing it against `dequantized` — they agree to one bfloat16 ULP, which is the arithmetic and not the layout. A tensor whose name does **not** have `.scales` beside it in the same file is not quantized: read it as-is. `transformer/config.json` and `text_encoder/config.json` each carry an added ```json "quantization": { "group_size": 64, "bits": 4 } ``` so a loader can discover the format from the config rather than sniffing key names. It is the only edit made to any config file: the block is appended as text and every other byte of both files is upstream's, so no value was re-serialized and no precision was rewritten. `model_index.json`, `vae/config.json`, `scheduler/scheduler_config.json` and all four tokenizer files are byte-for-byte upstream's. ### One trap for anyone writing the loader The packed weights are `uint32`. A loader that casts every incoming tensor to the model dtype — the usual `value.astype(dtype)` in a load loop — will reinterpret eight packed weights as one float and produce silent garbage. Cast floating-point tensors only. ## Files ``` model_index.json LICENSE scheduler/scheduler_config.json transformer/config.json (+ quantization block) transformer/diffusion_pytorch_model-0000{1,2,3}-of-00003.safetensors 3.465 GB transformer/diffusion_pytorch_model.safetensors.index.json text_encoder/config.json (+ quantization block) text_encoder/generation_config.json text_encoder/model-0000{1,2,3}-of-00003.safetensors 2.263 GB text_encoder/model.safetensors.index.json vae/config.json vae/diffusion_pytorch_model.safetensors 0.168 GB tokenizer/{tokenizer.json,tokenizer_config.json,vocab.json,merges.txt} 0.016 GB ``` | Component | Original | Here | |---|---|---| | Transformer (6.15 B) | 24.620 GB, float32 | **3.465 GB** | | Text encoder, Qwen3-4B (4.02 B) | 8.045 GB, bfloat16 | **2.263 GB** | | VAE | 0.168 GB, bfloat16 | **0.168 GB** (copied) | | Tokenizer + configs | 0.012 GB | 0.016 GB | | **Total** | **32.844 GB** | **5.912 GB** | Each shard keeps the membership of the shard it came from, so the shard a tensor lives in is the same one it lived in upstream. Both `*.safetensors.index.json` files were regenerated: `weight_map` covers all 1071 transformer and 904 text encoder tensors, and `metadata.total_size` is the exact summed payload (3,465,052,288 and 2,262,920,192 bytes). ## Verification Everything below was run on an Apple M2 (16 GB) against this directory. **It loads with every key checked.** The converted weights were loaded through the reference implementation's own loader with all three checks on — no unused keys, every model key set, no shape mismatch. The check is not vacuous: removing the `quantization` block from `transformer/config.json`, so the loader builds a dense module tree, makes it fail immediately with `mismatchedSize(path: ["x_embedder", "weight"], expected [3840, 64], actual [3840, 8])`. **It produces a bit-identical image to the original.** Same prompt, seed 42, 512 px, 8 steps, flow-match Euler, guidance 0: | | PNG MD5 | pixels | |---|---|---| | `Tongyi-MAI/Z-Image-Turbo` (32.84 GB), quantized to 4-bit at load | `8366dc53f2958d0ef946efa3f07a2f69` | mean 132.09, sd 56.14, 0.08% clipped | | This repository | `8366dc53f2958d0ef946efa3f07a2f69` | mean 132.09, sd 56.14, 0.08% clipped | Max absolute per-channel difference over all 786,432 float32 values: **0.0**. This is expected rather than lucky — the runtime path is float32 → bfloat16 → `quantize`, and this conversion is the same three steps done ahead of time — but it is the claim worth checking, and it was checked rather than assumed. **Per-tensor quantization error**, dequantized and measured against the float32 original: | Tensor | Relative L2 | SNR | |---|---|---| | `layers.10.feed_forward.w1.weight` `[10240, 3840]` | 0.0947 | 20.5 dB | | `cap_embedder.1.weight` `[3840, 2560]` | 0.0911 | 20.8 dB | | `all_x_embedder.2-1.weight` `[3840, 64]` | 0.0762 | 22.4 dB | | `model.embed_tokens.weight` `[151936, 2560]` | 0.0911 | 20.8 dB | | `model.layers.20.self_attn.q_proj.weight` `[4096, 2560]` | 0.0928 | 20.7 dB | Consistent across components and shapes, which is what argues against a packing error — a wrong nibble order or a mismatched group size shows up as one tensor far off the others, not as a uniform 9%. **Runtime**, 512 px, 8 steps, staged so the encoder is released before the transformer loads: 20 s to load and forward once, about 9.0 s per denoising step, 2.4 GB peak resident. ### What 4 bits costs, honestly This conversion adds nothing to what the runtime already did, so the real question is what 4-bit costs against a wider format. The same prompt and seed at 8-bit (group 32) produces a **different but equally sharp and photorealistic sample** — the fox is framed differently, the fur and whisker detail is comparable, neither is degraded. Per-pixel that reads as PSNR 20.6 dB and correlation 0.91 against the 4-bit image, but those numbers are measuring two different samples of the same prompt, not fidelity loss. Sampling is chaotic: any change to the velocity field moves the trajectory. **Not measured:** an unquantized bfloat16 reference. The transformer alone is 12.3 GB at bfloat16 and the attempt drove this 16 GB machine to 10.2 GB of swap without completing a single step. So there is no measurement here of 4-bit against full precision — only against 8-bit, and only on one prompt and one seed. If you need that comparison, run it on a machine with the memory for it. ## Attribution Derivative of [`Tongyi-MAI/Z-Image-Turbo`](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo), Apache License 2.0. The full licence text ships as `LICENSE` in this repository. The model card, the model itself and every architectural decision in it are Tongyi-MAI's; see the original repository, the [project site](https://tongyi-mai.github.io/Z-Image-blog/) and the [technical report](https://arxiv.org/abs/2511.22699). The text encoder is Qwen3-4B, also Apache-2.0. The tokenizer files are the originals, unmodified. Everything Tongyi-MAI states about the model applies unchanged: it is an 8-step distilled model, guidance should be 0, and the recommended resolution is 1024. Quantized with mlx-swift's `MLX.quantized`. No content, capability or safety property of the original was evaluated here.