diff --git a/src/maxdiffusion/configs/base_zimage.yml b/src/maxdiffusion/configs/base_zimage.yml new file mode 100644 index 000000000..01d60e1cd --- /dev/null +++ b/src/maxdiffusion/configs/base_zimage.yml @@ -0,0 +1,39 @@ +# Official Z-Image inference defaults. Z-Image and Turbo share the denoiser. +pretrained_model_name_or_path: "Tongyi-MAI/Z-Image" +prompt: "A red fox reading a book in a quiet library, cinematic light" +resolution: 1024 +num_inference_steps: 50 +seed: 42 +output_file: "/tmp/zimage.png" +weights_dtype: "bfloat16" +activations_dtype: "bfloat16" +attention: "flash" +flash_block_sizes: {} +hardware: "tpu" +skip_jax_distributed_system: true +run_name: "zimage" +output_dir: "/tmp/maxdiffusion" +save_config_to_gcs: false +jax_cache_dir: "" +unet_checkpoint: "" +dataset_name: "" +dataset_save_location: "/tmp" +per_device_batch_size: 1 +learning_rate_schedule_steps: 1 +max_train_steps: 1 +compile_topology_num_slices: -1 +quantization_local_shard_count: -1 +use_qwix_quantization: false +attention_sharding_uniform: true +logical_axis_rules: [["batch", "data"], ["activation_length", "context"], ["activation_heads", "tensor"], ["heads", "tensor"], ["mlp", "tensor"], ["embed", "fsdp"], ["norm", "tensor"], ["out_channels", "tensor"]] +data_sharding: [["data", "fsdp", "context", "tensor"]] +mesh_axes: ["data", "fsdp", "context", "tensor"] +dcn_data_parallelism: 1 +dcn_fsdp_parallelism: 1 +dcn_context_parallelism: 1 +dcn_tensor_parallelism: 1 +ici_data_parallelism: 1 +ici_fsdp_parallelism: 1 +ici_context_parallelism: 1 +ici_tensor_parallelism: -1 +allow_split_physical_axes: 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..89bb648f5 --- /dev/null +++ b/src/maxdiffusion/configs/base_zimage_turbo.yml @@ -0,0 +1,40 @@ +# Official Z-Image-Turbo inference defaults. The base model uses the same +# transformer; only its recommended denoising-step count differs. +pretrained_model_name_or_path: "Tongyi-MAI/Z-Image-Turbo" +prompt: "A red fox reading a book in a quiet library, cinematic light" +resolution: 1024 +num_inference_steps: 9 +seed: 42 +output_file: "/tmp/zimage-turbo.png" +weights_dtype: "bfloat16" +activations_dtype: "bfloat16" +attention: "flash" +flash_block_sizes: {} +hardware: "tpu" +skip_jax_distributed_system: true +run_name: "zimage-turbo" +output_dir: "/tmp/maxdiffusion" +save_config_to_gcs: false +jax_cache_dir: "" +unet_checkpoint: "" +dataset_name: "" +dataset_save_location: "/tmp" +per_device_batch_size: 1 +learning_rate_schedule_steps: 1 +max_train_steps: 1 +compile_topology_num_slices: -1 +quantization_local_shard_count: -1 +use_qwix_quantization: false +attention_sharding_uniform: true +logical_axis_rules: [["batch", "data"], ["activation_length", "context"], ["activation_heads", "tensor"], ["heads", "tensor"], ["mlp", "tensor"], ["embed", "fsdp"], ["norm", "tensor"], ["out_channels", "tensor"]] +data_sharding: [["data", "fsdp", "context", "tensor"]] +mesh_axes: ["data", "fsdp", "context", "tensor"] +dcn_data_parallelism: 1 +dcn_fsdp_parallelism: 1 +dcn_context_parallelism: 1 +dcn_tensor_parallelism: 1 +ici_data_parallelism: 1 +ici_fsdp_parallelism: 1 +ici_context_parallelism: 1 +ici_tensor_parallelism: -1 +allow_split_physical_axes: false diff --git a/src/maxdiffusion/generate_zimage.py b/src/maxdiffusion/generate_zimage.py new file mode 100644 index 000000000..5a634d118 --- /dev/null +++ b/src/maxdiffusion/generate_zimage.py @@ -0,0 +1,37 @@ +"""Generate an image with Z-Image or Z-Image-Turbo.""" + +from absl import app +import jax + +from maxdiffusion import pyconfig +from maxdiffusion.max_utils import create_device_mesh, get_flash_block_sizes +from maxdiffusion.pipelines.z_image import ZImagePipeline +from jax.sharding import Mesh + + +def main(argv): + pyconfig.initialize(argv) + config = pyconfig.config + mesh = Mesh(create_device_mesh(config), config.mesh_axes) + pipeline = ZImagePipeline.from_pretrained( + config.pretrained_model_name_or_path, + jax.random.key(config.seed), + attention_kernel=config.attention, + mesh=mesh, + flash_block_sizes=get_flash_block_sizes(config), + dtype=config.activations_dtype, + weights_dtype=config.weights_dtype, + logical_axis_rules=config.logical_axis_rules, + ) + image = pipeline( + config.prompt, + height=config.resolution, + width=config.resolution, + num_inference_steps=config.num_inference_steps, + seed=config.seed, + ) + image.save(config.output_file) + + +if __name__ == "__main__": + app.run(main) 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/z_image/__init__.py b/src/maxdiffusion/models/z_image/__init__.py new file mode 100644 index 000000000..0e759aeb7 --- /dev/null +++ b/src/maxdiffusion/models/z_image/__init__.py @@ -0,0 +1,3 @@ +"""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..cdb887141 --- /dev/null +++ b/src/maxdiffusion/models/z_image/transformer_z_image.py @@ -0,0 +1,578 @@ +"""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 BlockSizes +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, + 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) -> 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(x.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).astype(hidden_states.dtype) + k = self.norm_k(k).astype(hidden_states.dtype) + q = self._apply_rope(q, freqs_cis) + k = self._apply_rope(k, freqs_cis) + # 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: + q, k, v = (item.reshape(batch, length, -1) for item in (q, k, v)) + output = self.attention_op.apply_attention(q, k, v, 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 = dict(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 = dict( + 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) + image_features, image_positions, image_masks, sizes = [], [], [], [] + caption_features, caption_positions, caption_masks = [], [], [] + for image, caption in zip(x, cap_feats): + caption, cap_position, cap_mask = self._pad(caption, (caption.shape[0], 1, 1), (1, 0, 0)) + patch, size, tokens = self._patchify(image) + patch, image_position, image_mask = self._pad(patch, tokens, (caption.shape[0] + 1, 0, 0)) + image_features.append(patch) + image_positions.append(image_position) + image_masks.append(image_mask) + caption_features.append(caption) + caption_positions.append(cap_position) + caption_masks.append(cap_mask) + sizes.append(size) + + image_lengths = [item.shape[0] for item in image_features] + image_embeddings = self.x_embedder(jnp.concatenate(image_features, axis=0)) + image_embeddings = list(jnp.split(image_embeddings, jnp.cumsum(jnp.asarray(image_lengths[:-1])))) + image_embeddings, image_freqs, image_attention_mask, image_lengths = self._prepare( + image_embeddings, image_positions, image_masks, self.x_pad_token[...] + ) + for block in self.noise_refiner: + image_embeddings = block(image_embeddings, image_freqs, image_attention_mask, adaln) + + caption_lengths = [item.shape[0] for item in caption_features] + captions = self.cap_embedder(self.cap_embedder_norm(jnp.concatenate(caption_features, axis=0))) + captions = list(jnp.split(captions, jnp.cumsum(jnp.asarray(caption_lengths[:-1])))) + captions, caption_freqs, caption_attention_mask, caption_lengths = self._prepare( + captions, caption_positions, caption_masks, self.cap_pad_token[...] + ) + for block in self.context_refiner: + captions = block(captions, caption_freqs, caption_attention_mask) + + joint_lengths = [image_length + caption_length for image_length, caption_length in zip(image_lengths, caption_lengths)] + joint_max = max(joint_lengths) + joint, joint_freqs, joint_mask = [], [], [] + for index, length in enumerate(joint_lengths): + value = jnp.concatenate((image_embeddings[index, : image_lengths[index]], captions[index, : caption_lengths[index]])) + freqs = jnp.concatenate((image_freqs[index, : image_lengths[index]], caption_freqs[index, : caption_lengths[index]])) + joint.append(jnp.pad(value, ((0, joint_max - length), (0, 0)))) + joint_freqs.append(jnp.pad(freqs, ((0, joint_max - length), (0, 0)))) + joint_mask.append(jnp.arange(joint_max) < length) + hidden_states, freqs_cis = jnp.stack(joint), jnp.stack(joint_freqs) + joint_attention_mask = None if all(length == joint_max for length in joint_lengths) else jnp.stack(joint_mask) + for block in self.layers: + hidden_states = block(hidden_states, freqs_cis, joint_attention_mask, adaln) + hidden_states = self.final_layer(hidden_states, c=adaln) + + 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..9c926df81 --- /dev/null +++ b/src/maxdiffusion/models/z_image/z_image_utils.py @@ -0,0 +1,108 @@ +"""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..4725d66af --- /dev/null +++ b/src/maxdiffusion/pipelines/z_image/__init__.py @@ -0,0 +1,3 @@ +"""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..3431fe0ec --- /dev/null +++ b/src/maxdiffusion/pipelines/z_image/z_image_pipeline.py @@ -0,0 +1,164 @@ +"""Inference pipeline for the official Z-Image and Z-Image-Turbo checkpoints.""" + +from contextlib import nullcontext +from typing import Optional + +from flax import nnx +import flax.linen as nn +from flax.traverse_util import flatten_dict +import jax +import jax.numpy as jnp +from jax.sharding import NamedSharding, PartitionSpec as P +import numpy as np +from PIL import Image +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from maxdiffusion import FlaxAutoencoderKL +from maxdiffusion.models.z_image.transformer_z_image import ZImageTransformer2DModel +from maxdiffusion.models.z_image.z_image_utils import load_z_image_transformer + + +def create_z_image_transformer( + model_id: str, + rngs: nnx.Rngs, + attention_kernel: str = "dot_product", + mesh: Optional[jax.sharding.Mesh] = None, + flash_block_sizes=None, + dtype: jnp.dtype = jnp.bfloat16, + weights_dtype: jnp.dtype = jnp.bfloat16, + logical_axis_rules=None, +) -> ZImageTransformer2DModel: + """Instantiate and stream a Diffusers Z-Image checkpoint into an NNX model.""" + config = ZImageTransformer2DModel.load_config(model_id, subfolder="transformer") + + def factory(init_rngs): + return ZImageTransformer2DModel( + rngs=init_rngs, + attention_kernel=attention_kernel, + mesh=mesh, + flash_block_sizes=flash_block_sizes, + dtype=dtype, + weights_dtype=weights_dtype, + **config, + ) + + model_shape = nnx.eval_shape(factory, rngs) + graphdef, state, rest = nnx.split(model_shape, nnx.Param, ...) + target_shardings = None + if mesh is not None and logical_axis_rules is not None: + logical_specs = nnx.get_partition_spec(state) + sharding_tree = nn.logical_to_mesh_sharding(logical_specs, mesh, logical_axis_rules) + target_shardings = {path: variable.value for path, variable in nnx.to_flat_state(sharding_tree)} + loaded_params = load_z_image_transformer(model_id, state.to_pure_dict(), target_shardings=target_shardings) + flat_state = dict(nnx.to_flat_state(state)) + for path, value in flatten_dict(loaded_params).items(): + flat_state[path].value = value + return nnx.merge(graphdef, nnx.from_flat_state(flat_state), rest) + + +class ZImagePipeline: + """A small, composable JAX denoising pipeline. + + Text encoding stays in PyTorch because Transformers has no Flax Qwen3 model; + the VAE and the 12B denoiser execute in JAX/MaxDiffusion. + """ + + def __init__(self, transformer, vae, vae_params, tokenizer, text_encoder, dtype=jnp.bfloat16, mesh=None): + 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 + 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. + self.vae_shift_factor = getattr(vae.config, "shift_factor", 0.1159) + + @classmethod + def from_pretrained( + cls, + model_id: str, + rng: jax.Array, + attention_kernel: str = "dot_product", + mesh: Optional[jax.sharding.Mesh] = None, + flash_block_sizes=None, + dtype: jnp.dtype = jnp.bfloat16, + weights_dtype: jnp.dtype = jnp.bfloat16, + text_device: Optional[str] = None, + logical_axis_rules=None, + ): + transformer = create_z_image_transformer( + model_id, nnx.Rngs(rng), attention_kernel, mesh, flash_block_sizes, dtype, weights_dtype, logical_axis_rules + ) + vae, vae_params = FlaxAutoencoderKL.from_pretrained( + model_id, subfolder="vae", from_pt=True, use_safetensors=True, dtype=weights_dtype + ) + if mesh is not None: + replicated = NamedSharding(mesh, P()) + vae_params = jax.tree.map(lambda value: jax.device_put(value, replicated), vae_params) + tokenizer = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", use_fast=True) + text_encoder = AutoModelForCausalLM.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16) + text_device = text_device or ("cuda" if torch.cuda.is_available() else "cpu") + text_encoder.to(text_device).eval() + return cls(transformer, vae, vae_params, tokenizer, text_encoder, dtype, mesh) + + def encode_prompt(self, prompt: str, max_sequence_length: int = 512) -> list[jax.Array]: + chat = self.tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True, enable_thinking=True + ) + # Z-Image consumes only unmasked Qwen states. Avoiding 512-token right + # padding is materially faster for its CPU text-encoder handoff. + inputs = self.tokenizer(chat, padding=True, max_length=max_sequence_length, truncation=True, return_tensors="pt") + device = next(self.text_encoder.parameters()).device + inputs = {name: value.to(device) for name, value in inputs.items()} + with torch.no_grad(): + embeddings = self.text_encoder(**inputs, output_hidden_states=True).hidden_states[-2] + mask = inputs["attention_mask"].bool() + return [ + jnp.asarray(embeddings[index][mask[index]].float().cpu().numpy(), dtype=self.dtype) for index in range(len(mask)) + ] + + @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: 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, + output_type: str = "pil", + ): + 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}.") + prompt_embeds = self.encode_prompt(prompt, max_sequence_length) + latent_height, latent_width = 2 * (height // vae_scale), 2 * (width // vae_scale) + latents = jax.random.normal( + jax.random.key(seed), (1, self.transformer.in_channels, latent_height, latent_width), self.dtype + ) + # The released scheduler is FlowMatchEulerDiscreteScheduler(shift=3.0). + sigmas = self._sigmas(num_inference_steps, shift=3.0) + with self.mesh if self.mesh is not None else nullcontext(): + for index in range(num_inference_steps): + timestep = jnp.full((1,), 1.0 - sigmas[index], dtype=self.dtype) + prediction = self.transformer([latents[0, :, None]], timestep, prompt_embeds, return_dict=False)[0][0][:, 0] + latents = latents + (sigmas[index + 1] - sigmas[index]) * (-prediction[None].astype(latents.dtype)) + + decoded = self.vae.apply( + {"params": self.vae_params}, + latents / self.vae.config.scaling_factor + self.vae_shift_factor, + deterministic=True, + method=self.vae.decode, + ).sample + image = np.asarray((decoded[0].transpose(1, 2, 0) / 2 + 0.5).clip(0, 1) * 255, dtype=np.uint8) + return Image.fromarray(image) if output_type == "pil" else image 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..fe3b4d849 --- /dev/null +++ b/src/maxdiffusion/tests/z_image/z_image_module_parity_test.py @@ -0,0 +1,139 @@ +"""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 = dict( + 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..9229b3d78 --- /dev/null +++ b/src/maxdiffusion/tests/z_image/z_image_transformer_test.py @@ -0,0 +1,34 @@ +"""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)