---
license: apache-2.0
language:
- en
library_name: diffusers
pipeline_tag: text-to-image
base_model: Tongyi-MAI/Z-Image
tags:
- z-image
- anyflow2
- distillation
- text-to-image
- custom-pipeline
- safetensors
---

# ZImage — 4 steps, checkpoint 3000

An experimental **3000-training-step EMA** checkpoint for 1024×1024 text-to-image generation in **4 model evaluations**. The student is initialized from [Tongyi-MAI/Z-Image](https://huggingface.co/Tongyi-MAI/Z-Image), then trained with single-stage AnyFlow2 independent-query interval TDM.

The LoRA has been **merged** into the BF16 transformer. This repository contains a complete Diffusers directory with sharded Safetensors, tokenizer, text encoder, VAE, scheduler, and custom pipeline/model source. No FAR training repository, PEFT adapter loading, or optimizer checkpoint is required.

## Example

![Four-step generation from the 3000-step EMA](example.png)

Prompt: `A red poster with the white words "HELLO WORLD", clean bold typography.` Four steps, seed 0, 1024×1024.

## Quick start

Use Python 3.10–3.12 and a CUDA-enabled PyTorch installation. The pinned dependencies record the versions used to validate this export. Install the dependencies in `requirements.txt`, then:

```python
import torch
from diffusers import DiffusionPipeline

repo_id = "kimi000/ZImage-4Step-3000"
pipe = DiffusionPipeline.from_pretrained(
    repo_id,
    custom_pipeline=repo_id,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    use_safetensors=True,
).to("cuda")

image = pipe(
    prompt='A cozy bookstore at dusk, warm amber light spilling through the windows. A large sign above the door clearly reads "OPEN BOOKS" in elegant white lettering. Photorealistic street photography, rich detail.',
    height=1024,
    width=1024,
    num_inference_steps=4,
    generator=torch.Generator("cpu").manual_seed(0),
).images[0]
image.save("zimage_anyflow2_4step.png")
```

The custom code implements the interval-conditioned global-basis transformer and the matching sampler; a stock `ZImagePipeline` cannot load these weights directly. `trust_remote_code=True` opts into loading the Python files in this repository. The code is supplied for inspection.

To reduce GPU memory, replace `.to("cuda")` with `pipe.enable_model_cpu_offload()` after loading. The full model also needs enough host RAM for loading. Exact peak GPU memory depends on prompt length, image size and the selected attention backend.

## Command-line demo

Download the small example and dependency file:

```bash
hf download kimi000/ZImage-4Step-3000 demo.py requirements.txt --local-dir zimage-demo
cd zimage-demo
pip install -r requirements.txt
python demo.py --prompt 'A red poster with the large white words "HELLO WORLD", clean bold typography, crisp studio lighting.' --output hello.png
# Optional memory saving:
python demo.py --cpu-offload --output example.png
```

The first run downloads approximately 21 GB. Set `HF_HUB_CACHE` to a directory on a disk with enough space if needed.  For a local snapshot, use `python demo.py --model /path/to/snapshot`. No authentication is needed for this public repository.

## Checkpoint and sampling settings

| Setting | Value |
| --- | --- |
| Training checkpoint | Exactly step 3000; EMA weights |
| Base | Tongyi-MAI/Z-Image |
| Training resolution / global batch | 1024×1024 / 128 |
| Student rollout budget | 4 |
| Continuous basis | 128 triangular bases, global coordinates |
| Fixed time shift | 6.0 |
| Final interval scale | 0.5 |
| Training objective | Velocity regression + independent-query interval TDM, weight 10 |
| Teacher CFG during training | 4.0 |
| Inference CFG | No additional CFG |
| Prompt sampling mixture | MJHQ 50%, English TextAtlas subset 25%, English reward-model prompt dataset 25% |
| LoRA training / export | Rank 256, alpha 256; merged for inference |

The four-step raw-time grid is `[1, 5/7, 3/7, 1/7, 0]`; after shift 6 it is `[1, 0.9375, 0.8181818, 0.5, 0]`. Use the supplied scheduler to preserve this grid. Larger step counts are supported by the continuous head, but this experiment trains four-step trajectories.

`export_metadata.json` records the exact source checkpoint SHA256 and export settings. `training_config.yml` records the training recipe. The frozen VAE, tokenizer and text encoder come from the base model.

## Limitations

This is a research checkpoint with no claimed benchmark score. Images and rendered text may contain errors. The base model's capabilities and limitations still apply. It was trained from English prompts; other languages have not been evaluated. This checkpoint is for text-to-image generation.

## License and attribution

Apache-2.0. The base model is [Tongyi-MAI/Z-Image](https://huggingface.co/Tongyi-MAI/Z-Image); Diffusers implements its native transformer and pipeline components. The included AnyFlow2 adaptation and inference code are distributed under the included Apache-2.0 license. This is a community research adaptation.