# tongyi-mai/z-image architecture

> Architecture diagram of the GitHub repository tongyi-mai/z-image, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-16). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/tongyi-mai/z-image
- Repository: https://github.com/tongyi-mai/z-image

## Overview

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: [`src/zimage/pipeline.py`](https://github.com/tongyi-mai/z-image/blob/main/src/zimage/pipeline.py) (Pipeline), [`src/utils/loader.py`](https://github.com/tongyi-mai/z-image/blob/main/src/utils/loader.py) (Loader), [`src/zimage/transformer.py`](https://github.com/tongyi-mai/z-image/blob/main/src/zimage/transformer.py) (Transformer), [`src/__init__.py`](https://github.com/tongyi-mai/z-image/blob/main/src/__init__.py) (Package Init)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_group_entrypoints["Entry points"]
  node_node_inference_script["Inference<br/>CLI script<br/>[inference.py]"]
  node_node_batch_inference_script["Batch Inference<br/>CLI script<br/>[batch_inference.py]"]
  node_node_diffusers_bridge["Diffusers API<br/>compat layer<br/>[pipeline.py]"]
end

subgraph group_group_core["Core Z-Image"]
  node_node_pipeline["Pipeline<br/>orchestrator<br/>[pipeline.py]"]
  node_node_transformer["Transformer<br/>DiT backbone<br/>[transformer.py]"]
  node_node_scheduler["Scheduler<br/>step controller<br/>[scheduler.py]"]
  node_node_autoencoder["Autoencoder<br/>VAE boundary<br/>[autoencoder.py]"]
end

subgraph group_group_support["Support layer"]
  node_node_attention_utils["Attention<br/>backend utils<br/>[attention.py]"]
  node_node_loader_utils["Loader<br/>runtime utils<br/>[loader.py]"]
  node_node_import_utils["Imports<br/>runtime utils<br/>[import_utils.py]"]
  node_node_helpers_utils["Helpers<br/>runtime utils<br/>[helpers.py]"]
  node_node_model_config["Model Config<br/>[model.py]"]
  node_node_inference_config["Inference Config<br/>[inference.py]"]
  node_node_manifest_files["Manifests<br/>artifact metadata<br/>[z-image-turbo.txt]"]
  node_node_manifest_tool["Manifest Tool<br/>packaging tool"]
  node_node_public_init["Package Init<br/>package export<br/>[__init__.py]"]
end

node_node_inference_script -->|"runs"| node_node_pipeline
node_node_batch_inference_script -->|"runs"| node_node_pipeline
node_node_diffusers_bridge -->|"wraps"| node_node_pipeline
node_node_pipeline -->|"loads"| node_node_loader_utils
node_node_pipeline -->|"reads"| node_node_model_config
node_node_pipeline -->|"reads"| node_node_inference_config
node_node_pipeline -->|"calls"| node_node_transformer
node_node_pipeline -->|"drives"| node_node_scheduler
node_node_pipeline -->|"encodes/decodes"| node_node_autoencoder
node_node_transformer -->|"uses"| node_node_attention_utils
node_node_loader_utils -->|"depends on"| node_node_import_utils
node_node_loader_utils -->|"depends on"| node_node_helpers_utils
node_node_model_config -->|"maps to"| node_node_manifest_files
node_node_manifest_tool -->|"generates"| node_node_manifest_files
node_node_public_init -->|"exports"| node_node_pipeline
node_node_public_init -->|"exports"| node_node_transformer

click node_node_inference_script "https://github.com/tongyi-mai/z-image/blob/main/inference.py"
click node_node_batch_inference_script "https://github.com/tongyi-mai/z-image/blob/main/batch_inference.py"
click node_node_diffusers_bridge "https://github.com/tongyi-mai/z-image/blob/main/src/zimage/pipeline.py"
click node_node_pipeline "https://github.com/tongyi-mai/z-image/blob/main/src/zimage/pipeline.py"
click node_node_transformer "https://github.com/tongyi-mai/z-image/blob/main/src/zimage/transformer.py"
click node_node_scheduler "https://github.com/tongyi-mai/z-image/blob/main/src/zimage/scheduler.py"
click node_node_autoencoder "https://github.com/tongyi-mai/z-image/blob/main/src/zimage/autoencoder.py"
click node_node_attention_utils "https://github.com/tongyi-mai/z-image/blob/main/src/utils/attention.py"
click node_node_loader_utils "https://github.com/tongyi-mai/z-image/blob/main/src/utils/loader.py"
click node_node_import_utils "https://github.com/tongyi-mai/z-image/blob/main/src/utils/import_utils.py"
click node_node_helpers_utils "https://github.com/tongyi-mai/z-image/blob/main/src/utils/helpers.py"
click node_node_model_config "https://github.com/tongyi-mai/z-image/blob/main/src/config/model.py"
click node_node_inference_config "https://github.com/tongyi-mai/z-image/blob/main/src/config/inference.py"
click node_node_manifest_files "https://github.com/tongyi-mai/z-image/blob/main/src/config/manifests/z-image-turbo.txt"
click node_node_manifest_tool "https://github.com/tongyi-mai/z-image/blob/main/src/tools/generate_manifest.py"
click node_node_public_init "https://github.com/tongyi-mai/z-image/blob/main/src/__init__.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_node_inference_script,node_node_batch_inference_script,node_node_diffusers_bridge toneBlue
class node_node_pipeline,node_node_transformer,node_node_scheduler,node_node_autoencoder toneAmber
class node_node_attention_utils,node_node_loader_utils,node_node_import_utils,node_node_helpers_utils,node_node_model_config,node_node_inference_config,node_node_manifest_files,node_node_manifest_tool,node_node_public_init toneMint
```

## Components

### Entry points

User-facing inference scripts and runtime adapters.

- [Inference](https://github.com/tongyi-mai/z-image/blob/main/inference.py): CLI script · Single-image generation entrypoint for prompt-driven inference. · `inference.py`
- [Batch Inference](https://github.com/tongyi-mai/z-image/blob/main/batch_inference.py): CLI script · Batched generation entrypoint for multiple prompts or jobs. · `batch_inference.py`
- [Diffusers API](https://github.com/tongyi-mai/z-image/blob/main/src/zimage/pipeline.py): compat layer · 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](https://github.com/tongyi-mai/z-image/blob/main/src/zimage/pipeline.py): orchestrator · End-to-end inference assembly for conditioning, denoising, and decoding. · `src/zimage/pipeline.py`
- [Transformer](https://github.com/tongyi-mai/z-image/blob/main/src/zimage/transformer.py): DiT backbone · Single-stream denoising transformer operating over concatenated text, semantic, and image tokens. · `src/zimage/transformer.py`
- [Scheduler](https://github.com/tongyi-mai/z-image/blob/main/src/zimage/scheduler.py): step controller · Controls the diffusion or flow-matching iteration schedule and update steps. · `src/zimage/scheduler.py`
- [Autoencoder](https://github.com/tongyi-mai/z-image/blob/main/src/zimage/autoencoder.py): VAE boundary · Encodes images into latents and decodes latents back to pixels. · `src/zimage/autoencoder.py`

### Support layer

Configuration, utilities, and packaging helpers.

- [Attention](https://github.com/tongyi-mai/z-image/blob/main/src/utils/attention.py): backend utils · Attention backend selection and performance-sensitive kernel wiring. · `src/utils/attention.py`
- [Loader](https://github.com/tongyi-mai/z-image/blob/main/src/utils/loader.py): runtime utils · Model component loading and checkpoint assembly helpers. · `src/utils/loader.py`
- [Imports](https://github.com/tongyi-mai/z-image/blob/main/src/utils/import_utils.py): runtime utils · Optional dependency handling and import compatibility helpers. · `src/utils/import_utils.py`
- [Helpers](https://github.com/tongyi-mai/z-image/blob/main/src/utils/helpers.py): runtime utils · General-purpose runtime conveniences used across inference code. · `src/utils/helpers.py`
- [Model Config](https://github.com/tongyi-mai/z-image/blob/main/src/config/model.py): config · Model variant defaults and checkpoint-related settings. · `src/config/model.py`
- [Inference Config](https://github.com/tongyi-mai/z-image/blob/main/src/config/inference.py): config · Generation defaults such as steps, guidance, size, and seed behavior. · `src/config/inference.py`
- [Manifests](https://github.com/tongyi-mai/z-image/blob/main/src/config/manifests/z-image-turbo.txt): artifact metadata · Checkpoint manifest definitions for reproducible model loading. · `src/config/manifests/z-image-turbo.txt`
- [Manifest Tool](https://github.com/tongyi-mai/z-image/blob/main/src/tools/generate_manifest.py): packaging tool · Generates or updates checkpoint manifests. · `src/tools/generate_manifest.py`
- [Package Init](https://github.com/tongyi-mai/z-image/blob/main/src/__init__.py): package export · Top-level package export surface for the repository. · `src/__init__.py`

## Connections

- 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.