feat: Add Z-Image and Z-Image-Turbo support to MaxDiffusion

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

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

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)

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

src/maxdiffusion/models/z_image/__init__.py

@@ -0,0 +1,3 @@
+"""Z-Image model components."""
+
+from .transformer_z_image import ZImageTransformer2DModel