Publish ZImage step-3000 EMA as merged Diffusers pipeline with four-step demo

a7509ea verified 1 day ago

5 kB

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

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