kimi000/ZImage-AnyFlow2-4Step-1000

ZImage AnyFlow2 — 4 steps, checkpoint 1000

An experimental 1000-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.

Quick start

Use Python 3.10+ and a CUDA-enabled PyTorch installation. Install the dependencies in requirements.txt, then:

import torch
from diffusers import DiffusionPipeline

repo_id = "kimi000/ZImage-AnyFlow2-4Step-1000"
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-AnyFlow2-4Step-1000 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 the model. 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 1000; 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 an intermediate 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.