feat: Add Z-Image and Z-Image-Turbo support to MaxDiffusion#446
feat: Add Z-Image and Z-Image-Turbo support to MaxDiffusion#446csgoogle wants to merge 1 commit into
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
f56b6a1 to
b6c4e39
Compare
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.
c176505 to
7d5e1d0
Compare
|
|
||
| if self.offload_encoders: | ||
| # Restore parameters from CPU back to original device shardings | ||
| state = jax.tree_util.tree_map(lambda x, sharding: jax.device_put(x, sharding), state, self._text_encoder_shardings) |
There was a problem hiding this comment.
Can we use device_put_replicated (e.g., maxdiffusion.max_utils.device_put_replicated) instead of naive device_put tree map to ensure performant multi-host model swapping?
| def _apply_rope(x: jax.Array, freqs_cis: jax.Array, out_dtype: jnp.dtype) -> jax.Array: | ||
| """Apply complex RoPE; x is B,S,H,D and frequencies B,S,D/2.""" | ||
| x_float = x.astype(jnp.float32).reshape(*x.shape[:-1], -1, 2) | ||
| real, imag = x_float[..., 0], x_float[..., 1] | ||
| cos, sin = jnp.real(freqs_cis)[:, :, None], jnp.imag(freqs_cis)[:, :, None] | ||
| out = jnp.stack((real * cos - imag * sin, real * sin + imag * cos), axis=-1) | ||
| return out.reshape(x.shape).astype(out_dtype) |
There was a problem hiding this comment.
Similar to how we do in LTX2 and WAN, we can precompute real-valued cos and sin tensors directly and pass them as real arrays instead of relying on JAX to unroll complex tensors during execution.
| for dim, length in zip(self.axes_dims, self.axes_lens): | ||
| angular = 1.0 / (theta ** (np.arange(0, dim, 2, dtype=np.float32) / dim)) | ||
| phase = np.arange(length, dtype=np.float32)[:, None] * angular[None] | ||
| freqs.append(np.exp(1j * phase).astype(np.complex64)) |
There was a problem hiding this comment.
Instead of storing in np.complex64 and then splitting in _apply_rope, we should refactor to use real cos and sin to be performant
| self._text_encoder_split = (graphdef, offloaded_state, rest) | ||
| max_logging.log(f"Offloaded Qwen3 text encoder parameters to CPU in {(time.perf_counter() - s_offload):.4f}s.") | ||
|
|
||
| return [jnp.asarray(embeddings[index, : lengths[index]], dtype=self.dtype) for index in range(len(lengths))] |
There was a problem hiding this comment.
As we are slicing the embeddings to lengths[index], the shape of these arrays will change depending on exactly how many tokens the user's prompt has. Since these arrays get passed directly into our @jax.jit compiled denoise_step, I am worried whether it will force a full (and very slow) XLA recompilation of the denoiser each time.
To keep our serving loop blazing fast, I think we should avoid slicing here and instead pass the fixed-length padded tensors (e.g. [B, 512, D]) into the denoiser, along with an explicit boolean attention mask.
| image_lengths = [item.shape[0] for item in image_features] | ||
| caption_lengths = [item.shape[0] for item in caption_features] |
There was a problem hiding this comment.
To prevent above mentioned jax recompilation, we should tweak this to accept an explicit attention_mask from the pipeline. We can use that mask to know exactly where the valid tokens end
2 participants
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.