kimi000/ZImage-4Step-3000

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

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:

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:

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