Metadata-Version: 2.4
Name: zimgturbo
Version: 0.1.0
Summary: Z-Image-Turbo text-to-image on Apple Silicon, accelerated with int8 custom Metal kernels
Project-URL: Homepage, https://github.com/zimgturbo/zimgturbo
Author: zimgturbo
License: Apache-2.0
License-File: LICENSE
Keywords: apple-silicon,diffusion,int8,metal,mlx,text-to-image,z-image
Classifier: Environment :: GPU
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: huggingface-hub>=0.24
Requires-Dist: jinja2>=3.1
Requires-Dist: mlx>=0.32.0
Requires-Dist: numpy>=1.26
Requires-Dist: pillow>=10.0
Requires-Dist: safetensors>=0.4
Requires-Dist: transformers>=4.44
Provides-Extra: all
Requires-Dist: diffusers>=0.36; extra == 'all'
Requires-Dist: gradio>=4.44; extra == 'all'
Requires-Dist: torch>=2.4; extra == 'all'
Provides-Extra: convert
Requires-Dist: diffusers>=0.36; extra == 'convert'
Requires-Dist: torch>=2.4; extra == 'convert'
Provides-Extra: web
Requires-Dist: gradio>=4.44; extra == 'web'
Description-Content-Type: text/markdown

# zimgturbo

**Z-Image-Turbo text-to-image on Apple Silicon, in a fraction of the time.**

A from-scratch int8 inference engine for the 6B-parameter [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo)
diffusion transformer, built on custom Metal kernels that drive the M5 tensor units
(Metal 4 `MetalPerformancePrimitives`) and are hosted inside [MLX](https://github.com/ml-explore/mlx).

> **1024×1024, the model's native 8 steps, ~13.7 s end to end on an M5 Pro (16-core GPU)**
> — versus **~76 s** for the reference PyTorch/MPS pipeline (**5.6×**) —
> with no quality regression (pixel cosine 0.986–0.997 vs. the MLX int8 baseline over 10 scenes).

![sample](assets/sample_panda.png)

---

## Install & run (one command each)

Requires an Apple Silicon Mac (M-series), macOS 26+, and [`uv`](https://docs.astral.sh/uv/).

```bash
# 1. get the weights (downloads a prepared int8 bundle, ~10 GB)
uvx zimgturbo setup            # downloads the int8 bundle (~10 GB) from HuggingFace

# 2. generate
uvx zimgturbo "a fluffy tabby cat on a windowsill, sunset backlight, photorealistic" -o cat.png
```

That's it. Subsequent images are just the second command.

<details>
<summary>Don't have a prepared bundle? Convert from the original checkpoint (no PyTorch needed)</summary>

```bash
# download the original Z-Image-Turbo (from HuggingFace) once, then:
uvx zimgturbo setup --from-original \
  --transformer   /path/to/Z-Image-Turbo/transformer \
  --vae           /path/to/vae.safetensors \
  --vae-config    /path/to/vae_config.json \
  --text-encoder  /path/to/Z-Image-Turbo/text_encoder \
  --tokenizer     /path/to/Z-Image-Turbo/tokenizer
```

The transformer is quantized to int8 on the fly using the calibration data shipped in the
package — no `torch`, no re-calibration. The resulting bundle lands in `~/.cache/zimgturbo`.
</details>

---

## Usage

### CLI

```bash
zimgturbo "a red rose covered in dew, macro"          # bare prompt = generate
zimgturbo generate "a kingfisher" -o bird.png --steps 8 --seed 42
zimgturbo bench                                        # per-stage timings
zimgturbo web  --reference /path/to/Z-Image-Turbo      # side-by-side demo (needs torch)
```

Options: `--steps` (default 8, the model's native count), `--seed`, `--height/--width`,
`-n` (batch), `--vae-dtype {bfloat16,float32}` (bf16 is default and 1.85× faster).

### Python

```python
from zimgturbo import ZImageTurbo

m = ZImageTurbo("~/.cache/zimgturbo")
r = m.generate("a single origami crane on a dark table, dramatic side light",
               steps=8, seed=1234)
print(r.seconds)      # {'encode': 0.31, 'denoise': 12.6, 'vae': 0.73}
r.save("crane.png")
```

### As a dependency

```bash
uv add zimgturbo && uv sync
```

---

## Performance

M5 Pro (16-core GPU, 48 GB), macOS 26.6, 1024×1024, 8 steps, medians over 5 runs:

| pipeline | encode | denoise | s/step | VAE | **total** | speedup |
|---|---|---|---|---|---|---|
| PyTorch/MPS bf16 (9 steps) | – | – | ~8.4 | – | **76.0 s** | 1.0× |
| MLX group-int8 (8 steps) | 0.3 | ~50 | ~6.2 | 1.7 | **~52 s** | 1.5× |
| **zimgturbo (8 steps)** | 0.31 | 12.61 | 1.575 | 0.73 | **13.7 s** | **5.6×** |

This is within **4% of the hardware roofline**: the transformer needs 62 TFLOP/step and the
measured int8 tensor-unit peak is 49 TOPS, so ~1.55 s/step is the floor; we hit 1.575.

## Quality

Same prompt and initial noise. Left → right: PyTorch bf16 (9 steps), MLX group-int8 (8),
our int8 scheme simulated in MLX (8), **our engine (8)**:

![comparison](assets/compare_cat.jpg)

Over a 10-scene set (portraits, fur, feathers, macro, English/Chinese text, architecture,
low light) the engine matches the MLX int8 baseline at pixel cosine **0.986–0.997** and is
**0.88** against the original PyTorch — marginally *closer* to PyTorch than the MLX baseline
itself (0.864). The residual is trajectory divergence from int8 rounding, not blur.

## How it works (short version)

Seven custom Metal kernels do the heavy lifting; everything else stays in MLX:

| kernel | what it does | measured |
|---|---|---|
| fused int8 GEMM | int8×int8→int32 with dequant+bias epilogue in registers | **49 TOPS** (2× fp16) |
| int8 flash attention | scores stay on-chip, vectorized softmax | **26 TOPS** |
| 5 fused row kernels | RMSNorm / modulation / SwiGLU / RoPE / transpose / quantize | 0.3–1.2 ms each |

Two things that are load-bearing and easy to get wrong:

- **GEMM layout.** With `transpose_right=true`, `matmul2d` wants the reduction dim K on the
  x axis and N on the y axis, so plain row-major `W[N][K]` works with `.slice(0, n0)`. Getting
  those axes backwards produces a read pattern that *looks* like a hardware aliasing bug and
  tempts an elaborate multi-pass workaround. There is none — the single dispatch is bit-exact.
- **Range.** The SwiGLU intermediate and the `to_out` input reach ~2e5–4e5, and the VAE
  overflows in fp16. Those are kept in fp32/bf16 and quantized straight to int8; never
  materialized as fp16.

Full write-up: [`OPTIMIZATION_GUIDE.md`](OPTIMIZATION_GUIDE.md) (a from-first-principles tour
for readers new to perf work) and the [technical report](paper/zimgturbo-report.pdf).

## Limits

- **8 steps under 5 s is not reachable** on this class of hardware: 8 steps need 496 TFLOP and
  the int8 peak is 49 TOPS, a 10.1 s hard floor. This is a compute limit, not an engineering gap.
- **int4 weights are supported but *slower*** (41 vs 49 TOPS) — this workload is compute bound,
  not weight-bandwidth bound, so shrinking weights doesn't help. int4 only saves memory.
- Numbers are specific to the **M5 Pro (16-core)** and to Z-Image-Turbo's shapes.

## For maintainers: publishing a weight bundle

Code lives on GitHub; the ~10 GB int8 bundle lives on the HuggingFace Hub:

```bash
zimgturbo publish --repo <your-hf-repo>     # uploads ~/.cache/zimgturbo
```

Then users just `zimgturbo setup`. The default bundle is [`yunfengwang/zimgturbo`](https://huggingface.co/yunfengwang/zimgturbo).

## Acknowledgements

Builds on [SmoothQuant](https://arxiv.org/abs/2211.10438),
[FlashAttention](https://arxiv.org/abs/2205.14135),
[LLM.int8()](https://arxiv.org/abs/2208.07339), the
[roofline model](https://dl.acm.org/doi/10.1145/1498765.1498785),
Apple's [MLX](https://github.com/ml-explore/mlx), and
[Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo).

## License

Apache-2.0 (code). Model weights are subject to the Z-Image-Turbo license.
