diff --git a/src/maxdiffusion/checkpointing/checkpointing_utils.py b/src/maxdiffusion/checkpointing/checkpointing_utils.py index 81aede97b..72414d9b8 100644 --- a/src/maxdiffusion/checkpointing/checkpointing_utils.py +++ b/src/maxdiffusion/checkpointing/checkpointing_utils.py @@ -36,6 +36,7 @@ STABLE_DIFFUSION_XL_CHECKPOINT = "STABLE_DIFUSSION_XL_CHECKPOINT" FLUX_CHECKPOINT = "FLUX_CHECKPOINT" WAN_CHECKPOINT = "WAN_CHECKPOINT" +Z_IMAGE_CHECKPOINT = "Z_IMAGE_CHECKPOINT" def create_orbax_checkpoint_manager( @@ -62,6 +63,16 @@ def create_orbax_checkpoint_manager( item_handlers = None if checkpoint_type == FLUX_CHECKPOINT: item_names = ("flux_state", "flux_config", "vae_state", "vae_config", "scheduler", "scheduler_config") + elif checkpoint_type == Z_IMAGE_CHECKPOINT: + # Only `transformer_state` is trainable; the VAE and the Qwen3 text encoder + # are frozen and kept as separate non-trainable items. + item_names = ("transformer_state", "vae_state", "text_encoder_state", "z_image_config") + item_handlers = { + "z_image_config": ocp.JsonCheckpointHandler(), + "transformer_state": ocp.StandardCheckpointHandler(), + "vae_state": ocp.StandardCheckpointHandler(), + "text_encoder_state": ocp.StandardCheckpointHandler(), + } elif checkpoint_type == WAN_CHECKPOINT: item_names = ("low_noise_transformer_state", "high_noise_transformer_state", "wan_state", "wan_config") item_handlers = { diff --git a/src/maxdiffusion/checkpointing/z_image_checkpointer.py b/src/maxdiffusion/checkpointing/z_image_checkpointer.py new file mode 100644 index 000000000..527a2be81 --- /dev/null +++ b/src/maxdiffusion/checkpointing/z_image_checkpointer.py @@ -0,0 +1,415 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Loading and Orbax caching for the Z-Image pipeline. + +All of the pipeline's loading logic lives here: `ZImagePipeline` itself is a +pure runtime object. The Orbax cache holds the whole pipeline -- the denoiser, +the VAE and the Qwen3 text encoder -- so a warm start skips the PyTorch to +Flax conversion entirely. + +Every component goes through the same three steps, whichever way it is loaded: + + 1. build it abstractly (shapes only, no weights), + 2. resolve its logical annotations against the mesh into shardings, + 3. fill it, either from safetensors or from the Orbax cache, placing each + parameter directly into its target sharding. + +`_nnx_component` does 1 and 2 for the denoiser and the text encoder; +`_linen_component` does the same for the Diffusers VAE, which is still Linen. +""" + +import inspect +import json +import os +from typing import Optional + +from etils import epath +import flax +from flax import nnx +import flax.linen as nn +from flax.linen import spmd as flax_spmd +from flax.traverse_util import flatten_dict, unflatten_dict +import jax +import jax.numpy as jnp +import orbax.checkpoint as ocp +from transformers import AutoConfig, AutoTokenizer + +from .. import max_logging, max_utils +from ..models import FlaxAutoencoderKL +from ..models.qwen3_flax import FlaxQwen3Config, NNXFlaxQwen3Model +from ..models.qwen3_utils import load_qwen3_weights +from ..models.z_image.transformer_z_image import ZImageTransformer2DModel +from ..models.z_image.z_image_utils import load_z_image_transformer +from ..pipelines.z_image import ZImagePipeline +from .checkpointing_utils import create_orbax_checkpoint_manager + + +Z_IMAGE_CHECKPOINT = "Z_IMAGE_CHECKPOINT" + +CONFIG_ITEM = "z_image_config" +# The denoiser is the only component a Z-Image workflow ever updates. The VAE +# and the text encoder are frozen, so they are saved as separate items and are +# never restored into the transformer's trainable `nnx.Param` collection. +TRAINABLE_ITEMS = ("transformer_state",) +NON_TRAINABLE_ITEMS = ("vae_state", "text_encoder_state") + +# Diffusers' Flax VAE config loader drops keys its constructor does not +# declare, so the published shift factor is carried in the cache metadata. +_DEFAULT_VAE_SHIFT_FACTOR = 0.1159 + + +def _plain(tree): + """Plain dict of plain arrays, which is all Orbax's Standard handler takes.""" + tree = flax.core.unfreeze(tree) if isinstance(tree, flax.core.FrozenDict) else tree + return jax.tree_util.tree_map( + lambda leaf: leaf.unbox() if isinstance(leaf, flax_spmd.LogicallyPartitioned) else leaf, + tree, + is_leaf=lambda leaf: isinstance(leaf, flax_spmd.LogicallyPartitioned), + ) + + +def _abstract_tree(shapes, shardings): + """Nested tree of ShapeDtypeStructs carrying each leaf's target sharding.""" + return unflatten_dict( + { + path: jax.ShapeDtypeStruct(leaf.shape, leaf.dtype, sharding=shardings[path] if shardings else None) + for path, leaf in flatten_dict(shapes).items() + } + ) + + +def _nnx_component(factory, mesh, axis_rules, seed=0): + """Build an NNX module abstractly: its graph, params and target shardings. + + Shardings come from the module's own `nnx.with_partitioning` annotations, + resolved against the mesh. + """ + graphdef, state, rest = nnx.split(nnx.eval_shape(factory, nnx.Rngs(jax.random.key(seed))), nnx.Param, ...) + shardings = None + if mesh is not None: + sharding_tree = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), mesh, axis_rules) + shardings = {path: variable.value for path, variable in nnx.to_flat_state(sharding_tree)} + return graphdef, state, rest, shardings + + +def _linen_component(module, init_args, mesh, axis_rules): + """The Linen equivalent of `_nnx_component`, for the Diffusers VAE. + + A module with no logical annotations (`vae_flax.FlaxAutoencoderKL` has none) + resolves to fully replicated. + """ + abstract = jax.eval_shape(lambda: module.init(jax.random.key(0), *init_args)) + shapes = _plain(abstract["params"]) + if mesh is None: + return shapes, None + mesh_shardings = nn.logical_to_mesh_sharding(nn.get_partition_spec(abstract), mesh, axis_rules) + return shapes, dict(flatten_dict(_plain(mesh_shardings["params"]))) + + +def _fill(graphdef, state, rest, params): + """Put loaded params into an abstract NNX state and rebuild the module.""" + flat_state = dict(nnx.to_flat_state(state)) + for path, value in flatten_dict(params).items(): + flat_state[path].value = value + return nnx.merge(graphdef, nnx.from_flat_state(flat_state), rest) + + +def _transformer_factory(transformer_config, config, mesh): + def factory(rngs): + return ZImageTransformer2DModel( + rngs=rngs, + attention_kernel=config.attention, + mesh=mesh, + flash_block_sizes=max_utils.get_flash_block_sizes(config), + dtype=config.activations_dtype, + weights_dtype=config.weights_dtype, + **transformer_config, + ) + + return factory + + +def create_z_image_transformer(model_id: str, config, mesh: Optional[jax.sharding.Mesh] = None): + """Instantiate and stream a Diffusers Z-Image checkpoint into an NNX model.""" + transformer_config = ZImageTransformer2DModel.load_config(model_id, subfolder="transformer") + graphdef, state, rest, shardings = _nnx_component( + _transformer_factory(transformer_config, config, mesh), mesh, config.logical_axis_rules, config.seed + ) + params = load_z_image_transformer(model_id, state.to_pure_dict(), target_shardings=shardings) + return _fill(graphdef, state, rest, params) + + +def _vae_constructor_kwargs(raw_config: dict) -> dict: + """Keep only the keys FlaxAutoencoderKL declares, so it can be rebuilt offline.""" + accepted = set(inspect.signature(FlaxAutoencoderKL.__init__).parameters) - {"self", "parent", "name"} + return {key: value for key, value in raw_config.items() if key in accepted and key not in ("dtype", "weights_dtype")} + + +def _text_encoder_dir(model_id: str) -> str: + """Local directory holding the text encoder's safetensors shards.""" + if os.path.isdir(model_id): + return os.path.join(model_id, "text_encoder") + from huggingface_hub import snapshot_download + + return os.path.join(snapshot_download(model_id, allow_patterns=["text_encoder/*"]), "text_encoder") + + +class ZImageCheckpointer: + """Builds a `ZImagePipeline`, using an Orbax cache when one is configured. + + `config.pretrained_orbax_dir` turns the cache on. The first run loads from + Diffusers/HuggingFace (slow: safetensors are converted to Flax on the host) + and writes the cache; later runs restore straight into the target shardings. + """ + + def __init__(self, config, mesh: Optional[jax.sharding.Mesh] = None, checkpoint_type: str = Z_IMAGE_CHECKPOINT): + self.config = config + self.mesh = mesh + self.checkpoint_type = checkpoint_type + + @property + def pretrained_orbax_dir(self) -> str: + return getattr(self.config, "pretrained_orbax_dir", "") + + def load_pipeline(self) -> ZImagePipeline: + """Restore the pipeline from the Orbax cache, else load it from Diffusers.""" + orbax_dir = self.pretrained_orbax_dir + if orbax_dir: + pipeline = self._restore_pipeline(orbax_dir) + if pipeline is not None: + return pipeline + + max_logging.log("Loading Z-Image pipeline from Diffusers.") + pipeline = self.load_diffusers_pipeline() + if orbax_dir: + self._save_pipeline(orbax_dir, pipeline) + return pipeline + + # ------------------------------------------------------- component builders + + def build_vae(self, vae_config: dict): + """VAE module, its abstract params and their target shardings.""" + vae = FlaxAutoencoderKL(**vae_config, dtype=self.config.activations_dtype, weights_dtype=self.config.weights_dtype) + sample = jnp.ones((1, vae.config.in_channels, 64, 64), jnp.float32) + return vae, *_linen_component(vae, (sample,), self.mesh, self._vae_axis_rules()) + + def build_text_encoder(self, text_encoder_config: dict): + """Qwen3 graph, its abstract params and their target shardings.""" + qwen3_config = FlaxQwen3Config( + vocab_size=text_encoder_config["vocab_size"], + hidden_size=text_encoder_config["hidden_size"], + intermediate_size=text_encoder_config["intermediate_size"], + num_hidden_layers=text_encoder_config["num_hidden_layers"], + num_attention_heads=text_encoder_config["num_attention_heads"], + num_key_value_heads=text_encoder_config["num_key_value_heads"], + head_dim=text_encoder_config["head_dim"], + rms_norm_eps=text_encoder_config["rms_norm_eps"], + rope_theta=text_encoder_config["rope_theta"], + max_position_embeddings=text_encoder_config["max_position_embeddings"], + dtype=self.config.weights_dtype, + ) + # The norms declare float32 scales, so the abstract params already carry the + # right dtypes for both the loader and the Orbax restore target. + return _nnx_component( + lambda rngs: NNXFlaxQwen3Model(rngs=rngs, config=qwen3_config), + self.mesh, + self._text_encoder_axis_rules(), + self.config.seed, + ) + + def _vae_axis_rules(self): + """`vae_logical_axis_rules` if the config declares them, as WAN does.""" + return tuple(getattr(self.config, "vae_logical_axis_rules", None) or self.config.logical_axis_rules) + + def _text_encoder_axis_rules(self): + """Qwen3 annotates its embedding table with `vocab`, which Z-Image's rules + do not name; leaving it unmapped keeps that axis replicated.""" + rules = tuple(self.config.logical_axis_rules) + if not any(rule[0] == "vocab" for rule in rules): + rules += (("vocab", None),) + return rules + + # ------------------------------------------------------------------- load + + def load_diffusers_pipeline(self) -> ZImagePipeline: + """Load every component from the Diffusers checkpoint, converting to Flax.""" + config = self.config + model_id = config.pretrained_model_name_or_path + raw_vae_config = self._raw_vae_config(model_id) + + transformer = create_z_image_transformer(model_id, config, self.mesh) + + vae, _, vae_shardings = self.build_vae(_vae_constructor_kwargs(raw_vae_config)) + # from_pretrained rebuilds the module too, but only its weights are wanted: + # the module above already carries the shardings resolved from the mesh. + _, vae_params = FlaxAutoencoderKL.from_pretrained( + model_id, subfolder="vae", from_pt=True, use_safetensors=True, dtype=config.weights_dtype + ) + vae_params = self._place(vae_params, vae_shardings) + + graphdef, state, rest, shardings = self.build_text_encoder( + AutoConfig.from_pretrained(model_id, subfolder="text_encoder").to_dict() + ) + params = load_qwen3_weights(_text_encoder_dir(model_id), state.to_pure_dict(), target_shardings=shardings) + text_encoder = _fill(graphdef, state, rest, params) + + return self._assemble( + transformer, + vae, + vae_params, + AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", use_fast=True), + text_encoder, + raw_vae_config.get("shift_factor", _DEFAULT_VAE_SHIFT_FACTOR), + ) + + def _restore_pipeline(self, orbax_dir: str) -> Optional[ZImagePipeline]: + try: + manager = self._checkpoint_manager(orbax_dir) + step = manager.latest_step() + if step is None: + max_logging.log(f"No pretrained orbax checkpoint found in {orbax_dir}") + return None + + metadata = manager.restore(step, args=ocp.args.Composite(**{CONFIG_ITEM: ocp.args.JsonRestore()}))[CONFIG_ITEM] + mismatch = self._metadata_mismatch(metadata) + if mismatch: + max_logging.log(f"Ignoring orbax checkpoint in {orbax_dir}: {mismatch}") + return None + + max_logging.log(f"Loading Z-Image pipeline from orbax checkpoint step {step} in {orbax_dir}") + graphdef, state, rest, shardings = _nnx_component( + _transformer_factory(metadata["transformer_config"], self.config, self.mesh), + self.mesh, + self.config.logical_axis_rules, + self.config.seed, + ) + vae, vae_shapes, vae_shardings = self.build_vae(metadata["vae_config"]) + text_graphdef, text_state, text_rest, text_shardings = self.build_text_encoder(metadata["text_encoder_config"]) + + # Shapes come from the models themselves, so every item lands directly in + # the sharding its component was built for -- no host staging. + restored = manager.restore( + step, + args=ocp.args.Composite( + **{ + "transformer_state": ocp.args.StandardRestore(_abstract_tree(state.to_pure_dict(), shardings)), + "vae_state": ocp.args.StandardRestore(_abstract_tree(vae_shapes, vae_shardings)), + "text_encoder_state": ocp.args.StandardRestore(_abstract_tree(text_state.to_pure_dict(), text_shardings)), + } + ), + ) + + return self._assemble( + _fill(graphdef, state, rest, restored["transformer_state"]), + vae, + restored["vae_state"], + self._load_tokenizer(orbax_dir), + _fill(text_graphdef, text_state, text_rest, restored["text_encoder_state"]), + metadata["vae_shift_factor"], + ) + except Exception as e: # pylint: disable=broad-except + max_logging.log(f"Failed to load orbax checkpoint from {orbax_dir}, falling back to Diffusers: {e}") + return None + + def _assemble(self, transformer, vae, vae_params, tokenizer, text_encoder, vae_shift_factor): + return ZImagePipeline( + transformer, + vae, + vae_params, + tokenizer, + text_encoder, + dtype=self.config.activations_dtype, + mesh=self.mesh, + logical_axis_rules=self.config.logical_axis_rules, + vae_shift_factor=vae_shift_factor, + offload_encoders=getattr(self.config, "offload_encoders", False), + ) + + def _place(self, params, shardings): + """Move host params onto the mesh in their target shardings.""" + if self.mesh is None or shardings is None: + return params + flat = flatten_dict(_plain(params)) + return unflatten_dict({path: jax.device_put(value, shardings[path]) for path, value in flat.items()}) + + def _metadata_mismatch(self, metadata) -> str: + """Cached weights are dtype-cast at save time, so a dtype change must miss.""" + expected = { + "model_id": self.config.pretrained_model_name_or_path, + "weights_dtype": str(self.config.weights_dtype), + } + for key, value in expected.items(): + if metadata.get(key) != value: + return f"{key} is {metadata.get(key)!r}, config asks for {value!r}" + return "" + + def _load_tokenizer(self, orbax_dir: str): + local = os.path.join(orbax_dir, "tokenizer") + if epath.Path(local).exists(): + return AutoTokenizer.from_pretrained(local, use_fast=True) + return AutoTokenizer.from_pretrained(self.config.pretrained_model_name_or_path, subfolder="tokenizer", use_fast=True) + + # ------------------------------------------------------------------- save + + def _save_pipeline(self, orbax_dir: str, pipeline: ZImagePipeline): + try: + max_logging.log(f"Saving Z-Image pipeline to orbax at {orbax_dir}") + manager = self._checkpoint_manager(orbax_dir) + model_id = self.config.pretrained_model_name_or_path + _, transformer_state, _ = nnx.split(pipeline.transformer, nnx.Param, ...) + metadata = { + "model_id": model_id, + "weights_dtype": str(self.config.weights_dtype), + "transformer_config": dict(ZImageTransformer2DModel.load_config(model_id, subfolder="transformer")), + "vae_config": _vae_constructor_kwargs(self._raw_vae_config(model_id)), + "vae_shift_factor": pipeline.vae_shift_factor, + "text_encoder_config": AutoConfig.from_pretrained(model_id, subfolder="text_encoder").to_dict(), + "trainable_items": list(TRAINABLE_ITEMS), + "non_trainable_items": list(NON_TRAINABLE_ITEMS), + } + manager.save( + 0, + args=ocp.args.Composite( + **{ + "transformer_state": ocp.args.StandardSave(transformer_state.to_pure_dict()), + "vae_state": ocp.args.StandardSave(_plain(pipeline.vae_params)), + "text_encoder_state": ocp.args.StandardSave( + nnx.split(pipeline.text_encoder, nnx.Param, ...)[1].to_pure_dict() + ), + CONFIG_ITEM: ocp.args.JsonSave(json.loads(json.dumps(metadata, default=str))), + } + ), + ) + manager.wait_until_finished() + pipeline.tokenizer.save_pretrained(os.path.join(orbax_dir, "tokenizer")) + max_logging.log(f"Z-Image pipeline saved to {orbax_dir}") + except Exception as e: # pylint: disable=broad-except + max_logging.log(f"Failed to save orbax checkpoint to {orbax_dir}: {e}") + + # ---------------------------------------------------------------- helpers + + def _checkpoint_manager(self, directory: str) -> ocp.CheckpointManager: + return create_orbax_checkpoint_manager( + directory, + enable_checkpointing=True, + save_interval_steps=1, + checkpoint_type=self.checkpoint_type, + use_async=False, + ) + + def _raw_vae_config(self, model_id: str) -> dict: + return dict(FlaxAutoencoderKL.load_config(model_id, subfolder="vae")) diff --git a/src/maxdiffusion/common_types.py b/src/maxdiffusion/common_types.py index 77da6900b..58e536686 100644 --- a/src/maxdiffusion/common_types.py +++ b/src/maxdiffusion/common_types.py @@ -53,6 +53,7 @@ WAN2_2 = "wan2.2" LTX2_VIDEO = "ltx2_video" LTX2_3 = "ltx2.3" +Z_IMAGE = "z_image" WAN_MODEL = WAN2_1 diff --git a/src/maxdiffusion/configs/base_zimage.yml b/src/maxdiffusion/configs/base_zimage.yml new file mode 100644 index 000000000..f6a0c24f7 --- /dev/null +++ b/src/maxdiffusion/configs/base_zimage.yml @@ -0,0 +1,157 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Official Z-Image inference defaults. Z-Image and Z-Image-Turbo share the +# denoiser; only the recommended denoising-step count differs (see +# base_zimage_turbo.yml). + +run_name: 'zimage' + +# If true save config to GCS in {base_output_directory}/{run_name}/ +save_config_to_gcs: False + +pretrained_model_name_or_path: 'Tongyi-MAI/Z-Image' +model_name: z_image +model_type: 'T2I' + +unet_checkpoint: '' +# This will convert the weights to this dtype. +# When running inference on TPUv5e, use weights_dtype: 'bfloat16' +weights_dtype: 'bfloat16' +# This sets the layer's dtype in the model. Ex: nn.Dense(dtype=activations_dtype) +activations_dtype: 'bfloat16' + +# Maximum sequence length for the text encoder +max_sequence_length: 512 +# offloads text encoder after text encoding to save memory. +offload_encoders: True + +# Attention +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +# Maxdiffusion has 2 types of attention sharding strategies: +# 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention) +# 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is +# sharded in cross attention q. +attention_sharding_uniform: True +# Empty dict lets the kernel pick defaults. Example v6e override: +# flash_block_sizes: { +# "block_q" : 1024, +# "block_kv_compute" : 1024, +# "block_kv" : 1024, +# "block_q_dkv" : 1024, +# "block_kv_dkv" : 1024, +# "block_kv_dkv_compute" : 1024, +# "block_q_dq" : 1024, +# "block_kv_dq" : 1024, +# "use_fused_bwd_kernel": False, +# } +flash_block_sizes: {} + +# Output directory +# Create a GCS bucket, e.g. my-maxdiffusion-outputs and set this to "gs://my-maxdiffusion-outputs/" +base_output_directory: "" +output_dir: "/mnt/disks/data/tmp/maxdiffusion" +# Local path for the generated image. If empty, an image named +# {run_name}_output_{seed}.png is written to the working directory. +output_file: "/mnt/disks/data/tmp/zimage.png" + +# Hardware +hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu' +skip_jax_distributed_system: True + +# Parallelism +mesh_axes: ['data', 'fsdp', 'context', 'tensor'] +# Z-Image runs sequence parallel (activations sharded over `context`) with +# data parallelism; the tensor axis is left at 1. `attention_sharding_uniform` +# adds the matching q/kv sequence rules on top of these. +logical_axis_rules: [ + ['batch', ['data', 'fsdp']], + ['activation_batch', ['data', 'fsdp']], + ['activation_length', 'context'], + ['activation_heads', 'tensor'], + ['heads', 'tensor'], + ['mlp', 'tensor'], + ['embed', 'fsdp'], + ['norm', 'tensor'], + ['out_channels', 'tensor'], + ] +data_sharding: [['data', 'fsdp', 'context', 'tensor']] + +# One axis for each parallelism type may hold a placeholder (-1) +# value to auto-shard based on available slices and devices. +# +# Measured on v6e-8, Turbo, 1024x1024, 9 steps, 8 images (per_device_batch_size 1.0): +# data=8 2.45s 0.307s/image <- default +# data=4 ctx=2 2.60s 0.325s/image +# data=2 ctx=4 3.13s 0.391s/image +# ctx=8 4.15s 0.519s/image +# fsdp=8 9.43s 1.179s/image +# The 12B of weights fit in HBM, so replicating them and giving each chip a +# whole image beats every sharded layout: no collectives at all. fsdp re-gathers +# every weight on all 9 denoising steps with no backward pass to amortize it. +# +# For single-image *latency* instead of throughput, generate one image +# (per_device_batch_size: 0.125) with ici_data_parallelism: 1 and +# ici_context_parallelism: -1, which splits one image's sequence 8 ways: 0.587s. +dcn_data_parallelism: 1 +dcn_fsdp_parallelism: 1 +dcn_context_parallelism: 1 +dcn_tensor_parallelism: 1 +ici_data_parallelism: -1 # recommended ICI axis to be auto-sharded +ici_fsdp_parallelism: 1 +ici_context_parallelism: 1 +ici_tensor_parallelism: 1 + +allow_split_physical_axes: False + +# Generation parameters +prompt: "A red fox reading a book in a quiet library, cinematic light" +height: 1024 +width: 1024 +num_inference_steps: 50 +# The released Z-Image configuration runs without classifier free guidance. +guidance_scale: 0.0 +seed: 42 +# Images per device. Fractional values are allowed: on 8 devices, +# 0.125 generates a single image. +per_device_batch_size: 1.0 +# Latents decoded per VAE call. The decoder is replicated and its +# activations are full-resolution, so raising this trades memory for speed. +vae_decode_chunk: 1 + +# Profiling +# generate_zimage runs an un-profiled warmup pass (compile), then a clean +# generation pass, and only then a profiled pass of profiler_steps steps. +enable_profiler: False +profiler_steps: 5 + +# ML Diagnostics settings +enable_ml_diagnostics: False +profiler_gcs_path: "" +enable_ondemand_xprof: False + +# Directory for the Orbax cache of the converted pipeline (transformer + VAE + +# text encoder). '' disables it. The first run loads from Diffusers and writes +# the cache; later runs restore from it and skip the torch -> flax conversion. +pretrained_orbax_dir: "" + +# Keys below are required by pyconfig but unused for Z-Image inference. +jax_cache_dir: "" +dataset_name: "" +dataset_save_location: "/mnt/disks/data/tmp" +learning_rate_schedule_steps: 1 +max_train_steps: 1 +compile_topology_num_slices: -1 +quantization_local_shard_count: -1 +use_qwix_quantization: False diff --git a/src/maxdiffusion/configs/base_zimage_turbo.yml b/src/maxdiffusion/configs/base_zimage_turbo.yml new file mode 100644 index 000000000..c99bdfb9b --- /dev/null +++ b/src/maxdiffusion/configs/base_zimage_turbo.yml @@ -0,0 +1,156 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Official Z-Image-Turbo inference defaults. The base model uses the same +# transformer; only its recommended denoising-step count differs. + +run_name: 'zimage-turbo' + +# If true save config to GCS in {base_output_directory}/{run_name}/ +save_config_to_gcs: False + +pretrained_model_name_or_path: 'Tongyi-MAI/Z-Image-Turbo' +model_name: z_image +model_type: 'T2I' + +unet_checkpoint: '' +# This will convert the weights to this dtype. +# When running inference on TPUv5e, use weights_dtype: 'bfloat16' +weights_dtype: 'bfloat16' +# This sets the layer's dtype in the model. Ex: nn.Dense(dtype=activations_dtype) +activations_dtype: 'bfloat16' + +# Maximum sequence length for the text encoder +max_sequence_length: 512 +# offloads text encoder after text encoding to save memory. +offload_encoders: True + +# Attention +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +# Maxdiffusion has 2 types of attention sharding strategies: +# 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention) +# 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is +# sharded in cross attention q. +attention_sharding_uniform: True +# Empty dict lets the kernel pick defaults. Example v6e override: +# flash_block_sizes: { +# "block_q" : 1024, +# "block_kv_compute" : 1024, +# "block_kv" : 1024, +# "block_q_dkv" : 1024, +# "block_kv_dkv" : 1024, +# "block_kv_dkv_compute" : 1024, +# "block_q_dq" : 1024, +# "block_kv_dq" : 1024, +# "use_fused_bwd_kernel": False, +# } +flash_block_sizes: {} + +# Output directory +# Create a GCS bucket, e.g. my-maxdiffusion-outputs and set this to "gs://my-maxdiffusion-outputs/" +base_output_directory: "" +output_dir: "/mnt/disks/data/tmp/maxdiffusion" +# Local path for the generated image. If empty, an image named +# {run_name}_output_{seed}.png is written to the working directory. +output_file: "/mnt/disks/data/tmp/zimage-turbo.png" + +# Hardware +hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu' +skip_jax_distributed_system: True + +# Parallelism +mesh_axes: ['data', 'fsdp', 'context', 'tensor'] +# Z-Image runs sequence parallel (activations sharded over `context`) with +# data parallelism; the tensor axis is left at 1. `attention_sharding_uniform` +# adds the matching q/kv sequence rules on top of these. +logical_axis_rules: [ + ['batch', ['data', 'fsdp']], + ['activation_batch', ['data', 'fsdp']], + ['activation_length', 'context'], + ['activation_heads', 'tensor'], + ['heads', 'tensor'], + ['mlp', 'tensor'], + ['embed', 'fsdp'], + ['norm', 'tensor'], + ['out_channels', 'tensor'], + ] +data_sharding: [['data', 'fsdp', 'context', 'tensor']] + +# One axis for each parallelism type may hold a placeholder (-1) +# value to auto-shard based on available slices and devices. +# +# Measured on v6e-8, Turbo, 1024x1024, 9 steps, 8 images (per_device_batch_size 1.0): +# data=8 2.45s 0.307s/image <- default +# data=4 ctx=2 2.60s 0.325s/image +# data=2 ctx=4 3.13s 0.391s/image +# ctx=8 4.15s 0.519s/image +# fsdp=8 9.43s 1.179s/image +# The 12B of weights fit in HBM, so replicating them and giving each chip a +# whole image beats every sharded layout: no collectives at all. fsdp re-gathers +# every weight on all 9 denoising steps with no backward pass to amortize it. +# +# For single-image *latency* instead of throughput, generate one image +# (per_device_batch_size: 0.125) with ici_data_parallelism: 1 and +# ici_context_parallelism: -1, which splits one image's sequence 8 ways: 0.587s. +dcn_data_parallelism: 1 +dcn_fsdp_parallelism: 1 +dcn_context_parallelism: 1 +dcn_tensor_parallelism: 1 +ici_data_parallelism: -1 # recommended ICI axis to be auto-sharded +ici_fsdp_parallelism: 1 +ici_context_parallelism: 1 +ici_tensor_parallelism: 1 + +allow_split_physical_axes: False + +# Generation parameters +prompt: "A red fox reading a book in a quiet library, cinematic light" +height: 1024 +width: 1024 +num_inference_steps: 9 +# The published Turbo configuration uses guidance_scale=0. +guidance_scale: 0.0 +seed: 42 +# Images per device. Fractional values are allowed: on 8 devices, +# 0.125 generates a single image. +per_device_batch_size: 1.0 +# Latents decoded per VAE call. The decoder is replicated and its +# activations are full-resolution, so raising this trades memory for speed. +vae_decode_chunk: 1 + +# Profiling +# generate_zimage runs an un-profiled warmup pass (compile), then a clean +# generation pass, and only then a profiled pass of profiler_steps steps. +enable_profiler: False +profiler_steps: 5 + +# ML Diagnostics settings +enable_ml_diagnostics: False +profiler_gcs_path: "" +enable_ondemand_xprof: False + +# Directory for the Orbax cache of the converted pipeline (transformer + VAE + +# text encoder). '' disables it. The first run loads from Diffusers and writes +# the cache; later runs restore from it and skip the torch -> flax conversion. +pretrained_orbax_dir: "" + +# Keys below are required by pyconfig but unused for Z-Image inference. +jax_cache_dir: "" +dataset_name: "" +dataset_save_location: "/mnt/disks/data/tmp" +learning_rate_schedule_steps: 1 +max_train_steps: 1 +compile_topology_num_slices: -1 +quantization_local_shard_count: -1 +use_qwix_quantization: False diff --git a/src/maxdiffusion/generate_flux2klein.py b/src/maxdiffusion/generate_flux2klein.py index b70fb637c..7956c850d 100644 --- a/src/maxdiffusion/generate_flux2klein.py +++ b/src/maxdiffusion/generate_flux2klein.py @@ -37,7 +37,8 @@ from maxdiffusion.models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel from maxdiffusion.models.vae_flax import FlaxAutoencoderKL -from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model, load_and_convert_qwen3_weights +from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model +from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler diff --git a/src/maxdiffusion/generate_ltx2.py b/src/maxdiffusion/generate_ltx2.py index 0445913ef..baeededd2 100644 --- a/src/maxdiffusion/generate_ltx2.py +++ b/src/maxdiffusion/generate_ltx2.py @@ -17,67 +17,14 @@ import jax.numpy as jnp import time import os -import subprocess from maxdiffusion.checkpointing.ltx2_checkpointer import LTX2Checkpointer from maxdiffusion import pyconfig, max_logging, max_utils from absl import app -from google.cloud import storage -from google.api_core.exceptions import GoogleAPIError import flax from maxdiffusion.utils.export_utils import export_to_video_with_audio from maxdiffusion.loaders.ltx2_lora_nnx_loader import LTX2NNXLoraLoader -def upload_video_to_gcs(output_dir: str, video_path: str): - """ - Uploads a local video file to a specified Google Cloud Storage bucket. - """ - try: - path_without_scheme = output_dir.removeprefix("gs://") - parts = path_without_scheme.split("/", 1) - bucket_name = parts[0] - folder_name = parts[1] if len(parts) > 1 else "" - - storage_client = storage.Client() - bucket = storage_client.bucket(bucket_name) - - source_file_path = f"./{video_path}" - destination_blob_name = os.path.join(folder_name, "videos", video_path) - - blob = bucket.blob(destination_blob_name) - - max_logging.log(f"Uploading {source_file_path} to {bucket_name}/{destination_blob_name}...") - blob.upload_from_filename(source_file_path) - max_logging.log(f"Upload complete {source_file_path}.") - - except GoogleAPIError as e: - max_logging.log(f"A storage error occurred during upload: {e}") - - -def delete_file(file_path: str): - if os.path.exists(file_path): - try: - os.remove(file_path) - max_logging.log(f"Successfully deleted file: {file_path}") - except OSError as e: - max_logging.log(f"Error deleting file '{file_path}': {e}") - else: - max_logging.log(f"The file '{file_path}' does not exist.") - - -def get_git_commit_hash(): - """Tries to get the current Git commit hash.""" - try: - commit_hash = subprocess.check_output(["git", "rev-parse", "HEAD"]).strip().decode("utf-8") - return commit_hash - except subprocess.CalledProcessError: - max_logging.log("Warning: 'git rev-parse HEAD' failed. Not running in a git repo?") - return None - except FileNotFoundError: - max_logging.log("Warning: 'git' command not found.") - return None - - jax.config.update("jax_use_shardy_partitioner", True) @@ -253,7 +200,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): saved_video_path.append(video_path) if config.output_dir.startswith("gs://"): - upload_video_to_gcs(os.path.join(config.output_dir, config.run_name), video_path) + max_utils.upload_file_to_gcs(os.path.join(config.output_dir, config.run_name), video_path, subdir="videos") timing_str = ( f"\n{'=' * 50}\n" @@ -322,7 +269,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): def main(argv: Sequence[str]) -> None: - commit_hash = get_git_commit_hash() + commit_hash = max_utils.get_git_commit_hash() pyconfig.initialize(argv) try: flax.config.update("flax_always_shard_variable", False) diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py index 348bef486..01199a303 100644 --- a/src/maxdiffusion/generate_wan.py +++ b/src/maxdiffusion/generate_wan.py @@ -16,7 +16,6 @@ import jax import time import os -import subprocess from maxdiffusion.checkpointing.wan_checkpointer_2_1 import WanCheckpointer2_1 from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p1 import WanCheckpointerI2V_2_1 @@ -26,7 +25,6 @@ from maxdiffusion.train_utils import transformer_engine_context from maxdiffusion.utils import export_to_video from maxdiffusion.utils.loading_utils import load_image -from google.cloud import storage import flax from maxdiffusion.common_types import WAN2_1, WAN2_2 from maxdiffusion.loaders.wan_lora_nnx_loader import Wan2_1NNXLoraLoader, Wan2_2NNXLoraLoader @@ -36,56 +34,6 @@ from maxdiffusion.pipelines.wan.wan_pipeline_i2v_2p2 import WanPipelineI2V_2_2 -def upload_video_to_gcs(output_dir: str, video_path: str): - """ - Uploads a local video file to a specified Google Cloud Storage bucket. - """ - try: - path_without_scheme = output_dir.removeprefix("gs://") - parts = path_without_scheme.split("/", 1) - bucket_name = parts[0] - folder_name = parts[1] if len(parts) > 1 else "" - - storage_client = storage.Client() - bucket = storage_client.bucket(bucket_name) - - source_file_path = f"./{video_path}" - destination_blob_name = os.path.join(folder_name, "videos", video_path) - - blob = bucket.blob(destination_blob_name) - - max_logging.log(f"Uploading {source_file_path} to {bucket_name}/{destination_blob_name}...") - blob.upload_from_filename(source_file_path) - max_logging.log(f"Upload complete {source_file_path}.") - - except Exception as e: - max_logging.log(f"An error occurred: {e}") - - -def delete_file(file_path: str): - if os.path.exists(file_path): - try: - os.remove(file_path) - max_logging.log(f"Successfully deleted file: {file_path}") - except OSError as e: - max_logging.log(f"Error deleting file '{file_path}': {e}") - else: - max_logging.log(f"The file '{file_path}' does not exist.") - - -def get_git_commit_hash(): - """Tries to get the current Git commit hash.""" - try: - commit_hash = subprocess.check_output(["git", "rev-parse", "HEAD"]).strip().decode("utf-8") - return commit_hash - except subprocess.CalledProcessError: - max_logging.log("Warning: 'git rev-parse HEAD' failed. Not running in a git repo?") - return None - except FileNotFoundError: - max_logging.log("Warning: 'git' command not found.") - return None - - jax.config.update("jax_use_shardy_partitioner", True) @@ -189,9 +137,9 @@ def inference_generate_video(config, pipeline, filename_prefix=""): video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4" export_to_video(videos[i], video_path, fps=config.fps) if config.output_dir.startswith("gs://"): - upload_video_to_gcs(os.path.join(config.output_dir, config.run_name), video_path) + max_utils.upload_file_to_gcs(os.path.join(config.output_dir, config.run_name), video_path, subdir="videos") # Delete local files to avoid storing too manys videos - delete_file(f"./{video_path}") + max_utils.delete_file(f"./{video_path}") return @@ -414,7 +362,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): export_to_video(videos[i], video_path, fps=config.fps) saved_video_path.append(video_path) if config.output_dir.startswith("gs://"): - upload_video_to_gcs(os.path.join(config.output_dir, config.run_name), video_path) + max_utils.upload_file_to_gcs(os.path.join(config.output_dir, config.run_name), video_path, subdir="videos") max_logging.log(f"generation_time: {generation_time}") if writer and jax.process_index() == 0: writer.add_scalar("inference/generation_time", generation_time, global_step=0) @@ -480,7 +428,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): def main(argv: Sequence[str]) -> None: - commit_hash = get_git_commit_hash() + commit_hash = max_utils.get_git_commit_hash() pyconfig.initialize(argv) try: flax.config.update("flax_always_shard_variable", False) diff --git a/src/maxdiffusion/generate_zimage.py b/src/maxdiffusion/generate_zimage.py new file mode 100644 index 000000000..2402aa9b5 --- /dev/null +++ b/src/maxdiffusion/generate_zimage.py @@ -0,0 +1,167 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Generate an image with Z-Image or Z-Image-Turbo.""" + +import time +from typing import Sequence + +from absl import app +import jax +from jax.sharding import Mesh + +from maxdiffusion import max_logging, max_utils, pyconfig +from maxdiffusion.checkpointing.z_image_checkpointer import ZImageCheckpointer +from maxdiffusion.max_utils import create_device_mesh + + +def call_pipeline(config, pipeline, num_inference_steps=None, return_timings=False): + if num_inference_steps is None: + num_inference_steps = config.num_inference_steps + # Using global_batch_size_to_train_on so not to create more config variables: + # per_device_batch_size images are generated on every device. + return pipeline( + [config.prompt] * config.global_batch_size_to_train_on, + height=config.height, + width=config.width, + num_inference_steps=num_inference_steps, + guidance_scale=config.guidance_scale, + seed=config.seed, + max_sequence_length=config.max_sequence_length, + vae_decode_chunk=config.vae_decode_chunk, + return_timings=return_timings, + ) + + +def run(config, commit_hash=None): + writer = max_utils.initialize_summary_writer(config) + if jax.process_index() == 0 and writer: + max_logging.log(f"TensorBoard logs will be written to: {config.tensorboard_dir}") + + if commit_hash: + writer.add_text("inference/git_commit_hash", commit_hash, global_step=0) + max_logging.log(f"Git Commit Hash: {commit_hash}") + else: + max_logging.log("Could not retrieve Git commit hash.") + + load_start = time.perf_counter() + mesh = Mesh(create_device_mesh(config), config.mesh_axes) + pipeline = ZImageCheckpointer(config, mesh).load_pipeline() + load_time = time.perf_counter() - load_start + max_logging.log(f"load_time: {load_time:.1f}s") + + max_logging.log("===================== Model details =======================") + max_logging.log(f"model name: {config.model_name}") + max_logging.log(f"model path: {config.pretrained_model_name_or_path}") + max_logging.log(f"model type: {config.model_type}") + max_logging.log(f"hardware: {jax.devices()[0].platform}") + max_logging.log(f"number of devices: {jax.device_count()}") + max_logging.log(f"per_device_batch_size: {config.per_device_batch_size}") + max_logging.log("============================================================") + + max_logging.log( + f"Num steps: {config.num_inference_steps}, height: {config.height}, width: {config.width}," + f" images: {config.global_batch_size_to_train_on}" + ) + + # Nothing in the pipeline reads the profiler flags -- the profiler is started + # explicitly around run 3 below -- so the config is never mutated here. + profiling_requested = config.get_keys().get("enable_profiler", False) or config.get_keys().get( + "enable_ml_diagnostics", False + ) + + # --------------------------------------------------------- + # Run 1: warmup compilation, nothing profiled. + # --------------------------------------------------------- + s0 = time.perf_counter() + # Warm up at the real step count. A shorter warmup would leave the + # step-count-dependent work (the sigma schedule, and anything XLA has not + # already cached at these shapes) to compile inside the timed run. + max_logging.log(f"🚀 Starting warmup compilation pass ({config.num_inference_steps} denoising steps)...") + _ = call_pipeline(config, pipeline) + compile_time = time.perf_counter() - s0 + max_logging.log(f"compile_time: {compile_time}") + if writer and jax.process_index() == 0: + writer.add_scalar("inference/compile_time", compile_time, global_step=0) + + # --------------------------------------------------------- + # Run 2: the real generation, profiling still disabled. + # --------------------------------------------------------- + s0 = time.perf_counter() + max_logging.log("🚀 Starting full-length generation pass...") + images, trace = call_pipeline(config, pipeline, return_timings=True) + generation_time = time.perf_counter() - s0 + max_logging.log(f"generation_time: {generation_time}") + generation_time_per_image = generation_time / len(images) + max_logging.log(f"generation time per image: {generation_time_per_image}") + if writer and jax.process_index() == 0: + writer.add_scalar("inference/generation_time", generation_time, global_step=0) + writer.add_scalar("inference/generation_time_per_image", generation_time_per_image, global_step=0) + + image_paths = max_utils.save_images(config, images) + + summary = [ + f"\n{'=' * 50}", + " TIMING SUMMARY", + f"{'=' * 50}", + f" Load (checkpoint): {load_time:>7.1f}s", + f" Compile: {compile_time:>7.1f}s", + f" Inference: {generation_time:>7.1f}s", + f" Per image ({len(images):>2d}): {generation_time_per_image:>7.1f}s", + ] + if trace: + summary.extend([ + f" {'─' * 40}", + f" Text Encoding: {trace.get('text_encode', 0.0):>7.1f}s", + f" Denoise Total: {trace.get('denoise', 0.0):>7.1f}s", + f" VAE Decode: {trace.get('vae_decode', 0.0):>7.1f}s", + f" Host Formatting: {trace.get('host_post', 0.0):>7.1f}s", + ]) + summary.append(f"{'=' * 50}") + max_logging.log("\n".join(summary)) + + # --------------------------------------------------------- + # Run 3: profiled pass, only if profiling was originally enabled. + # --------------------------------------------------------- + if profiling_requested: + profiling_steps = config.get_keys().get("profiler_steps", 5) + max_logging.log(f"🚀 Warmup for profiling pass ({profiling_steps} denoising steps)...") + _ = call_pipeline(config, pipeline, num_inference_steps=profiling_steps) + + max_logging.log(f"🚀 Starting profiling run ({profiling_steps} denoising steps)...") + profiler = max_utils.Profiler(config, session_name=f"denoise_profile_{profiling_steps}_steps") + profiler.start() + s0 = time.perf_counter() + _ = call_pipeline(config, pipeline, num_inference_steps=profiling_steps) + generation_time_with_profiler = time.perf_counter() - s0 + profiler.stop() + + max_logging.log(f"generation_time_with_profiler: {generation_time_with_profiler}") + if writer and jax.process_index() == 0: + writer.add_scalar("inference/generation_time_with_profiler", generation_time_with_profiler, global_step=0) + + return image_paths + + +def main(argv: Sequence[str]) -> None: + commit_hash = max_utils.get_git_commit_hash() + pyconfig.initialize(argv) + max_utils.ensure_machinelearning_job_runs(pyconfig.config) + run(pyconfig.config, commit_hash=commit_hash) + + +if __name__ == "__main__": + app.run(main) diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index 1ab5407ac..28b747894 100644 --- a/src/maxdiffusion/max_utils.py +++ b/src/maxdiffusion/max_utils.py @@ -296,6 +296,70 @@ def upload_blob(destination_gcs_name, source_file_name): blob.upload_from_filename(source_file_name) +def save_images(config, images): + """Writes each image locally, and to GCS when output_dir is a bucket.""" + stem, extension = os.path.splitext(config.output_file or f"{config.run_name}_output_{config.seed}.png") + paths = [] + for index, image in enumerate(images): + path = f"{stem}{extension}" if len(images) == 1 else f"{stem}_{index}{extension}" + parent_dir = os.path.dirname(path) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + image.save(path) + paths.append(path) + if config.output_dir.startswith("gs://"): + upload_file_to_gcs(os.path.join(config.output_dir, config.run_name), path, subdir="images") + max_logging.log(f"Saved {len(paths)} image(s), first: {paths[0]}") + return paths + + +def upload_file_to_gcs(output_dir: str, file_path: str, subdir: str = ""): + """Uploads one generated file to {output_dir}/{subdir}/, logging failures. + + Shared by the generate_* drivers, which each write their own media type. + """ + try: + path_without_scheme = output_dir.removeprefix("gs://") + parts = path_without_scheme.split("/", 1) + bucket_name = parts[0] + folder_name = parts[1] if len(parts) > 1 else "" + destination_blob_name = os.path.join(folder_name, subdir, os.path.basename(file_path)) + + storage_client = storage.Client() + bucket = storage_client.bucket(bucket_name) + blob = bucket.blob(destination_blob_name) + + max_logging.log(f"Uploading {file_path} to {bucket_name}/{destination_blob_name}...") + blob.upload_from_filename(file_path) + max_logging.log(f"Upload complete {file_path}.") + except Exception as e: # pylint: disable=broad-except + max_logging.log(f"An error occurred: {e}") + + +def delete_file(file_path: str): + """Removes a local file, e.g. after it has been uploaded to GCS.""" + if os.path.exists(file_path): + try: + os.remove(file_path) + max_logging.log(f"Successfully deleted file: {file_path}") + except OSError as e: + max_logging.log(f"Error deleting file '{file_path}': {e}") + else: + max_logging.log(f"The file '{file_path}' does not exist.") + + +def get_git_commit_hash(): + """Tries to get the current Git commit hash, for run provenance.""" + try: + return subprocess.check_output(["git", "rev-parse", "HEAD"]).strip().decode("utf-8") + except subprocess.CalledProcessError: + max_logging.log("Warning: 'git rev-parse HEAD' failed. Not running in a git repo?") + return None + except FileNotFoundError: + max_logging.log("Warning: 'git' command not found.") + return None + + def walk_and_upload_blobs(config, output_dir): user_dir = os.path.expanduser("~") uploaded_files = set() diff --git a/src/maxdiffusion/models/__init__.py b/src/maxdiffusion/models/__init__.py index 7ff8fd8fb..d82305aa3 100644 --- a/src/maxdiffusion/models/__init__.py +++ b/src/maxdiffusion/models/__init__.py @@ -21,12 +21,14 @@ _import_structure["controlnet_flax"] = ["FlaxControlNetModel"] _import_structure["unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"] _import_structure["vae_flax"] = ["FlaxAutoencoderKL"] +_import_structure["z_image.transformer_z_image"] = ["ZImageTransformer2DModel"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: from .controlnet_flax import FlaxControlNetModel from .unet_2d_condition_flax import FlaxUNet2DConditionModel from .vae_flax import FlaxAutoencoderKL + from .z_image.transformer_z_image import ZImageTransformer2DModel from .lora import * from .flux.transformers.transformer_flux_flax import FluxTransformer2DModel from .ltx_video.transformers.transformer3d import Transformer3DModel diff --git a/src/maxdiffusion/models/qwen3_flax.py b/src/maxdiffusion/models/qwen3_flax.py index 0a76544dd..b3ca43003 100644 --- a/src/maxdiffusion/models/qwen3_flax.py +++ b/src/maxdiffusion/models/qwen3_flax.py @@ -20,9 +20,6 @@ import flax.linen as nn import jax import jax.numpy as jnp -import numpy as np - -from maxdiffusion import max_logging # ----------------------------------------------------------------------------- # Qwen3 Configuration @@ -67,12 +64,13 @@ class FlaxQwen3RMSNorm(nn.Module): dim: int eps: float = 1e-6 dtype: Any = jnp.float32 + param_dtype: Any = jnp.float32 @nn.compact def __call__(self, x): x_float = x.astype(jnp.float32) variance = jnp.mean(jnp.square(x_float), axis=-1, keepdims=True) - scale = self.param("weight", nn.initializers.ones, (self.dim,), self.dtype) + scale = self.param("weight", nn.initializers.ones, (self.dim,), self.param_dtype) normed = x_float * jax.lax.rsqrt(variance + self.eps) return (normed.astype(self.dtype)) * scale @@ -87,6 +85,7 @@ def __call__(self, x): use_bias=False, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), dtype=self.config.dtype, + param_dtype=self.config.dtype, name="gate_proj", ) up_proj = nn.Dense( @@ -94,6 +93,7 @@ def __call__(self, x): use_bias=False, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), dtype=self.config.dtype, + param_dtype=self.config.dtype, name="up_proj", ) down_proj = nn.Dense( @@ -101,6 +101,7 @@ def __call__(self, x): use_bias=False, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")), dtype=self.config.dtype, + param_dtype=self.config.dtype, name="down_proj", ) @@ -178,6 +179,7 @@ def __call__( use_bias=False, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "heads")), dtype=self.config.dtype, + param_dtype=self.config.dtype, name="q_proj", ) k_proj = nn.Dense( @@ -185,6 +187,7 @@ def __call__( use_bias=False, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "heads")), dtype=self.config.dtype, + param_dtype=self.config.dtype, name="k_proj", ) v_proj = nn.Dense( @@ -192,6 +195,7 @@ def __call__( use_bias=False, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "heads")), dtype=self.config.dtype, + param_dtype=self.config.dtype, name="v_proj", ) o_proj = nn.Dense( @@ -199,6 +203,7 @@ def __call__( use_bias=False, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("heads", "embed")), dtype=self.config.dtype, + param_dtype=self.config.dtype, name="o_proj", ) @@ -359,6 +364,7 @@ def __call__( nn.initializers.normal(stddev=self.config.hidden_size**-0.5), ("vocab", "embed") ), dtype=self.config.dtype, + param_dtype=self.config.dtype, name="embed_tokens", ) hidden_states = embed_tokens(input_ids) @@ -408,10 +414,18 @@ def __call__( class NNXFlaxQwen3RMSNorm(nnx.Module): - def __init__(self, rngs: nnx.Rngs, dim: int, eps: float = 1e-6, dtype: jnp.dtype = jnp.float32): + def __init__( + self, + rngs: nnx.Rngs, + dim: int, + eps: float = 1e-6, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): self.eps = eps self.dtype = dtype - self.weight = nnx.Param(jnp.ones((dim,), dtype=dtype)) + # Held in float32 for the same reason as the Linen variant above. + self.weight = nnx.Param(jnp.ones((dim,), dtype=param_dtype)) def __call__(self, x: jnp.ndarray) -> jnp.ndarray: x_float = x.astype(jnp.float32) @@ -430,6 +444,7 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config): use_bias=False, kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "mlp")), dtype=config.dtype, + param_dtype=config.dtype, rngs=rngs, ) self.up_proj = nnx.Linear( @@ -438,6 +453,7 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config): use_bias=False, kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "mlp")), dtype=config.dtype, + param_dtype=config.dtype, rngs=rngs, ) self.down_proj = nnx.Linear( @@ -446,6 +462,7 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config): use_bias=False, kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("mlp", "embed")), dtype=config.dtype, + param_dtype=config.dtype, rngs=rngs, ) @@ -467,6 +484,7 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config): use_bias=False, kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "heads")), dtype=config.dtype, + param_dtype=config.dtype, rngs=rngs, ) self.k_proj = nnx.Linear( @@ -475,6 +493,7 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config): use_bias=False, kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "heads")), dtype=config.dtype, + param_dtype=config.dtype, rngs=rngs, ) self.v_proj = nnx.Linear( @@ -483,6 +502,7 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config): use_bias=False, kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "heads")), dtype=config.dtype, + param_dtype=config.dtype, rngs=rngs, ) self.o_proj = nnx.Linear( @@ -491,6 +511,7 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config): use_bias=False, kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("heads", "embed")), dtype=config.dtype, + param_dtype=config.dtype, rngs=rngs, ) @@ -531,14 +552,28 @@ def __call__( k = jnp.transpose(k, (0, 2, 1, 3)) v = jnp.transpose(v, (0, 2, 1, 3)) - scale = 1.0 / math.sqrt(self.head_dim) - scores = jnp.matmul(q, jnp.transpose(k, (0, 1, 3, 2))) * scale + # TODO: this upcasts both matmuls to float32 only to stay bit-comparable + # with the Linen implementation above. Upstream (transformers' Qwen3) runs + # the QK and probs@V matmuls in the model dtype and casts only the softmax + # to float32, which is the part that actually needs the range. Doing the + # same here would keep both matmuls on the MXU at full bfloat16 rate; it is + # left alone for now so the two implementations stay comparable. + q_f = q.astype(jnp.float32) + k_f = k.astype(jnp.float32) + v_f = v.astype(jnp.float32) + scores = jnp.matmul(q_f, jnp.transpose(k_f, (0, 1, 3, 2))) / math.sqrt(self.head_dim) + # Qwen3 is a causal LM: without this every token attends to the future. + causal_mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_)) + scores = jnp.where(causal_mask, scores, -1e4) + + # `attention_mask` is a 0/1 padding mask, not an additive one. if attention_mask is not None: - scores = scores + attention_mask + p_mask = attention_mask[:, jnp.newaxis, jnp.newaxis, :].astype(jnp.bool_) + scores = jnp.where(p_mask, scores, -1e4) attn_probs = jax.nn.softmax(scores, axis=-1) - output = jnp.matmul(attn_probs, v) + output = jnp.matmul(attn_probs, v_f).astype(self.config.dtype) output = jnp.transpose(output, (0, 2, 1, 3)) output = output.reshape(batch_size, seq_len, -1) output = self.o_proj(output) @@ -589,6 +624,7 @@ def __init__(self, rngs: nnx.Rngs, config: FlaxQwen3Config): features=config.hidden_size, embedding_init=nnx.with_partitioning(nnx.initializers.normal(stddev=config.hidden_size**-0.5), ("vocab", "embed")), dtype=config.dtype, + param_dtype=config.dtype, rngs=rngs, ) self.layers = nnx.List([NNXFlaxQwen3DecoderLayer(rngs=rngs, config=config) for _ in range(config.num_hidden_layers)]) @@ -619,103 +655,3 @@ def __call__( hidden_states = self.norm(hidden_states) return hidden_states, all_hidden_states - - -# ----------------------------------------------------------------------------- -# Weight Mapping & Conversion Utilities -# ----------------------------------------------------------------------------- - - -def load_and_convert_qwen3_weights(safetensors_path: str, jax_params: dict, config: FlaxQwen3Config) -> dict: - """ - Loads weights from safetensors via zero-copy safetensors.numpy and converts them to JAX parameter dictionary. - """ - import glob - import os - from safetensors.numpy import load_file - - torch_weights: dict = {} - if os.path.isdir(safetensors_path): - # Find all safetensors shards - shards = glob.glob(os.path.join(safetensors_path, "*.safetensors")) - max_logging.log(f"Loading sharded Qwen3 weights from directory: {safetensors_path} (Found {len(shards)} shards)...") - for shard in sorted(shards): - max_logging.log(f"Loading shard: {shard}...") - torch_weights.update(load_file(shard)) - else: - # Single file path - max_logging.log(f"Loading Qwen3 weights from file: {safetensors_path}...") - torch_weights = load_file(safetensors_path) - max_logging.log("Safetensors weights loaded successfully. Starting JAX parameter mapping...") - - # Helper to transpose and cast weight - def get_w(name: str, transpose: bool = True) -> np.ndarray: - nonlocal torch_weights - if name not in torch_weights: - raise KeyError(f"Weight '{name}' not found in safetensors!") - t = torch_weights[name] - if len(t.shape) == 2 and transpose: - t = t.T - return t - - # Create mutable copy of JAX params to populate - import flax - - flat_params = flax.traverse_util.flatten_dict(jax_params) - converted_flat = {} - - for k, v in flat_params.items(): - # Reconstruct path string for debugging/matching - path_str = ".".join(k) - - # 1. Token Embeddings - if k[0] == "embed_tokens" and k[1] == "embedding": - converted_flat[k] = get_w("model.embed_tokens.weight", transpose=False) - - # 2. Decoder Layer Normalizations (RMSNorm) - elif "input_layernorm" in path_str and k[-1] == "weight": - layer_idx = k[0].split("_")[1] - converted_flat[k] = get_w(f"model.layers.{layer_idx}.input_layernorm.weight") - - elif "post_attention_layernorm" in path_str and k[-1] == "weight": - layer_idx = k[0].split("_")[1] - converted_flat[k] = get_w(f"model.layers.{layer_idx}.post_attention_layernorm.weight") - - # 3. Attention Projections & QK-Norm - elif "self_attn" in path_str and k[-1] == "kernel": - layer_idx = k[0].split("_")[1] - proj_name = k[2] # q_proj, k_proj, v_proj, o_proj - converted_flat[k] = get_w(f"model.layers.{layer_idx}.self_attn.{proj_name}.weight") - - elif "self_attn" in path_str and "q_norm" in path_str and k[-1] == "weight": - layer_idx = k[0].split("_")[1] - converted_flat[k] = get_w(f"model.layers.{layer_idx}.self_attn.q_norm.weight") - - elif "self_attn" in path_str and "k_norm" in path_str and k[-1] == "weight": - layer_idx = k[0].split("_")[1] - converted_flat[k] = get_w(f"model.layers.{layer_idx}.self_attn.k_norm.weight") - - # 4. MLP Block - elif "mlp" in path_str and k[-1] == "kernel": - layer_idx = k[0].split("_")[1] - proj_name = k[2] # gate_proj, up_proj, down_proj - converted_flat[k] = get_w(f"model.layers.{layer_idx}.mlp.{proj_name}.weight") - - # 5. Final RMSNorm - elif k[0] == "norm" and k[1] == "weight": - converted_flat[k] = get_w("model.norm.weight") - - else: - max_logging.log(f"WARNING: JAX parameter '{path_str}' did not match any PyTorch weights!") - converted_flat[k] = np.zeros(v.shape, dtype=np.float32) if hasattr(v, "shape") and not isinstance(v, np.ndarray) else v - - # Clean up PyTorch memory immediately - del torch_weights - import gc - - gc.collect() - - res = flax.traverse_util.unflatten_dict(converted_flat) - return jax.tree_util.tree_map( - lambda leaf: jnp.zeros(leaf.shape, dtype=leaf.dtype) if isinstance(leaf, jax.ShapeDtypeStruct) else leaf, res - ) diff --git a/src/maxdiffusion/models/qwen3_utils.py b/src/maxdiffusion/models/qwen3_utils.py new file mode 100644 index 000000000..4c3182cac --- /dev/null +++ b/src/maxdiffusion/models/qwen3_utils.py @@ -0,0 +1,204 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Checkpoint conversion helpers for PyTorch Qwen3 checkpoints. + +Kept apart from `qwen3_flax`, which is the model definition alone. +""" + +from typing import Optional, Tuple + +from flax.traverse_util import flatten_dict, unflatten_dict +import jax +import jax.numpy as jnp +import numpy as np + +from maxdiffusion import max_logging +from maxdiffusion.models.qwen3_flax import FlaxQwen3Config + + +def qwen3_flax_key_to_pytorch_key(path: Tuple[str, ...]) -> Tuple[str, bool]: + """Return the PyTorch key feeding a Flax param path, and whether it transposes. + + Handles both module layouts: NNX nests its blocks in an `nnx.List`, giving + ('layers', 0, ...), while Linen names them flatly, giving ('layers_0', ...). + Every Flax `.kernel` is a torch `nn.Linear` weight and needs a transpose; + norm scales and the embedding table do not. + """ + if path[0] == "embed_tokens": + return "model.embed_tokens.weight", False + if path[0] == "norm": + return "model.norm.weight", False + if path[0] == "layers": # NNX: ('layers', index, ...) + layer_index, inner = path[1], path[2:] + elif path[0].startswith("layers_"): # Linen: ('layers_{index}', ...) + layer_index, inner = path[0].split("_")[1], path[1:] + else: + raise KeyError(f"Unknown Qwen3 Flax parameter path `{path}`.") + return f"model.layers.{layer_index}." + ".".join(inner[:-1]) + ".weight", path[-1] == "kernel" + + +def load_qwen3_weights( + safetensors_path: str, + eval_shapes: dict, + target_shardings: Optional[dict] = None, + device: str = "cpu", +) -> dict: + """Stream a PyTorch Qwen3 checkpoint into a Flax parameter tree. + + Streaming: tensors are read and converted one at a time through the torch + backend, so the full PyTorch state dict is never held on the host alongside + the converted tree, and each parameter is placed on device as soon as it is + read. The torch backend also means bfloat16 checkpoints load, which + `load_and_convert_qwen3_weights` below cannot do (numpy has no bfloat16). + Each parameter takes the dtype and sharding of its entry in `eval_shapes`. + """ + import glob + import os + from safetensors import safe_open + + if os.path.isdir(safetensors_path): + shards = sorted(glob.glob(os.path.join(safetensors_path, "*.safetensors"))) + else: + shards = [safetensors_path] + if not shards: + raise ValueError(f"No safetensors found in {safetensors_path}") + + expected = flatten_dict(eval_shapes) + sources = {} + for path in expected: + source_key, transpose = qwen3_flax_key_to_pytorch_key(path) + sources[source_key] = (path, transpose) + + converted = {} + cpu = jax.local_devices(backend=device)[0] + for shard in shards: + max_logging.log(f"Loading Qwen3 shard: {os.path.basename(shard)}...") + with safe_open(shard, framework="pt", device="cpu") as tensors: + for source_key in tensors.keys(): + # Checkpoints carry tensors this text-encoder stack does not use + # (lm_head, and rotary buffers), which are skipped here. + if source_key not in sources: + continue + target_key, transpose = sources[source_key] + value = tensors.get_tensor(source_key).float().numpy() + if transpose: + value = value.T + value = jnp.asarray(value, dtype=expected[target_key].dtype) + if value.shape != expected[target_key].shape: + raise ValueError(f"Shape mismatch for `{source_key}`: {value.shape} != {expected[target_key].shape}.") + target = target_shardings.get(target_key) if target_shardings is not None else cpu + converted[target_key] = jax.device_put(value, target) + + missing = set(expected) - set(converted) + if missing: + raise ValueError(f"Qwen3 checkpoint is missing {len(missing)} parameters, e.g. {sorted(missing)[:5]}.") + return unflatten_dict(converted) + + +def load_and_convert_qwen3_weights(safetensors_path: str, jax_params: dict, config: FlaxQwen3Config) -> dict: + """ + Loads weights from safetensors via zero-copy safetensors.numpy and converts them to JAX parameter dictionary. + """ + import glob + import os + from safetensors.numpy import load_file + + torch_weights: dict = {} + if os.path.isdir(safetensors_path): + # Find all safetensors shards + shards = glob.glob(os.path.join(safetensors_path, "*.safetensors")) + max_logging.log(f"Loading sharded Qwen3 weights from directory: {safetensors_path} (Found {len(shards)} shards)...") + for shard in sorted(shards): + max_logging.log(f"Loading shard: {shard}...") + torch_weights.update(load_file(shard)) + else: + # Single file path + max_logging.log(f"Loading Qwen3 weights from file: {safetensors_path}...") + torch_weights = load_file(safetensors_path) + max_logging.log("Safetensors weights loaded successfully. Starting JAX parameter mapping...") + + # Helper to transpose and cast weight + def get_w(name: str, transpose: bool = True) -> np.ndarray: + nonlocal torch_weights + if name not in torch_weights: + raise KeyError(f"Weight '{name}' not found in safetensors!") + t = torch_weights[name] + if len(t.shape) == 2 and transpose: + t = t.T + return t + + # Create mutable copy of JAX params to populate + import flax + + flat_params = flax.traverse_util.flatten_dict(jax_params) + converted_flat = {} + + for k, v in flat_params.items(): + # Reconstruct path string for debugging/matching + path_str = ".".join(k) + + # 1. Token Embeddings + if k[0] == "embed_tokens" and k[1] == "embedding": + converted_flat[k] = get_w("model.embed_tokens.weight", transpose=False) + + # 2. Decoder Layer Normalizations (RMSNorm) + elif "input_layernorm" in path_str and k[-1] == "weight": + layer_idx = k[0].split("_")[1] + converted_flat[k] = get_w(f"model.layers.{layer_idx}.input_layernorm.weight") + + elif "post_attention_layernorm" in path_str and k[-1] == "weight": + layer_idx = k[0].split("_")[1] + converted_flat[k] = get_w(f"model.layers.{layer_idx}.post_attention_layernorm.weight") + + # 3. Attention Projections & QK-Norm + elif "self_attn" in path_str and k[-1] == "kernel": + layer_idx = k[0].split("_")[1] + proj_name = k[2] # q_proj, k_proj, v_proj, o_proj + converted_flat[k] = get_w(f"model.layers.{layer_idx}.self_attn.{proj_name}.weight") + + elif "self_attn" in path_str and "q_norm" in path_str and k[-1] == "weight": + layer_idx = k[0].split("_")[1] + converted_flat[k] = get_w(f"model.layers.{layer_idx}.self_attn.q_norm.weight") + + elif "self_attn" in path_str and "k_norm" in path_str and k[-1] == "weight": + layer_idx = k[0].split("_")[1] + converted_flat[k] = get_w(f"model.layers.{layer_idx}.self_attn.k_norm.weight") + + # 4. MLP Block + elif "mlp" in path_str and k[-1] == "kernel": + layer_idx = k[0].split("_")[1] + proj_name = k[2] # gate_proj, up_proj, down_proj + converted_flat[k] = get_w(f"model.layers.{layer_idx}.mlp.{proj_name}.weight") + + # 5. Final RMSNorm + elif k[0] == "norm" and k[1] == "weight": + converted_flat[k] = get_w("model.norm.weight") + + else: + max_logging.log(f"WARNING: JAX parameter '{path_str}' did not match any PyTorch weights!") + converted_flat[k] = np.zeros(v.shape, dtype=np.float32) if hasattr(v, "shape") and not isinstance(v, np.ndarray) else v + + # Clean up PyTorch memory immediately + del torch_weights + import gc + + gc.collect() + + res = flax.traverse_util.unflatten_dict(converted_flat) + return jax.tree_util.tree_map( + lambda leaf: jnp.zeros(leaf.shape, dtype=leaf.dtype) if isinstance(leaf, jax.ShapeDtypeStruct) else leaf, res + ) diff --git a/src/maxdiffusion/models/z_image/__init__.py b/src/maxdiffusion/models/z_image/__init__.py new file mode 100644 index 000000000..8fea11219 --- /dev/null +++ b/src/maxdiffusion/models/z_image/__init__.py @@ -0,0 +1,19 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Z-Image model components.""" + +from .transformer_z_image import ZImageTransformer2DModel diff --git a/src/maxdiffusion/models/z_image/transformer_z_image.py b/src/maxdiffusion/models/z_image/transformer_z_image.py new file mode 100644 index 000000000..d7f009316 --- /dev/null +++ b/src/maxdiffusion/models/z_image/transformer_z_image.py @@ -0,0 +1,670 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""JAX/NNX implementation of the Diffusers Z-Image transformer. + +Z-Image and Z-Image-Turbo use the same denoiser. The model deliberately +keeps the variable-length image/text representation from Diffusers: it avoids +padding an image before patchification and only pads the joint sequence when a +batch contains different resolutions or prompt lengths. +""" + +from typing import Optional, Sequence + +import math + +import flax +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxdiffusion.common_types import ( + BATCH, + BlockSizes, + D_KV, + SELF_ATTN_HEAD, + SELF_ATTN_KV_LENGTH, + SELF_ATTN_Q_LENGTH, +) +from maxdiffusion.configuration_utils import ConfigMixin, register_to_config +from maxdiffusion.models.attention_flax import NNXAttentionOp + + +ADALN_EMBED_DIM = 256 +SEQ_MULTI_OF = 32 + + +def _linear( + rngs, + in_features, + out_features, + *, + use_bias=True, + dtype=jnp.float32, + weights_dtype=jnp.float32, + kernel_axes=("embed", "heads"), + bias_axes=("heads",), +): + """Create a logically sharded NNX dense layer for the distributed denoiser.""" + return nnx.Linear( + in_features, + out_features, + use_bias=use_bias, + rngs=rngs, + dtype=dtype, + param_dtype=weights_dtype, + kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), kernel_axes), + bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), bias_axes), + ) + + +@flax.struct.dataclass +class ZImageTransformer2DModelOutput: + sample: list[jax.Array] + + +class ZImageTimestepEmbedder(nnx.Module): + """Sinusoidal timestep embedding followed by the upstream two-layer MLP.""" + + def __init__( + self, + rngs: nnx.Rngs, + out_size: int, + mid_size: Optional[int] = None, + frequency_embedding_size: int = 256, + dtype: jnp.dtype = jnp.float32, + weights_dtype: jnp.dtype = jnp.float32, + ): + self.frequency_embedding_size = frequency_embedding_size + mid_size = out_size if mid_size is None else mid_size + self.mlp_in = _linear( + rngs, + frequency_embedding_size, + mid_size, + dtype=dtype, + weights_dtype=weights_dtype, + kernel_axes=("embed", "mlp"), + bias_axes=("mlp",), + ) + self.mlp_out = _linear( + rngs, + mid_size, + out_size, + dtype=dtype, + weights_dtype=weights_dtype, + kernel_axes=("mlp", "embed"), + bias_axes=("embed",), + ) + + @staticmethod + def timestep_embedding(timestep: jax.Array, dim: int, max_period: float = 10000.0) -> jax.Array: + half = dim // 2 + freqs = jnp.exp(-math.log(max_period) * jnp.arange(half, dtype=jnp.float32) / half) + args = timestep[:, None].astype(jnp.float32) * freqs[None] + embedding = jnp.concatenate((jnp.cos(args), jnp.sin(args)), axis=-1) + if dim % 2: + embedding = jnp.concatenate((embedding, jnp.zeros_like(embedding[:, :1])), axis=-1) + return embedding + + def __call__(self, timestep: jax.Array) -> jax.Array: + x = self.timestep_embedding(timestep, self.frequency_embedding_size) + return self.mlp_out(nnx.silu(self.mlp_in(x))) + + +class ZImageFeedForward(nnx.Module): + """The bias-free SwiGLU MLP used by every Z-Image transformer block.""" + + def __init__( + self, + rngs: nnx.Rngs, + dim: int, + hidden_dim: int, + dtype: jnp.dtype = jnp.float32, + weights_dtype: jnp.dtype = jnp.float32, + ): + self.w1 = _linear( + rngs, dim, hidden_dim, use_bias=False, dtype=dtype, weights_dtype=weights_dtype, kernel_axes=("embed", "mlp") + ) + self.w2 = _linear( + rngs, hidden_dim, dim, use_bias=False, dtype=dtype, weights_dtype=weights_dtype, kernel_axes=("mlp", "embed") + ) + self.w3 = _linear( + rngs, dim, hidden_dim, use_bias=False, dtype=dtype, weights_dtype=weights_dtype, kernel_axes=("embed", "mlp") + ) + + def __call__(self, x: jax.Array) -> jax.Array: + return self.w2(nnx.silu(self.w1(x)) * self.w3(x)) + + +class ZImageAttention(nnx.Module): + """Z-Image self-attention backed by MaxDiffusion's shared attention operator.""" + + def __init__( + self, + rngs: nnx.Rngs, + dim: int, + heads: int, + qk_norm: bool = True, + eps: float = 1e-5, + attention_kernel: str = "dot_product", + mesh: Optional[jax.sharding.Mesh] = None, + flash_block_sizes: Optional[BlockSizes] = None, + flash_min_seq_length: int = 4096, + dtype: jnp.dtype = jnp.float32, + weights_dtype: jnp.dtype = jnp.float32, + ): + if dim % heads: + raise ValueError(f"dim ({dim}) must be divisible by heads ({heads}).") + self.heads = heads + self.dim_head = dim // heads + self.qk_norm = qk_norm + self.attention_kernel = attention_kernel + self.flash_min_seq_length = flash_min_seq_length + self.to_q = _linear(rngs, dim, dim, use_bias=False, dtype=dtype, weights_dtype=weights_dtype) + self.to_k = _linear(rngs, dim, dim, use_bias=False, dtype=dtype, weights_dtype=weights_dtype) + self.to_v = _linear(rngs, dim, dim, use_bias=False, dtype=dtype, weights_dtype=weights_dtype) + self.to_out = _linear( + rngs, dim, dim, use_bias=False, dtype=dtype, weights_dtype=weights_dtype, kernel_axes=("heads", "embed") + ) + if qk_norm: + self.norm_q = nnx.RMSNorm( + self.dim_head, epsilon=eps, use_scale=True, rngs=rngs, dtype=jnp.float32, param_dtype=jnp.float32 + ) + self.norm_k = nnx.RMSNorm( + self.dim_head, epsilon=eps, use_scale=True, rngs=rngs, dtype=jnp.float32, param_dtype=jnp.float32 + ) + self.attention_op = NNXAttentionOp( + mesh=mesh, + attention_kernel=attention_kernel, + scale=self.dim_head**-0.5, + heads=heads, + dim_head=self.dim_head, + split_head_dim=True, + float32_qk_product=False, + # Z-Image attends over one joint image+caption sequence, so every + # block is self-attention. Naming the axes as such is what lets the + # sequence-parallel / ring rules shard queries while keeping keys and + # values whole; the generic names would shard both and each device + # would attend only to its own slice. + axis_names_q=(BATCH, SELF_ATTN_HEAD, SELF_ATTN_Q_LENGTH, D_KV), + axis_names_kv=(BATCH, SELF_ATTN_HEAD, SELF_ATTN_KV_LENGTH, D_KV), + flash_min_seq_length=flash_min_seq_length, + flash_block_sizes=flash_block_sizes, + dtype=dtype, + ) + + @staticmethod + 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) + + def __call__( + self, hidden_states: jax.Array, freqs_cis: jax.Array, attention_mask: Optional[jax.Array] = None + ) -> jax.Array: + batch, length, _ = hidden_states.shape + q = self.to_q(hidden_states).reshape(batch, length, self.heads, self.dim_head) + k = self.to_k(hidden_states).reshape(batch, length, self.heads, self.dim_head) + v = self.to_v(hidden_states).reshape(batch, length, self.heads, self.dim_head) + if self.qk_norm: + q = self.norm_q(q) + k = self.norm_k(k) + q = self._apply_rope(q, freqs_cis, hidden_states.dtype) + k = self._apply_rope(k, freqs_cis, hidden_states.dtype) + # The local dot-product path is deliberately used for CPU parity tests. + # Production TPU execution uses NNXAttentionOp below (Flash/Splash/etc.). + if self.attention_kernel == "dot_product" or length < self.flash_min_seq_length: + scores = jnp.einsum("bqhd,bkhd->bhqk", q.astype(jnp.float32), k.astype(jnp.float32)) + scores = scores * (self.dim_head**-0.5) + if attention_mask is not None: + scores = jnp.where(attention_mask[:, None, None, :], scores, -jnp.inf) + output = jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1).astype(v.dtype), v).reshape(batch, length, -1) + else: + # NNXAttentionOp expects shape (batch, heads, sequence, dim_head) for 4D inputs. + # Since our local q, k, v are (batch, sequence, heads, dim_head), transposing them + # to (batch, heads, sequence, dim_head) bypasses any 3D flattening/unflattening. + q_4d = jnp.transpose(q, (0, 2, 1, 3)) + k_4d = jnp.transpose(k, (0, 2, 1, 3)) + v_4d = jnp.transpose(v, (0, 2, 1, 3)) + output = self.attention_op.apply_attention(q_4d, k_4d, v_4d, attention_mask=attention_mask) + return self.to_out(output) + + +def _select_per_token(noisy: jax.Array, clean: jax.Array, noise_mask: jax.Array) -> jax.Array: + return jnp.where(noise_mask[..., None] == 1, noisy[:, None, :], clean[:, None, :]) + + +class ZImageTransformerBlock(nnx.Module): + + def __init__( + self, + rngs: nnx.Rngs, + layer_id: int, + dim: int, + n_heads: int, + norm_eps: float, + qk_norm: bool, + modulation: bool = True, + attention_kernel: str = "dot_product", + mesh: Optional[jax.sharding.Mesh] = None, + flash_block_sizes: Optional[BlockSizes] = None, + flash_min_seq_length: int = 4096, + dtype: jnp.dtype = jnp.float32, + weights_dtype: jnp.dtype = jnp.float32, + ): + self.layer_id = layer_id + self.modulation = modulation + self.attention = ZImageAttention( + rngs, + dim, + n_heads, + qk_norm, + attention_kernel=attention_kernel, + mesh=mesh, + flash_block_sizes=flash_block_sizes, + flash_min_seq_length=flash_min_seq_length, + dtype=dtype, + weights_dtype=weights_dtype, + ) + self.feed_forward = ZImageFeedForward(rngs, dim, int(dim / 3 * 8), dtype=dtype, weights_dtype=weights_dtype) + norm_args = { + "epsilon": norm_eps, + "use_scale": True, + "rngs": rngs, + "dtype": jnp.float32, + "param_dtype": jnp.float32, + } + self.attention_norm1 = nnx.RMSNorm(dim, **norm_args) + self.ffn_norm1 = nnx.RMSNorm(dim, **norm_args) + self.attention_norm2 = nnx.RMSNorm(dim, **norm_args) + self.ffn_norm2 = nnx.RMSNorm(dim, **norm_args) + if modulation: + self.adaln_modulation = _linear( + rngs, + min(dim, ADALN_EMBED_DIM), + 4 * dim, + dtype=dtype, + weights_dtype=weights_dtype, + kernel_axes=("embed", "mlp"), + bias_axes=("mlp",), + ) + + def __call__( + self, + x: jax.Array, + freqs_cis: jax.Array, + attention_mask: Optional[jax.Array] = None, + adaln_input: Optional[jax.Array] = None, + noise_mask: Optional[jax.Array] = None, + adaln_noisy: Optional[jax.Array] = None, + adaln_clean: Optional[jax.Array] = None, + ) -> jax.Array: + if not self.modulation: + attn_out = self.attention(self.attention_norm1(x), freqs_cis, attention_mask) + x = x + self.attention_norm2(attn_out) + return x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x))) + + if noise_mask is None: + modulation = self.adaln_modulation(adaln_input) + scale_msa, gate_msa, scale_mlp, gate_mlp = jnp.split(modulation[:, None], 4, axis=-1) + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + gate_msa, gate_mlp = jnp.tanh(gate_msa), jnp.tanh(gate_mlp) + else: + noisy = self.adaln_modulation(adaln_noisy) + clean = self.adaln_modulation(adaln_clean) + noisy_parts = jnp.split(noisy, 4, axis=-1) + clean_parts = jnp.split(clean, 4, axis=-1) + scale_msa = 1.0 + _select_per_token(noisy_parts[0], clean_parts[0], noise_mask) + gate_msa = _select_per_token(jnp.tanh(noisy_parts[1]), jnp.tanh(clean_parts[1]), noise_mask) + scale_mlp = 1.0 + _select_per_token(noisy_parts[2], clean_parts[2], noise_mask) + gate_mlp = _select_per_token(jnp.tanh(noisy_parts[3]), jnp.tanh(clean_parts[3]), noise_mask) + + attn_out = self.attention(self.attention_norm1(x) * scale_msa, freqs_cis, attention_mask) + x = x + gate_msa * self.attention_norm2(attn_out) + return x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp)) + + +class ZImageFinalLayer(nnx.Module): + + def __init__( + self, + rngs: nnx.Rngs, + hidden_size: int, + out_channels: int, + dtype: jnp.dtype = jnp.float32, + weights_dtype: jnp.dtype = jnp.float32, + ): + self.norm_final = nnx.LayerNorm( + hidden_size, + epsilon=1e-6, + use_bias=False, + use_scale=False, + rngs=rngs, + dtype=jnp.float32, + param_dtype=jnp.float32, + ) + self.linear = _linear( + rngs, + hidden_size, + out_channels, + dtype=dtype, + weights_dtype=weights_dtype, + kernel_axes=("embed", "out_channels"), + bias_axes=("out_channels",), + ) + self.adaln_modulation = _linear( + rngs, + min(hidden_size, ADALN_EMBED_DIM), + hidden_size, + dtype=dtype, + weights_dtype=weights_dtype, + kernel_axes=("embed", "mlp"), + bias_axes=("mlp",), + ) + + def __call__( + self, + x: jax.Array, + c: Optional[jax.Array] = None, + noise_mask: Optional[jax.Array] = None, + c_noisy: Optional[jax.Array] = None, + c_clean: Optional[jax.Array] = None, + ) -> jax.Array: + if noise_mask is None: + scale = 1.0 + self.adaln_modulation(nnx.silu(c))[:, None] + else: + scale = 1.0 + _select_per_token( + self.adaln_modulation(nnx.silu(c_noisy)), self.adaln_modulation(nnx.silu(c_clean)), noise_mask + ) + return self.linear(self.norm_final(x) * scale) + + +class ZImageRopeEmbedder: + """CPU-precomputed multi-axis RoPE frequencies, identical to Diffusers' RopeEmbedder.""" + + def __init__( + self, theta: float = 256.0, axes_dims: Sequence[int] = (32, 48, 48), axes_lens: Sequence[int] = (1536, 512, 512) + ): + if len(axes_dims) != len(axes_lens): + raise ValueError("axes_dims and axes_lens must have the same length.") + self.axes_dims = tuple(axes_dims) + self.axes_lens = tuple(axes_lens) + freqs = [] + 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)) + self.freqs_cis = tuple(freqs) + + def __call__(self, ids: jax.Array) -> jax.Array: + if ids.ndim != 3 or ids.shape[-1] != len(self.axes_dims): + raise ValueError("ids must have shape (batch, sequence, number_of_axes).") + return jnp.concatenate([jnp.asarray(freq)[ids[..., axis]] for axis, freq in enumerate(self.freqs_cis)], axis=-1) + + +class ZImageTransformer2DModel(nnx.Module, ConfigMixin): + """Z-Image / Z-Image-Turbo denoiser. + + The public call signature mirrors Diffusers for the base text-to-image model. + `x` and `cap_feats` are lists so prompts and images can have individual + lengths/resolutions; this is important for the upstream checkpoint format. + """ + + config_name = "config.json" + + @register_to_config + def __init__( + self, + rngs: nnx.Rngs, + all_patch_size: Sequence[int] = (2,), + all_f_patch_size: Sequence[int] = (1,), + in_channels: int = 16, + dim: int = 3840, + n_layers: int = 30, + n_refiner_layers: int = 2, + n_heads: int = 30, + n_kv_heads: int = 30, + norm_eps: float = 1e-5, + qk_norm: bool = True, + cap_feat_dim: int = 2560, + rope_theta: float = 256.0, + t_scale: float = 1000.0, + axes_dims: Sequence[int] = (32, 48, 48), + axes_lens: Sequence[int] = (1536, 512, 512), + attention_kernel: str = "dot_product", + mesh: Optional[jax.sharding.Mesh] = None, + flash_block_sizes: Optional[BlockSizes] = None, + flash_min_seq_length: int = 4096, + dtype: jnp.dtype = jnp.float32, + weights_dtype: jnp.dtype = jnp.float32, + **unused_kwargs, + ): + del n_kv_heads, unused_kwargs + if tuple(all_patch_size) != (2,) or tuple(all_f_patch_size) != (1,): + raise NotImplementedError("MaxDiffusion currently supports Z-Image's released 2x2x1 patching only.") + if dim // n_heads != sum(axes_dims): + raise ValueError("Z-Image head_dim must equal sum(axes_dims).") + self.in_channels = in_channels + self.out_channels = in_channels + self.dim = dim + self.n_heads = n_heads + self.t_scale = t_scale + self.patch_size, self.f_patch_size = 2, 1 + patch_dim = self.patch_size * self.patch_size * self.f_patch_size * in_channels + self.x_embedder = _linear(rngs, patch_dim, dim, dtype=dtype, weights_dtype=weights_dtype) + self.final_layer = ZImageFinalLayer(rngs, dim, patch_dim, dtype=dtype, weights_dtype=weights_dtype) + self.t_embedder = ZImageTimestepEmbedder( + rngs, min(dim, ADALN_EMBED_DIM), mid_size=1024, dtype=dtype, weights_dtype=weights_dtype + ) + self.cap_embedder_norm = nnx.RMSNorm( + cap_feat_dim, epsilon=norm_eps, use_scale=True, rngs=rngs, dtype=jnp.float32, param_dtype=jnp.float32 + ) + self.cap_embedder = _linear(rngs, cap_feat_dim, dim, dtype=dtype, weights_dtype=weights_dtype) + self.x_pad_token = nnx.Param(jnp.zeros((1, dim), dtype=weights_dtype)) + self.cap_pad_token = nnx.Param(jnp.zeros((1, dim), dtype=weights_dtype)) + + block_args = { + "dim": dim, + "n_heads": n_heads, + "norm_eps": norm_eps, + "qk_norm": qk_norm, + "attention_kernel": attention_kernel, + "mesh": mesh, + "flash_block_sizes": flash_block_sizes, + "flash_min_seq_length": flash_min_seq_length, + "dtype": dtype, + "weights_dtype": weights_dtype, + } + self.noise_refiner = nnx.List( + [ZImageTransformerBlock(rngs, 1000 + index, modulation=True, **block_args) for index in range(n_refiner_layers)] + ) + self.context_refiner = nnx.List( + [ZImageTransformerBlock(rngs, index, modulation=False, **block_args) for index in range(n_refiner_layers)] + ) + self.layers = nnx.List([ZImageTransformerBlock(rngs, index, modulation=True, **block_args) for index in range(n_layers)]) + self.rope_embedder = ZImageRopeEmbedder(rope_theta, axes_dims, axes_lens) + + @staticmethod + def _patchify(image: jax.Array) -> tuple[jax.Array, tuple[int, int, int], tuple[int, int, int]]: + channels, frames, height, width = image.shape + if frames % 1 or height % 2 or width % 2: + raise ValueError("Z-Image latents must be divisible by the released 1x2x2 patch size.") + ft, ht, wt = frames, height // 2, width // 2 + patches = image.reshape(channels, ft, 1, ht, 2, wt, 2).transpose(1, 3, 5, 2, 4, 6, 0) + return patches.reshape(ft * ht * wt, -1), (frames, height, width), (ft, ht, wt) + + @staticmethod + def _coordinate_grid(size: tuple[int, int, int], start: tuple[int, int, int]) -> jax.Array: + return jnp.stack( + jnp.meshgrid(*[jnp.arange(s, s + n, dtype=jnp.int32) for s, n in zip(start, size)], indexing="ij"), axis=-1 + ).reshape(-1, 3) + + def _pad( + self, feature: jax.Array, grid_size: tuple[int, int, int], start: tuple[int, int, int] + ) -> tuple[jax.Array, jax.Array, jax.Array]: + original = feature.shape[0] + pad = (-original) % SEQ_MULTI_OF + positions = self._coordinate_grid(grid_size, start) + if pad: + feature = jnp.concatenate((feature, jnp.repeat(feature[-1:], pad, axis=0))) + positions = jnp.concatenate((positions, jnp.repeat(jnp.zeros((1, 3), dtype=jnp.int32), pad, axis=0))) + return feature, positions, jnp.arange(original + pad) >= original + + def _prepare( + self, features: list[jax.Array], positions: list[jax.Array], pad_masks: list[jax.Array], pad_token: jax.Array + ): + lengths = [feature.shape[0] for feature in features] + max_length = max(lengths) + padded_features, padded_freqs, attention_mask = [], [], [] + for feature, position, pad_mask in zip(features, positions, pad_masks): + feature = jnp.where(pad_mask[:, None], pad_token, feature) + frequency = self.rope_embedder(position[None])[0] + missing = max_length - feature.shape[0] + padded_features.append(jnp.pad(feature, ((0, missing), (0, 0)))) + padded_freqs.append(jnp.pad(frequency, ((0, missing), (0, 0)))) + attention_mask.append(jnp.arange(max_length) < feature.shape[0]) + mask = jnp.stack(attention_mask) + return ( + jnp.stack(padded_features), + jnp.stack(padded_freqs), + None if all(length == max_length for length in lengths) else mask, + lengths, + ) + + def __call__( + self, + x: list[jax.Array], + t: jax.Array, + cap_feats: list[jax.Array], + return_dict: bool = True, + patch_size: int = 2, + f_patch_size: int = 1, + **unused_kwargs, + ): + del unused_kwargs + if patch_size != 2 or f_patch_size != 1: + raise NotImplementedError("Only Z-Image's released 2x2x1 patching is supported.") + if len(x) != len(cap_feats): + raise ValueError("x and cap_feats must have one item per batch element.") + adaln = self.t_embedder(t * self.t_scale).astype(x[0].dtype) + + # 1. Pad and patchify each element locally (on its device) + caption_features, caption_positions, caption_masks = [], [], [] + image_features, image_positions, image_masks = [], [], [] + sizes = [] + + for image, caption in zip(x, cap_feats): + caption_padded, cap_position, cap_mask = self._pad(caption, (caption.shape[0], 1, 1), (1, 0, 0)) + patch, size, tokens = self._patchify(image) + patch_padded, image_position, image_mask = self._pad(patch, tokens, (caption.shape[0] + 1, 0, 0)) + + caption_features.append(caption_padded) + caption_positions.append(cap_position) + caption_masks.append(cap_mask) + image_features.append(patch_padded) + image_positions.append(image_position) + image_masks.append(image_mask) + sizes.append(size) + + # 2. Find maximum lengths statically (Python values) + image_lengths = [item.shape[0] for item in image_features] + caption_lengths = [item.shape[0] for item in caption_features] + max_image_len = max(image_lengths) + max_caption_len = max(caption_lengths) + + # 3. Pad to maximum sequence lengths statically + caption_features_padded, caption_positions_padded, caption_masks_padded = [], [], [] + for cap, pos, mask in zip(caption_features, caption_positions, caption_masks): + missing = max_caption_len - cap.shape[0] + if missing == 0: + caption_features_padded.append(cap) + caption_positions_padded.append(pos) + caption_masks_padded.append(mask) + else: + caption_features_padded.append(jnp.pad(cap, ((0, missing), (0, 0)))) + caption_positions_padded.append(jnp.pad(pos, ((0, missing), (0, 0)))) + caption_masks_padded.append(jnp.pad(mask, (0, missing), constant_values=True)) + + image_features_padded, image_positions_padded, image_masks_padded = [], [], [] + for patch, pos, mask in zip(image_features, image_positions, image_masks): + missing = max_image_len - patch.shape[0] + if missing == 0: + image_features_padded.append(patch) + image_positions_padded.append(pos) + image_masks_padded.append(mask) + else: + image_features_padded.append(jnp.pad(patch, ((0, missing), (0, 0)))) + image_positions_padded.append(jnp.pad(pos, ((0, missing), (0, 0)))) + image_masks_padded.append(jnp.pad(mask, (0, missing), constant_values=True)) + + # Stack into 3D batch arrays (maintains data-parallel sharding natively with zero communication) + caption_features_stacked = jnp.stack(caption_features_padded, axis=0) + caption_masks_stacked = jnp.stack(caption_masks_padded, axis=0) + caption_positions_stacked = jnp.stack(caption_positions_padded, axis=0) + + image_features_stacked = jnp.stack(image_features_padded, axis=0) + image_masks_stacked = jnp.stack(image_masks_padded, axis=0) + image_positions_stacked = jnp.stack(image_positions_padded, axis=0) + + # 4. Embed images and captions in parallel + image_embeddings = self.x_embedder(image_features_stacked) + image_embeddings = jnp.where(image_masks_stacked[..., None], self.x_pad_token[...], image_embeddings) + + captions = self.cap_embedder(self.cap_embedder_norm(caption_features_stacked)) + captions = jnp.where(caption_masks_stacked[..., None], self.cap_pad_token[...], captions) + + # 5. Rope frequencies + image_freqs_stacked = self.rope_embedder(image_positions_stacked) + caption_freqs_stacked = self.rope_embedder(caption_positions_stacked) + + # 6. Refiner blocks + has_image_padding = any(l < max_image_len for l in image_lengths) + image_attention_mask = None if not has_image_padding else ~image_masks_stacked + for block in self.noise_refiner: + image_embeddings = block(image_embeddings, image_freqs_stacked, image_attention_mask, adaln) + + has_caption_padding = any(l < max_caption_len for l in caption_lengths) + caption_attention_mask = None if not has_caption_padding else ~caption_masks_stacked + for block in self.context_refiner: + captions = block(captions, caption_freqs_stacked, caption_attention_mask) + + # 7. Joint sequence construction (no physical cross-device communication) + joint = jnp.concatenate((image_embeddings, captions), axis=1) + freqs_cis = jnp.concatenate((image_freqs_stacked, caption_freqs_stacked), axis=1) + + # 8. Model layers + joint_lengths = [img_l + cap_l for img_l, cap_l in zip(image_lengths, caption_lengths)] + joint_max = max_caption_len + max_image_len + has_joint_padding = any(l < joint_max for l in joint_lengths) + joint_attention_mask = ( + None if not has_joint_padding else jnp.concatenate((~image_masks_stacked, ~caption_masks_stacked), axis=1) + ) + + for block in self.layers: + joint = block(joint, freqs_cis, joint_attention_mask, adaln) + hidden_states = self.final_layer(joint, c=adaln) + + # 9. Unpatchify outputs + outputs = [] + for index, (frames, height, width) in enumerate(sizes): + tokens = frames * (height // 2) * (width // 2) + patches = hidden_states[index, :tokens].reshape(frames, height // 2, width // 2, 1, 2, 2, self.out_channels) + outputs.append(patches.transpose(6, 0, 3, 1, 4, 2, 5).reshape(self.out_channels, frames, height, width)) + return ZImageTransformer2DModelOutput(sample=outputs) if return_dict else (outputs,) diff --git a/src/maxdiffusion/models/z_image/z_image_utils.py b/src/maxdiffusion/models/z_image/z_image_utils.py new file mode 100644 index 000000000..00e6c9987 --- /dev/null +++ b/src/maxdiffusion/models/z_image/z_image_utils.py @@ -0,0 +1,124 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Checkpoint conversion helpers for Diffusers Z-Image checkpoints.""" + +import json +import os + +from flax.traverse_util import flatten_dict, unflatten_dict +import jax +import jax.numpy as jnp +from huggingface_hub import hf_hub_download +from safetensors import safe_open + + +_PREFIX_RENAMES = ( + ("all_x_embedder.2-1.", "x_embedder."), + ("all_final_layer.2-1.", "final_layer."), + ("t_embedder.mlp.0.", "t_embedder.mlp_in."), + ("t_embedder.mlp.2.", "t_embedder.mlp_out."), + ("mlp.0.", "mlp_in."), + ("mlp.2.", "mlp_out."), + ("cap_embedder.0.", "cap_embedder_norm."), + ("cap_embedder.1.", "cap_embedder."), + ("adaln_modulation.0.", "adaln_modulation."), + ("adaln_modulation.1.", "adaln_modulation."), + ("adaLN_modulation.0.", "adaln_modulation."), + ("adaLN_modulation.1.", "adaln_modulation."), +) + + +def z_image_pytorch_key_to_nnx_key(key: str) -> tuple[tuple[str | int, ...], bool]: + """Return the target NNX state path and whether a linear weight needs transpose.""" + for source, destination in _PREFIX_RENAMES: + if key.startswith(source): + key = destination + key[len(source) :] + break + key = key.replace(".adaln_modulation.1.", ".adaln_modulation.") + key = key.replace(".adaln_modulation.0.", ".adaln_modulation.") + key = key.replace(".adaLN_modulation.1.", ".adaln_modulation.") + key = key.replace(".adaLN_modulation.0.", ".adaln_modulation.") + key = key.replace(".attention.to_out.0.", ".attention.to_out.") + key = key.replace("attention.to_out.0.", "attention.to_out.") + key = key.replace(".weight", ".kernel") + key = key.replace(".norm_final.kernel", ".norm_final.scale") + key = key.replace(".attention_norm1.kernel", ".attention_norm1.scale") + key = key.replace(".attention_norm2.kernel", ".attention_norm2.scale") + key = key.replace(".ffn_norm1.kernel", ".ffn_norm1.scale") + key = key.replace(".ffn_norm2.kernel", ".ffn_norm2.scale") + key = key.replace(".norm_q.kernel", ".norm_q.scale") + key = key.replace(".norm_k.kernel", ".norm_k.scale") + key = key.replace("cap_embedder_norm.kernel", "cap_embedder_norm.scale") + for norm_name in ("attention_norm1", "attention_norm2", "ffn_norm1", "ffn_norm2", "norm_q", "norm_k"): + if key == f"{norm_name}.kernel": + key = f"{norm_name}.scale" + path = tuple(int(part) if part.isdigit() else part for part in key.split(".")) + # Every converted .kernel is a torch Linear weight. Norm scales and pad + # tokens deliberately bypass this path. + return path, path[-1] == "kernel" + + +def load_z_image_transformer( + pretrained_model_name_or_path: str, + eval_shapes: dict, + device: str = "cpu", + hf_download: bool = True, + subfolder: str = "transformer", + target_shardings: dict | None = None, +) -> dict: + """Load and convert a Diffusers Z-Image transformer into an NNX parameter tree. + + Shards are streamed one at a time. This avoids retaining the 12B PyTorch + state dict alongside the converted JAX tree on the host. + """ + index_name = "diffusion_pytorch_model.safetensors.index.json" + if os.path.isdir(pretrained_model_name_or_path): + index_path = os.path.join(pretrained_model_name_or_path, subfolder, index_name) + elif hf_download: + index_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=index_name) + else: + raise ValueError("A local model path is required when hf_download is False.") + with open(index_path, encoding="utf-8") as handle: + weight_map = json.load(handle)["weight_map"] + + expected = flatten_dict(eval_shapes) + converted = {} + cpu = jax.local_devices(backend=device)[0] + for filename in sorted(set(weight_map.values())): + path = ( + os.path.join(pretrained_model_name_or_path, subfolder, filename) + if os.path.isdir(pretrained_model_name_or_path) + else hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=filename) + ) + with safe_open(path, framework="pt", device="cpu") as tensors: + for source_key in tensors.keys(): + target_key, transpose = z_image_pytorch_key_to_nnx_key(source_key) + if target_key not in expected: + raise KeyError(f"Z-Image checkpoint key `{source_key}` maps to unknown NNX key `{target_key}`.") + value = tensors.get_tensor(source_key).float().numpy() + if transpose: + value = value.T + value = jnp.asarray(value, dtype=expected[target_key].dtype) + if value.shape != expected[target_key].shape: + raise ValueError(f"Shape mismatch for `{source_key}`: {value.shape} != {expected[target_key].shape}.") + target = target_shardings.get(target_key) if target_shardings is not None else cpu + converted[target_key] = jax.device_put(value, target) + + missing = set(expected) - set(converted) + if missing: + raise ValueError(f"Z-Image checkpoint is missing {len(missing)} parameters, e.g. {sorted(missing)[:5]}.") + return unflatten_dict(converted) diff --git a/src/maxdiffusion/pipelines/z_image/__init__.py b/src/maxdiffusion/pipelines/z_image/__init__.py new file mode 100644 index 000000000..afe63827d --- /dev/null +++ b/src/maxdiffusion/pipelines/z_image/__init__.py @@ -0,0 +1,19 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Z-Image inference pipeline.""" + +from .z_image_pipeline import ZImagePipeline diff --git a/src/maxdiffusion/pipelines/z_image/z_image_pipeline.py b/src/maxdiffusion/pipelines/z_image/z_image_pipeline.py new file mode 100644 index 000000000..519682716 --- /dev/null +++ b/src/maxdiffusion/pipelines/z_image/z_image_pipeline.py @@ -0,0 +1,245 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Inference pipeline for the official Z-Image and Z-Image-Turbo checkpoints.""" + +from contextlib import nullcontext +from functools import partial +import time +from typing import Optional, Union + +from flax import nnx +from flax.linen import partitioning as nn_partitioning +import jax +import jax.numpy as jnp +import numpy as np +from PIL import Image + + +@jax.jit +def encode_step(graphdef, state, rest, input_ids, attention_mask): + """Qwen3 forward, returning the layer Z-Image conditions on.""" + text_encoder = nnx.merge(graphdef, state, rest) + _, all_hidden_states = text_encoder(input_ids=input_ids, attention_mask=attention_mask) + # `all_hidden_states` is [embeddings, layer_1, ..., layer_n], so -2 is the + # penultimate layer's raw output -- the same tensor Transformers hands back + # as `hidden_states[-2]`. + return all_hidden_states[-2] + + +@partial(jax.jit, static_argnames=("model_dtype",)) +def denoise_step(graphdef, state, rest, latents, timestep, prompt_embeds, sigma_delta, model_dtype): + """One flow-matching Euler step, compiled once and reused every step. + + Latents carry float32 so the Euler update accumulates in full precision; + only the denoiser's input is cast down. + """ + transformer = nnx.merge(graphdef, state, rest) + # Z-Image takes one array per batch element so prompts of different lengths + # can share a step; the frame axis it expects is the `None` below. + samples = [sample[:, None] for sample in latents.astype(model_dtype)] + predictions = transformer(samples, timestep, prompt_embeds, return_dict=False)[0] + prediction = jnp.stack([item[:, 0] for item in predictions]) + return latents + sigma_delta * (-prediction.astype(jnp.float32)) + + +class ZImagePipeline: + """A small, composable JAX denoising pipeline. + + Every component -- the Qwen3 text encoder, the VAE and the 6B denoiser -- + runs in JAX/MaxDiffusion. Loading lives in + `maxdiffusion.checkpointing.z_image_checkpointer.ZImageCheckpointer`. + """ + + def __init__( + self, + transformer, + vae, + vae_params, + tokenizer, + text_encoder, + dtype=jnp.bfloat16, + mesh=None, + logical_axis_rules=(), + vae_shift_factor: Optional[float] = None, + offload_encoders: bool = False, + ): + self.transformer = transformer + self.vae = vae + self.vae_params = vae_params + self.tokenizer = tokenizer + self.text_encoder = text_encoder + self.dtype = dtype + self.mesh = mesh + # The attention op constrains its activations by logical axis name, so the + # rules must be in scope while the denoise step is traced. + self.logical_axis_rules = logical_axis_rules + self.vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1) + # Z-Image's Flux-format VAE config carries this field, while the older + # Flax AutoencoderKL config loader drops unknown attributes, so the + # checkpointer reads it from the raw config and passes it in. + self.vae_shift_factor = vae_shift_factor if vae_shift_factor is not None else 0.1159 + self._jitted_decoder = None + # Pre-split models at initialization to avoid any runtime CPU/metadata overhead + self._text_encoder_split = nnx.split(self.text_encoder, nnx.Param, ...) + self._transformer_split = nnx.split(self.transformer, nnx.Param, ...) + + self.offload_encoders = offload_encoders + if self.offload_encoders: + # Save the original shardings of the text encoder parameters + graphdef, state, rest = self._text_encoder_split + self._text_encoder_shardings = jax.tree_util.tree_map(lambda x: x.sharding, state) + + def _decode_fn(self): + """Jitted VAE decode, including the latent denormalization.""" + if self._jitted_decoder is None: + + def decode(params, latents): + return self.vae.apply( + {"params": params}, + latents / self.vae.config.scaling_factor + self.vae_shift_factor, + deterministic=True, + method=self.vae.decode, + ).sample + + self._jitted_decoder = jax.jit(decode) + return self._jitted_decoder + + def encode_prompt(self, prompt: Union[str, list[str]], max_sequence_length: int = 512) -> list[jax.Array]: + prompts = [prompt] if isinstance(prompt, str) else list(prompt) + chats = [ + self.tokenizer.apply_chat_template( + [{"role": "user", "content": item}], tokenize=False, add_generation_prompt=True, enable_thinking=True + ) + for item in prompts + ] + # Padding to a fixed length keeps one compiled encoder for every prompt. + # Qwen3 is causal and the pad mask is applied, so right padding cannot + # change the real tokens' states -- they are sliced back out below, and + # Z-Image consumes only those unmasked states. + inputs = self.tokenizer( + chats, padding="max_length", max_length=max_sequence_length, truncation=True, return_tensors="np" + ) + graphdef, state, rest = self._text_encoder_split + + 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) + + with self.mesh if self.mesh is not None else nullcontext(): + embeddings = encode_step( + graphdef, + state, + rest, + jnp.asarray(inputs["input_ids"], jnp.int32), + jnp.asarray(inputs["attention_mask"], jnp.int32), + ) + lengths = np.asarray(inputs["attention_mask"]).sum(axis=-1) + + if self.offload_encoders: + from maxdiffusion import max_logging + + s_offload = time.perf_counter() + cpus = jax.devices("cpu") + offloaded_state = jax.tree_util.tree_map(lambda x: jax.device_put(x, device=cpus[0]), state) + 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))] + + @staticmethod + def _sigmas(num_inference_steps: int, shift: float) -> jax.Array: + sigmas = jnp.linspace(1.0, 0.0, num_inference_steps + 1, dtype=jnp.float32) + return shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + + def __call__( + self, + prompt: Union[str, list[str]], + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 9, + guidance_scale: float = 0.0, + seed: int = 0, + max_sequence_length: int = 512, + vae_decode_chunk: int = 1, + output_type: str = "pil", + return_timings: bool = False, + ): + del guidance_scale # The published Turbo configuration uses guidance_scale=0. + vae_scale = self.vae_scale_factor * 2 + if height % vae_scale or width % vae_scale: + raise ValueError(f"height and width must be divisible by {vae_scale}.") + trace = {} + stage_start = time.perf_counter() + + def stage(name, value): + """Close out a timed stage, if the caller asked for timings. + + Dispatch is async, so a stage's time is only its own if we block at the + boundary. That barrier costs a host-device round trip and stops the host + running ahead to enqueue the next stage, so it is skipped entirely when + no trace was requested. + """ + nonlocal stage_start + if return_timings: + jax.block_until_ready(value) + trace[name] = time.perf_counter() - stage_start + stage_start = time.perf_counter() + return value + + prompt_embeds = stage("text_encode", self.encode_prompt(prompt, max_sequence_length)) + batch_size = len(prompt_embeds) + latent_height, latent_width = 2 * (height // vae_scale), 2 * (width // vae_scale) + latents = jax.random.normal( + jax.random.key(seed), (batch_size, self.transformer.in_channels, latent_height, latent_width), self.dtype + ).astype(jnp.float32) + # The released scheduler is FlowMatchEulerDiscreteScheduler(shift=3.0). + sigmas = self._sigmas(num_inference_steps, shift=3.0) + # Split once, outside the loop: every step reuses the same compiled + # executable, so the step count only changes the Python trip count. + graphdef, state, rest = self._transformer_split + with self.mesh if self.mesh is not None else nullcontext(), nn_partitioning.axis_rules(self.logical_axis_rules): + for index in range(num_inference_steps): + # Timestep and step size stay traced arguments; making them static + # would compile a separate executable for every step. + latents = denoise_step( + graphdef, + state, + rest, + latents, + jnp.full((batch_size,), 1.0 - sigmas[index], dtype=self.dtype), + prompt_embeds, + (sigmas[index + 1] - sigmas[index]).astype(jnp.float32), + self.dtype, + ) + latents = stage("denoise", latents) + + # The decoder is replicated and its activations are full-resolution, so + # a whole batch at once needs many GB. Decode in chunks; raise + # vae_decode_chunk to trade memory for speed. + chunk = max(1, int(vae_decode_chunk)) + decoded = jnp.concatenate( + [ + self._decode_fn()(self.vae_params, latents[start : start + chunk].astype(self.dtype)) + for start in range(0, batch_size, chunk) + ] + ) + decoded = stage("vae_decode", decoded) + # The conversion to numpy blocks on its own, so this stage is always real. + frames = np.asarray((decoded.transpose(0, 2, 3, 1) / 2 + 0.5).clip(0, 1) * 255, dtype=np.uint8) + images = [Image.fromarray(frame) for frame in frames] if output_type == "pil" else list(frames) + stage("host_post", None) + return (images, trace) if return_timings else images diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index 793b8979b..e84571165 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/src/maxdiffusion/pyconfig.py @@ -35,13 +35,14 @@ WAN2_2, LTX2_VIDEO, LTX2_3, + Z_IMAGE, RING_ATTENTION_AXIS_RULES, SEQUENCE_PARALLEL_AXIS_RULES, ULYSSES_ATTENTION_AXIS_RULES, ULYSSES_RING_ATTENTION_AXIS_RULES, ) -_ALLOWED_MODEL_NAMES = {WAN2_1, WAN2_2, LTX2_VIDEO, LTX2_3} +_ALLOWED_MODEL_NAMES = {WAN2_1, WAN2_2, LTX2_VIDEO, LTX2_3, Z_IMAGE} _ALLOWED_TRAINING_MODEL_NAMES = {WAN2_1} @@ -203,6 +204,8 @@ def user_init(raw_keys): raw_keys["names_which_can_be_saved"] = [] if "names_which_can_be_offloaded" not in raw_keys: raw_keys["names_which_can_be_offloaded"] = [] + if "offload_encoders" not in raw_keys: + raw_keys["offload_encoders"] = False raw_keys["weights_dtype"] = jax.numpy.dtype(raw_keys["weights_dtype"]) raw_keys["activations_dtype"] = jax.numpy.dtype(raw_keys["activations_dtype"]) @@ -354,7 +357,5 @@ def initialize(argv, **kwargs): if __name__ == "__main__": initialize(sys.argv) - from maxdiffusion import max_logging - - max_logging.log(config.steps) + print(config.steps) r = range(config.steps) diff --git a/src/maxdiffusion/tests/z_image/z_image_module_parity_test.py b/src/maxdiffusion/tests/z_image/z_image_module_parity_test.py new file mode 100644 index 000000000..de3c6bfc8 --- /dev/null +++ b/src/maxdiffusion/tests/z_image/z_image_module_parity_test.py @@ -0,0 +1,155 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Parity coverage for every Z-Image transformer layer against Diffusers PyTorch.""" + +import os +import unittest + +os.environ["JAX_PLATFORMS"] = "cpu" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +import torch +from flax import nnx + +from diffusers.models.transformers.transformer_z_image import ( + FeedForward as HFFeedForward, + FinalLayer as HFFinalLayer, + RopeEmbedder as HFRopeEmbedder, + TimestepEmbedder as HFTimestepEmbedder, + ZImageTransformer2DModel as HFZImageTransformer2DModel, + ZImageTransformerBlock as HFZImageTransformerBlock, +) +from maxdiffusion.models.z_image.transformer_z_image import ( + ZImageFeedForward, + ZImageFinalLayer, + ZImageRopeEmbedder, + ZImageTimestepEmbedder, + ZImageTransformer2DModel, + ZImageTransformerBlock, +) +from maxdiffusion.models.z_image.z_image_utils import z_image_pytorch_key_to_nnx_key + + +def to_numpy(value): + if isinstance(value, torch.Tensor): + if value.dtype == torch.bfloat16: + value = value.float() + return value.detach().cpu().numpy() + return np.asarray(value) + + +def assert_close(test_case, actual, expected, atol=2e-5, rtol=2e-5): + test_case.assertEqual(to_numpy(actual).shape, to_numpy(expected).shape) + np.testing.assert_allclose(to_numpy(actual), to_numpy(expected), atol=atol, rtol=rtol) + + +def copy_parameters(local_module, torch_module): + _, state, rest = nnx.split(local_module, nnx.Param, ...) + flat_state = dict(nnx.to_flat_state(state)) + mapped = set() + for source_key, tensor in torch_module.state_dict().items(): + target_key, transpose = z_image_pytorch_key_to_nnx_key(source_key) + if target_key not in flat_state: + continue + value = to_numpy(tensor) + if transpose: + value = value.T + flat_state[target_key][...] = jnp.asarray(value) + mapped.add(target_key) + missing = set(flat_state) - mapped + if missing: + raise AssertionError(f"Unmapped NNX parameters: {sorted(missing)}") + return nnx.merge(nnx.graphdef(local_module), nnx.from_flat_state(flat_state), rest) + + +@pytest.mark.skipif(os.getenv("GITHUB_ACTIONS") == "true", reason="PyTorch parity tests are not run in GitHub Actions") +class ZImageModuleParityTest(unittest.TestCase): + + def setUp(self): + torch.manual_seed(0) + self.rngs = nnx.Rngs(jax.random.key(0)) + + def test_timestep_embedder_parity(self): + hf = HFTimestepEmbedder(8, mid_size=12, frequency_embedding_size=10).eval() + local = copy_parameters(ZImageTimestepEmbedder(self.rngs, 8, 12, 10), hf) + timestep = torch.tensor([0.1, 0.7]) + assert_close(self, local(jnp.asarray(to_numpy(timestep))), hf(timestep)) + + def test_feed_forward_parity(self): + hf = HFFeedForward(8, 20).eval() + local = copy_parameters(ZImageFeedForward(self.rngs, 8, 20), hf) + inputs = torch.randn(2, 7, 8) + assert_close(self, local(jnp.asarray(to_numpy(inputs))), hf(inputs)) + + def test_final_layer_parity(self): + hf = HFFinalLayer(8, 6).eval() + local = copy_parameters(ZImageFinalLayer(self.rngs, 8, 6), hf) + inputs, conditioning = torch.randn(2, 7, 8), torch.randn(2, 8) + assert_close(self, local(jnp.asarray(to_numpy(inputs)), jnp.asarray(to_numpy(conditioning))), hf(inputs, conditioning)) + + def test_rope_embedder_parity(self): + hf = HFRopeEmbedder(theta=256.0, axes_dims=[2, 2, 4], axes_lens=[64, 64, 64]) + local = ZImageRopeEmbedder(theta=256.0, axes_dims=[2, 2, 4], axes_lens=[64, 64, 64]) + ids = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) + assert_close(self, local(jnp.asarray(ids.numpy())[None])[0], hf(ids), atol=1e-6, rtol=1e-6) + + def test_transformer_block_parity(self): + hf = HFZImageTransformerBlock(0, 32, 4, 4, 1e-5, True, modulation=True).eval() + local = copy_parameters(ZImageTransformerBlock(self.rngs, 0, 32, 4, 1e-5, True, attention_kernel="dot_product"), hf) + inputs = torch.randn(2, 32, 32) + ids = torch.randint(0, 16, (2, 32, 3), dtype=torch.int32) + rope = HFRopeEmbedder(axes_dims=[2, 2, 4], axes_lens=[32, 32, 32]) + freqs = torch.stack([rope(ids[index]) for index in range(ids.shape[0])]) + conditioning = torch.randn(2, 32) + assert_close( + self, + local(jnp.asarray(to_numpy(inputs)), jnp.asarray(to_numpy(freqs)), None, jnp.asarray(to_numpy(conditioning))), + hf(inputs, None, freqs, conditioning), + atol=3e-5, + rtol=3e-5, + ) + + def test_full_transformer_parity(self): + config = { + "all_patch_size": (2,), + "all_f_patch_size": (1,), + "in_channels": 4, + "dim": 32, + "n_layers": 1, + "n_refiner_layers": 1, + "n_heads": 4, + "n_kv_heads": 4, + "cap_feat_dim": 8, + "axes_dims": [2, 2, 4], + "axes_lens": [64, 64, 64], + } + hf = HFZImageTransformer2DModel(**config).eval() + local = copy_parameters(ZImageTransformer2DModel(rngs=self.rngs, attention_kernel="dot_product", **config), hf) + images = [torch.randn(4, 1, 4, 4)] + captions = [torch.randn(32, 8)] + timestep = torch.tensor([0.3]) + expected = hf(images, timestep, captions, return_dict=False)[0][0] + actual = local( + [jnp.asarray(to_numpy(images[0]))], + jnp.asarray(to_numpy(timestep)), + [jnp.asarray(to_numpy(captions[0]))], + return_dict=False, + )[0][0] + assert_close(self, actual, expected, atol=5e-5, rtol=5e-5) diff --git a/src/maxdiffusion/tests/z_image/z_image_transformer_test.py b/src/maxdiffusion/tests/z_image/z_image_transformer_test.py new file mode 100644 index 000000000..56c864eaf --- /dev/null +++ b/src/maxdiffusion/tests/z_image/z_image_transformer_test.py @@ -0,0 +1,50 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Fast structural tests for the JAX Z-Image transformer.""" + +import jax +import jax.numpy as jnp +from flax import nnx + +from maxdiffusion.models.z_image.transformer_z_image import ZImageTransformer2DModel + + +def _tiny_model(): + return ZImageTransformer2DModel( + rngs=nnx.Rngs(jax.random.key(0)), + in_channels=4, + dim=32, + n_layers=1, + n_refiner_layers=1, + n_heads=4, + n_kv_heads=4, + cap_feat_dim=8, + axes_dims=(2, 2, 4), + axes_lens=(64, 64, 64), + attention_kernel="dot_product", + ) + + +def test_variable_prompt_and_image_lengths(): + model = _tiny_model() + output = model( + [jnp.ones((4, 1, 4, 4)), jnp.ones((4, 1, 4, 6))], + jnp.array([0.1, 0.9]), + [jnp.ones((5, 8)), jnp.ones((17, 8))], + ).sample + assert output[0].shape == (4, 1, 4, 4) + assert output[1].shape == (4, 1, 4, 6)