Z Image Inference From Scratch

A from-scratch, dependency-light inference implementation of Z-Image-Turbo — a flow-matching diffusion transformer for text-to-image generation.

---

🎯 Runs on 8 GB of VRAM

The defining feature of this project: you can generate 1024×1024 images on a single GPU with as little as 8 GB of VRAM.

Z-Image-Turbo is a large model — a 3840-dim, 30-layer 6B Diffusion Transformer, a 36-layer Qwen3 4B text encoder, and a Flux-style VAE decoder. Loaded naively, it does not fit in 8 GB. This implementation ships a custom three-state, per-layer VRAM manager that keeps weights on CPU and streams only the layer currently being computed onto the GPU, so peak VRAM stays tiny. The whole pipeline (text encoding → denoising → VAE decoding) runs end-to-end on a budget card.

No bitsandbytes. No quantization. No 4-bit tricks. Just exact bfloat16 weights, offloaded one layer at a time.

Set --vram_limit to control the ceiling. By default it auto-detects your total VRAM and subtracts 0.5 GB of headroom — so on an 8 GB card it targets ~7.5 GB and enables offloading automatically. Pass None to disable offloading entirely (for cards large enough to hold everything).

---

✨ Highlights

• From scratch. Every module — the DiT, the text encoder wrapper, the VAE decoder, the flow-match scheduler, the state-dict converters — is written by hand in plain PyTorch. No diffusers, no DiffSynth runtime dependency; the model definitions live in this repo.
• 8 GB VRAM friendly. Per-layer CPU↔GPU offloading with a three-state machine (offload → onload → preparing → computation) keeps peak memory low. See How the 8 GB trick works.
• Zero-cost model loading. Models are built on the meta device (no random init, no allocation) and weights are loaded with assign=True (pointer assignment, no copy). The first byte of VRAM is touched only when a layer is actually needed.
• Turbo speed. Default config is 8 denoising steps with cfg_scale=1.0 (no negative prompt needed) for fast, high-quality 1024×1024 outputs.
• Clean CLI + Python API. Generate from the command line or import ZImagePipeline directly.

---

🚀 Quick Start

Installation

pip install -r requirements.txt

Generate an image

python inference.py --prompt "a cute girl" --output output.jpg

On an 8 GB GPU this just works out of the box — VRAM management is enabled automatically.

Python API

import torch
from pipeline import ZImagePipeline

pipe = ZImagePipeline.from_pretrained(
    model_id="Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    device="cuda",
    vram_limit=7.5,          # GB; enables per-layer offloading. Use None to disable.
)

image = pipe(
    prompt="a cute girl",
    negative_prompt="",
    cfg_scale=1.0,
    height=1024,
    width=1024,
    seed=42,
    device="cuda",
    num_inference_steps=8,
)

image.save("output.jpg")

---

🧠 How the 8 GB trick works

Three techniques combine to collapse peak VRAM:

1. Meta-device construction (skip_model_initialization). Models are instantiated with all parameters placed on the meta device — no random initialization, no memory allocated. This avoids briefly allocating the full model in RAM/VRAM just to throw the random values away.

2. assign=True weight loading. Weights are loaded directly from safetensors into the model with load_state_dict(..., assign=True), which assigns tensor pointers rather than copying into pre-allocated buffers.

3. Three-state per-layer offloading (AutoWrappedLinear / AutoWrappedModule). Every nn.Linear, RMSNorm, LayerNorm, Embedding, Conv2d, and GroupNorm is wrapped in an auto-offload module with a state machine:

   • offload — weights live on CPU in bfloat16 (state 0).
   • onload — moved to the onload device when the parent model is active (state 1).
   • preparing → computation — promoted onto GPU only right before the forward pass, and only if free VRAM permits; otherwise computed in place.

   During inference the pipeline explicitly swaps models in and out of GPU: encode the prompt (text encoder on GPU, then offloaded) → denoise (DiT on GPU, then offloaded) → decode (VAE on GPU, then offloaded). Only one model is resident on the GPU at any moment, and within each model only the active layer. That is why an 8 GB card is enough.

The offloading machinery lives in utils/vram.py, and the model-swapping orchestration in pipeline.py.

---

🏗️ Project Structure

z_image_inference_from_scratch/
├── inference.py                       # CLI entry point
├── pipeline.py                        # ZImagePipeline: load + run end-to-end
├── requirements.txt
├── models/
│   ├── __init__.py
│   ├── dit.py                         # ZImageDiT diffusion transformer + model_fn_z_image_turbo
│   ├── modules.py                     # RMSNorm, attention, transformer blocks, timestep embedder
│   ├── text_encoder.py                # Qwen3-based text encoder + encode_prompt
│   ├── vae.py                         # FluxVAEDecoder
│   └── scheduler.py                   # FlowMatchScheduler (shift=3.0)
└── utils/
    ├── __init__.py
    ├── loading.py                     # safetensors → CPU loader + meta-device init helper
    ├── state_dict_converters.py       # DiffSynth/diffusers → from-scratch key remapping
    └── vram.py                        # AutoWrapped* + enable_vram_management (the 8 GB engine)

---

🔧 Technical Details

State-dict converters in utils/state_dict_converters.py remap the official DiffSynth/diffusers checkpoint keys into this repo's hand-written module names (strip the model.diffusion_model. DiT prefix, drop the text-encoder lm_head, and rebuild the VAE decoder block layout).

---

🙏 Acknowledgements

This project would not exist without two open-source efforts, which were studied closely and used as references throughout the implementation:

• DiffSynth-Studio — the three-state VRAM management design (offload / onload / preparing / computation) and the assign=True + meta-device loading pattern are directly inspired by DiffSynth's approach. The AutoWrapped* machinery in utils/vram.py follows the same philosophy that makes DiffSynth so memory-efficient.

• Z-Image by Tongyi-MAI — the official Z-Image implementation and the original Z-Image-test.py were the reference for the pipeline orchestration, the flow-match schedule, the DiT forward pass (model_fn_z_image_turbo), and the bfloat16 weight-loading details (including the subtle point of preserving the Qwen3 RoPE inv_freq buffer in float32 rather than downcasting it).

All model weights are downloaded from the official Tongyi-MAI/Z-Image-Turbo release on Hugging Face. Full credit for the model itself goes to the Z-Image team at Tongyi-MAI.

---

📄 License

This is an independent, from-scratch reimplementation provided for research and educational purposes. The Z-Image-Turbo model weights remain under their own license — please respect the terms set by Tongyi-MAI/Z-Image-Turbo when downloading and using them.