How tongyi-mai/z-image works
1. This repository is an inference-focused implementation of Z-Image, a single-stream diffusion transformer image generator. It is not a training repo: the code centers on loading released checkpoints and running text-to-image generation efficiently.
2. The main runtime entrypoints are inference.py and batch_inference.py. These are the user-facing scripts that wire prompts, resolution, step count, guidance, and random seed into the model stack for single-image or batched generation.
3. The core model lives under src/zimage/ and is split into three main subsystems:
- transformer.py: the DiT-style backbone that performs the denoising computation.
- scheduler.py: the step controller that drives the diffusion/flow-matching loop.
- autoencoder.py: the latent-space boundary that encodes/decodes between pixel images and model latents.
4. src/zimage/pipeline.py is the orchestration layer. It likely assembles text conditioning, latent preparation, iterative denoising, and final decoding into a single generation pipeline. Architecturally, this is the place where the model’s components become a usable end-to-end inference graph.
5. The README describes Z-Image as a “single-stream” model, and that shows up in the architecture boundary: text tokens, visual semantic tokens, and image VAE tokens are concatenated into one sequence before entering the transformer. That means the central computation is a unified sequence model rather than separate text/image streams.
6. src/utils/attention.py and the README’s attention backend notes indicate attention is a configurable performance boundary. The code supports different kernels/backends, with Flash Attention and PyTorch SDPA as important execution choices for speed and memory use. This is a key architectural lever for deployment.
7. src/utils/loader.py, import_utils.py, and helpers.py form the support layer around model construction. These utilities are the glue for optional dependencies, component loading, and general runtime conveniences needed to make the pipeline portable across environments.
8. src/config/ defines model and inference configuration separately from the core implementation. src/config/model.py and src/config/inference.py likely hold variant-specific defaults such as checkpoint names, input sizes, step counts, and guidance behavior, keeping model selection outside the execution code.
9. src/config/manifests/ plus src/tools/generate_manifest.py define a packaging boundary for checkpoints. The manifest files describe model artifacts and are likely used to make model loading reproducible across variants such as Z-Image and Z-Image-Turbo.
10. There is also a compatibility boundary for the Hugging Face diffusers ecosystem. The README shows a ZImagePipeline exposed through diffusers, so this repo is designed to support both native PyTorch inference and integration into the broader diffusers runtime.
11. The most important data flow is: prompt/configuration in, tokenization and latent preparation, iterative scheduler-driven transformer calls, then VAE decoding back to an image. In architecture-diagram terms, the central loop is “text conditioning → latent denoising → image decode,” with the transformer as the compute hotspot and the scheduler as the control plane.
12. Operationally, the repo targets checkpoint-based deployment from Hugging Face and ModelScope. Those external model hubs matter because the code assumes weights are loaded at runtime rather than baked into the repository, which makes the loading/configuration layer an important boundary in the architecture.
Main entry points: pipeline.py Pipeline · loader.py Loader · transformer.py Transformer · __init__.py Package Init
Components 16
Entry points
User-facing inference scripts and runtime adapters.
- Inference · Single-image generation entrypoint for prompt-driven inference.
inference.py - Batch Inference · Batched generation entrypoint for multiple prompts or jobs.
batch_inference.py - Diffusers API · Compatibility boundary exposing the pipeline through a Hugging Face Diffusers-style interface.
src/zimage/pipeline.py
Core Z-Image
Main inference pipeline and model components.
- Pipeline · End-to-end inference assembly for conditioning, denoising, and decoding.
src/zimage/pipeline.py - Transformer · Single-stream denoising transformer operating over concatenated text, semantic, and image tokens.
src/zimage/transformer.py - Scheduler · Controls the diffusion or flow-matching iteration schedule and update steps.
src/zimage/scheduler.py - Autoencoder · Encodes images into latents and decodes latents back to pixels.
src/zimage/autoencoder.py
Support layer
Configuration, utilities, and packaging helpers.
- Attention · Attention backend selection and performance-sensitive kernel wiring.
src/utils/attention.py - Loader · Model component loading and checkpoint assembly helpers.
src/utils/loader.py - Imports · Optional dependency handling and import compatibility helpers.
src/utils/import_utils.py - Helpers · General-purpose runtime conveniences used across inference code.
src/utils/helpers.py - Model Config · Model variant defaults and checkpoint-related settings.
src/config/model.py - Inference Config · Generation defaults such as steps, guidance, size, and seed behavior.
src/config/inference.py - Manifests · Checkpoint manifest definitions for reproducible model loading.
src/config/manifests/z-image-turbo.txt - Manifest Tool · Generates or updates checkpoint manifests.
src/tools/generate_manifest.py - Package Init · Top-level package export surface for the repository.
src/__init__.py
Connections 16
- Inference → Pipeline: runs · Single-image CLI delegates generation to the shared pipeline.
- Batch Inference → Pipeline: runs · Batch CLI reuses the same pipeline for repeated jobs.
- Diffusers API → Pipeline: wraps · Diffusers-facing interface adapts to the native pipeline implementation.
- Pipeline → Loader: loads · Pipeline constructs model components and checkpoints through loading helpers.
- Pipeline → Model Config: reads · Pipeline selects model variants and defaults from configuration.
- Pipeline → Inference Config: reads · Pipeline consumes generation parameters and inference defaults.
- Pipeline → Transformer: calls · Pipeline invokes the denoising backbone during iterative generation.
- Pipeline → Scheduler: drives · Pipeline uses the scheduler to control the denoising loop.
- Pipeline → Autoencoder: encodes/decodes · Pipeline moves between pixel space and latent space via the autoencoder.
- Transformer → Attention: uses · Transformer execution is shaped by the selected attention backend.
- Loader → Imports: depends on · Loading logic uses optional-import helpers for environment portability.
- Loader → Helpers: depends on · Loader reuses shared runtime helper functions.
- Model Config → Manifests: maps to · Model configuration references artifact manifests for checkpoint selection.
- Manifest Tool → Manifests: generates · Packaging tool produces the manifest files used for release artifacts.
- Package Init → Pipeline: exports · Package init exposes the main inference pipeline entry surface.
- Package Init → Transformer: exports · Package init makes core model components available to consumers.