Skip to content

Commit 21ea345

Browse files
csgooglecsgoogle
committed
Implement Z-Image and Z-Image-Turbo serving pipeline and NNX transformer
1 parent bda41e1 commit 21ea345

23 files changed

Lines changed: 2975 additions & 226 deletions

‎src/maxdiffusion/checkpointing/checkpointing_utils.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
STABLE_DIFFUSION_XL_CHECKPOINT = "STABLE_DIFUSSION_XL_CHECKPOINT"
3737
FLUX_CHECKPOINT = "FLUX_CHECKPOINT"
3838
WAN_CHECKPOINT = "WAN_CHECKPOINT"
39+
Z_IMAGE_CHECKPOINT = "Z_IMAGE_CHECKPOINT"
3940

4041

4142
def create_orbax_checkpoint_manager(
@@ -62,6 +63,16 @@ def create_orbax_checkpoint_manager(
6263
item_handlers = None
6364
if checkpoint_type == FLUX_CHECKPOINT:
6465
item_names = ("flux_state", "flux_config", "vae_state", "vae_config", "scheduler", "scheduler_config")
66+
elif checkpoint_type == Z_IMAGE_CHECKPOINT:
67+
# Only `transformer_state` is trainable; the VAE and the Qwen3 text encoder
68+
# are frozen and kept as separate non-trainable items.
69+
item_names = ("transformer_state", "vae_state", "text_encoder_state", "z_image_config")
70+
item_handlers = {
71+
"z_image_config": ocp.JsonCheckpointHandler(),
72+
"transformer_state": ocp.StandardCheckpointHandler(),
73+
"vae_state": ocp.StandardCheckpointHandler(),
74+
"text_encoder_state": ocp.StandardCheckpointHandler(),
75+
}
6576
elif checkpoint_type == WAN_CHECKPOINT:
6677
item_names = ("low_noise_transformer_state", "high_noise_transformer_state", "wan_state", "wan_config")
6778
item_handlers = {

‎src/maxdiffusion/checkpointing/z_image_checkpointer.py‎

Lines changed: 421 additions & 0 deletions
Large diffs are not rendered by default.

‎src/maxdiffusion/common_types.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
WAN2_2 = "wan2.2"
5454
LTX2_VIDEO = "ltx2_video"
5555
LTX2_3 = "ltx2.3"
56+
Z_IMAGE = "z_image"
5657

5758
WAN_MODEL = WAN2_1
5859

‎src/maxdiffusion/configs/base_zimage.yml‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# Official Z-Image inference defaults. Z-Image and Z-Image-Turbo share the
16+
# denoiser; only the recommended denoising-step count differs (see
17+
# base_zimage_turbo.yml).
18+
19+
run_name: 'zimage'
20+
21+
# If true save config to GCS in {base_output_directory}/{run_name}/
22+
save_config_to_gcs: False
23+
24+
pretrained_model_name_or_path: 'Tongyi-MAI/Z-Image'
25+
model_name: z_image
26+
model_type: 'T2I'
27+
28+
unet_checkpoint: ''
29+
# This will convert the weights to this dtype.
30+
# When running inference on TPUv5e, use weights_dtype: 'bfloat16'
31+
weights_dtype: 'bfloat16'
32+
# This sets the layer's dtype in the model. Ex: nn.Dense(dtype=activations_dtype)
33+
activations_dtype: 'bfloat16'
34+
35+
# Maximum sequence length for the text encoder
36+
max_sequence_length: 512
37+
# offloads text encoder after text encoding to save memory.
38+
offload_encoders: False
39+
40+
# Attention
41+
attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
42+
# Maxdiffusion has 2 types of attention sharding strategies:
43+
# 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
44+
# 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is
45+
# sharded in cross attention q.
46+
attention_sharding_uniform: True
47+
# Empty dict lets the kernel pick defaults. Example v6e override:
48+
# flash_block_sizes: {
49+
# "block_q" : 1024,
50+
# "block_kv_compute" : 1024,
51+
# "block_kv" : 1024,
52+
# "block_q_dkv" : 1024,
53+
# "block_kv_dkv" : 1024,
54+
# "block_kv_dkv_compute" : 1024,
55+
# "block_q_dq" : 1024,
56+
# "block_kv_dq" : 1024,
57+
# "use_fused_bwd_kernel": False,
58+
# }
59+
flash_block_sizes: {}
60+
61+
# Output directory
62+
# Create a GCS bucket, e.g. my-maxdiffusion-outputs and set this to "gs://my-maxdiffusion-outputs/"
63+
base_output_directory: ""
64+
output_dir: "/mnt/disks/data/tmp/maxdiffusion"
65+
# Local path for the generated image. If empty, an image named
66+
# {run_name}_output_{seed}.png is written to the working directory.
67+
output_file: "/mnt/disks/data/tmp/zimage.png"
68+
69+
# Hardware
70+
hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
71+
skip_jax_distributed_system: True
72+
73+
# Parallelism
74+
mesh_axes: ['data', 'fsdp', 'context', 'tensor']
75+
# Z-Image runs sequence parallel (activations sharded over `context`) with
76+
# data parallelism; the tensor axis is left at 1. `attention_sharding_uniform`
77+
# adds the matching q/kv sequence rules on top of these.
78+
logical_axis_rules: [
79+
['batch', ['data', 'fsdp']],
80+
['activation_batch', ['data', 'fsdp']],
81+
['activation_length', 'context'],
82+
['activation_heads', 'tensor'],
83+
['heads', 'tensor'],
84+
['mlp', 'tensor'],
85+
['embed', 'fsdp'],
86+
['norm', 'tensor'],
87+
['out_channels', 'tensor'],
88+
]
89+
data_sharding: [['data', 'fsdp', 'context', 'tensor']]
90+
91+
sharding:
92+
transformer: 'default'
93+
94+
# One axis for each parallelism type may hold a placeholder (-1)
95+
# value to auto-shard based on available slices and devices.
96+
# For v6e TPU, Data Parallelism (DP) is optimal for throughput when per_device_batch_size=1 (e.g. 0.3s/image on v6e-8).
97+
dcn_data_parallelism: 1
98+
dcn_fsdp_parallelism: 1
99+
dcn_context_parallelism: 1
100+
dcn_tensor_parallelism: 1
101+
ici_data_parallelism: -1 # recommended ICI axis to be auto-sharded
102+
ici_fsdp_parallelism: 1
103+
ici_context_parallelism: 1
104+
ici_tensor_parallelism: 1
105+
106+
allow_split_physical_axes: False
107+
108+
# Generation parameters
109+
prompt: "A red fox reading a book in a quiet library, cinematic light"
110+
height: 1024
111+
width: 1024
112+
num_inference_steps: 50
113+
# The released Z-Image configuration runs without classifier free guidance.
114+
guidance_scale: 0.0
115+
seed: 42
116+
# Images per device. Fractional values are allowed: on 8 devices,
117+
# 0.125 generates a single image.
118+
per_device_batch_size: 1.0
119+
# Latents decoded per VAE call. The decoder is replicated and its
120+
# activations are full-resolution, so raising this trades memory for speed.
121+
vae_decode_chunk: 1
122+
123+
# Profiling
124+
# generate_zimage runs an un-profiled warmup pass (compile), then a clean
125+
# generation pass, and only then a profiled pass of profiler_steps steps.
126+
enable_profiler: False
127+
profiler_steps: 5
128+
129+
# ML Diagnostics settings
130+
enable_ml_diagnostics: False
131+
profiler_gcs_path: ""
132+
enable_ondemand_xprof: False
133+
134+
# Directory for the Orbax cache of the converted pipeline (transformer + VAE +
135+
# text encoder). '' disables it. The first run loads from Diffusers and writes
136+
# the cache; later runs restore from it and skip the torch -> flax conversion.
137+
pretrained_orbax_dir: ""
138+
139+
# Keys below are required by pyconfig but unused for Z-Image inference.
140+
jax_cache_dir: ""
141+
dataset_name: ""
142+
dataset_save_location: "/mnt/disks/data/tmp"
143+
learning_rate_schedule_steps: 1
144+
max_train_steps: 1
145+
compile_topology_num_slices: -1
146+
quantization_local_shard_count: -1
147+
use_qwix_quantization: False

‎src/maxdiffusion/configs/base_zimage_turbo.yml‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# Official Z-Image-Turbo inference defaults. The base model uses the same
16+
# transformer; only its recommended denoising-step count differs.
17+
18+
run_name: 'zimage-turbo'
19+
20+
# If true save config to GCS in {base_output_directory}/{run_name}/
21+
save_config_to_gcs: False
22+
23+
pretrained_model_name_or_path: 'Tongyi-MAI/Z-Image-Turbo'
24+
model_name: z_image
25+
model_type: 'T2I'
26+
27+
unet_checkpoint: ''
28+
# This will convert the weights to this dtype.
29+
# When running inference on TPUv5e, use weights_dtype: 'bfloat16'
30+
weights_dtype: 'bfloat16'
31+
# This sets the layer's dtype in the model. Ex: nn.Dense(dtype=activations_dtype)
32+
activations_dtype: 'bfloat16'
33+
34+
# Maximum sequence length for the text encoder
35+
max_sequence_length: 512
36+
# offloads text encoder after text encoding to save memory.
37+
offload_encoders: False
38+
39+
# Attention
40+
attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
41+
# Maxdiffusion has 2 types of attention sharding strategies:
42+
# 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
43+
# 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is
44+
# sharded in cross attention q.
45+
attention_sharding_uniform: True
46+
# Empty dict lets the kernel pick defaults. Example v6e override:
47+
# flash_block_sizes: {
48+
# "block_q" : 1024,
49+
# "block_kv_compute" : 1024,
50+
# "block_kv" : 1024,
51+
# "block_q_dkv" : 1024,
52+
# "block_kv_dkv" : 1024,
53+
# "block_kv_dkv_compute" : 1024,
54+
# "block_q_dq" : 1024,
55+
# "block_kv_dq" : 1024,
56+
# "use_fused_bwd_kernel": False,
57+
# }
58+
flash_block_sizes: {}
59+
60+
# Output directory
61+
# Create a GCS bucket, e.g. my-maxdiffusion-outputs and set this to "gs://my-maxdiffusion-outputs/"
62+
base_output_directory: ""
63+
output_dir: "/mnt/disks/data/tmp/maxdiffusion"
64+
# Local path for the generated image. If empty, an image named
65+
# {run_name}_output_{seed}.png is written to the working directory.
66+
output_file: "/mnt/disks/data/tmp/zimage-turbo.png"
67+
68+
# Hardware
69+
hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
70+
skip_jax_distributed_system: True
71+
72+
# Parallelism
73+
mesh_axes: ['data', 'fsdp', 'context', 'tensor']
74+
# Z-Image runs sequence parallel (activations sharded over `context`) with
75+
# data parallelism; the tensor axis is left at 1. `attention_sharding_uniform`
76+
# adds the matching q/kv sequence rules on top of these.
77+
logical_axis_rules: [
78+
['batch', ['data', 'fsdp']],
79+
['activation_batch', ['data', 'fsdp']],
80+
['activation_length', 'context'],
81+
['activation_heads', 'tensor'],
82+
['heads', 'tensor'],
83+
['mlp', 'tensor'],
84+
['embed', 'fsdp'],
85+
['norm', 'tensor'],
86+
['out_channels', 'tensor'],
87+
]
88+
data_sharding: [['data', 'fsdp', 'context', 'tensor']]
89+
90+
sharding:
91+
transformer: 'default'
92+
93+
# One axis for each parallelism type may hold a placeholder (-1)
94+
# value to auto-shard based on available slices and devices.
95+
# For v6e TPU, Data Parallelism (DP) is optimal for throughput when per_device_batch_size=1 (e.g. 0.3s/image on v6e-8).
96+
dcn_data_parallelism: 1
97+
dcn_fsdp_parallelism: 1
98+
dcn_context_parallelism: 1
99+
dcn_tensor_parallelism: 1
100+
ici_data_parallelism: -1 # recommended ICI axis to be auto-sharded
101+
ici_fsdp_parallelism: 1
102+
ici_context_parallelism: 1
103+
ici_tensor_parallelism: 1
104+
105+
allow_split_physical_axes: False
106+
107+
# Generation parameters
108+
prompt: "A red fox reading a book in a quiet library, cinematic light"
109+
height: 1024
110+
width: 1024
111+
num_inference_steps: 9
112+
# The published Turbo configuration uses guidance_scale=0.
113+
guidance_scale: 0.0
114+
seed: 42
115+
# Images per device. Fractional values are allowed: on 8 devices,
116+
# 0.125 generates a single image.
117+
per_device_batch_size: 1.0
118+
# Latents decoded per VAE call. The decoder is replicated and its
119+
# activations are full-resolution, so raising this trades memory for speed.
120+
vae_decode_chunk: 1
121+
122+
# Profiling
123+
# generate_zimage runs an un-profiled warmup pass (compile), then a clean
124+
# generation pass, and only then a profiled pass of profiler_steps steps.
125+
enable_profiler: False
126+
profiler_steps: 5
127+
128+
# ML Diagnostics settings
129+
enable_ml_diagnostics: False
130+
profiler_gcs_path: ""
131+
enable_ondemand_xprof: False
132+
133+
# Directory for the Orbax cache of the converted pipeline (transformer + VAE +
134+
# text encoder). '' disables it. The first run loads from Diffusers and writes
135+
# the cache; later runs restore from it and skip the torch -> flax conversion.
136+
pretrained_orbax_dir: ""
137+
138+
# Keys below are required by pyconfig but unused for Z-Image inference.
139+
jax_cache_dir: ""
140+
dataset_name: ""
141+
dataset_save_location: "/mnt/disks/data/tmp"
142+
learning_rate_schedule_steps: 1
143+
max_train_steps: 1
144+
compile_topology_num_slices: -1
145+
quantization_local_shard_count: -1
146+
use_qwix_quantization: False

‎src/maxdiffusion/generate_flux2klein.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,8 @@
3737

3838
from maxdiffusion.models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel
3939
from maxdiffusion.models.vae_flax import FlaxAutoencoderKL
40-
from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model, load_and_convert_qwen3_weights
40+
from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model
41+
from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights
4142
from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler
4243

4344

0 commit comments

Comments
 (0)