feat: Add Z-Image and Z-Image-Turbo support to MaxDiffusion#446
Open
csgoogle wants to merge 1 commit into
Open
feat: Add Z-Image and Z-Image-Turbo support to MaxDiffusion#446csgoogle wants to merge 1 commit into
csgoogle wants to merge 1 commit into
Conversation
csgoogle
marked this pull request as ready for review
July 22, 2026 05:35
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
csgoogle
force-pushed
the
sagarchapara/zimage
branch
9 times, most recently
from
July 23, 2026 07:51
f56b6a1 to
b6c4e39
Compare
csgoogle
force-pushed
the
sagarchapara/zimage
branch
from
July 23, 2026 11:39
b6c4e39 to
c176505
Compare
…harding specs - Pre-split text_encoder and transformer during initialization to avoid host-side splitting CPU overhead. - Fully tensorize ZImageTransformer2DModel __call__ forward path using statically-padded 3D stacked sequences, bypassing cross-device all-reduce/all-gather collectives and loops. - Pass 4D transposed layouts directly to NNXAttentionOp to bypass unflatten/flatten transposes. - Maintain RMSNorm outputs in float32 and pass them directly into _apply_rope to enable seamless on-device loop/kernel fusion and avoid float32 <-> bfloat16 casting loops. - Inline latent generation logic inside pipeline call, removing redundant JIT overhead. - Add logical sharding specs and registry for Z-Image transformer.
csgoogle
force-pushed
the
sagarchapara/zimage
branch
from
July 23, 2026 11:48
c176505 to
7d5e1d0
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add Z-Image and Z-Image-Turbo support
Adds the Tongyi-MAI Z-Image family to MaxDiffusion: a 6.15B flow-matching DiT, its Flux-format VAE, and a Qwen3-4B text encoder — all running in JAX on TPU, with an Orbax cache for the whole pipeline.
What's here
models/z_image/transformer_z_image.pymodels/z_image/z_image_utils.pypipelines/z_image/z_image_pipeline.pycheckpointing/z_image_checkpointer.pymodels/qwen3_utils.pyqwen3_flax.pygenerate_zimage.pygenerate_wan.pyconfigs/base_zimage{,_turbo}.ymltests/z_image/Performance (v6e-8, Z-Image-Turbo, 1024², 9 steps)
per_device_batch_size: 1.0)Default mesh is
data=8. Measured alternatives at batch 8:data=4/ctx=20.325s/img,ctx=80.519s/img,fsdp=81.179s/img.Notable fixes made along the way
NNXFlaxQwen3Attentionwas numerically broken — no causal mask, and it readattention_maskas additive when callers pass 0/1. Cosine 0.242 vs HF Transformers; now 0.999996. Its existing test never checked numerics.param_dtype(float32 default) regardless of compute dtype, so the 4B encoder occupied 16.1 GB/chip instead of 8.0. Norm scales stay fp32 deliberately, now declared in the module.get_git_commit_hash/ GCS upload /delete_file/save_imagesmoved tomax_utils, deduping ~90 lines across thegenerate_wan,generate_ltx2andgenerate_zimagedrivers.Verification
pyink --checkclean across 319 files; ruff clean on all touched files.