diff --git a/src/maxdiffusion/base_flux2klein.yml b/src/maxdiffusion/base_flux2klein.yml new file mode 100644 index 000000000..033e8a578 --- /dev/null +++ b/src/maxdiffusion/base_flux2klein.yml @@ -0,0 +1,276 @@ +# 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. + +# This sentinel is a reminder to choose a real run name. +run_name: 'flux2klein_test_run' + +metrics_file: "" # for testing, local file that stores scalar metrics. If empty, no metrics are written. +# If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/ +write_metrics: True + +timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written. +write_timing_metrics: True + +gcs_metrics: False +# If true save config to GCS in {base_output_directory}/{run_name}/ +save_config_to_gcs: False +log_period: 100 + +pretrained_model_name_or_path: 'black-forest-labs/FLUX.2-klein-4B' +clip_model_name_or_path: 'ariG23498/clip-vit-large-patch14-text-flax' +t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax' + +# Flux params +flux_name: "flux2klein" +scale_shift_order: "scale_shift" +use_latents: False +latents_path: "" +max_sequence_length: 512 +time_shift: True +base_shift: 0.5 +max_shift: 1.15 + + +unet_checkpoint: '' +revision: 'refs/pr/95' +# 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' + +# matmul and conv precision from https://jax.readthedocs.io/en/latest/jax.lax.html#jax.lax.Precision +# Options are "DEFAULT", "HIGH", "HIGHEST" +# fp32 activations and fp32 weights with HIGHEST will provide the best precision +# at the cost of time. +precision: "DEFAULT" + +# if False state is not jitted and instead replicate is called. This is good for debugging on single host +# It must be True for multi-host. +jit_initializers: True + +# Set true to load weights from pytorch +from_pt: True +split_head_dim: True +attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te +# If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens. +# Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster. +# However, when padding tokens are significant, this will lead to worse quality and should be set to True. +mask_padding_tokens: True +# 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 + +flash_block_sizes: {} +# GroupNorm groups +norm_num_groups: 32 + +# If train_new_flux, flux weights will be randomly initialized to train flux from scratch +# else they will be loaded from pretrained_model_name_or_path +train_new_flux: False + +# train text_encoder - Currently not supported for SDXL +train_text_encoder: False +text_encoder_learning_rate: 4.25e-6 + +# https://arxiv.org/pdf/2305.08891.pdf +snr_gamma: -1.0 + +timestep_bias: { + # a value of later will increase the frequence of the model's final training steps. + # none, earlier, later, range + strategy: "none", + # multiplier for bias, a value of 2.0 will double the weight of the bias, 0.5 will halve it. + multiplier: 1.0, + # when using strategy=range, the beginning (inclusive) timestep to bias. + begin: 0, + # when using strategy=range, the final step (inclusive) to bias. + end: 1000, + # portion of timesteps to bias. + # 0.5 will bias one half of the timesteps. Value of strategy determines + # whether the biased portions are in the earlier or later timesteps. + portion: 0.25 +} + +# Override parameters from checkpoints's scheduler. +diffusion_scheduler_config: { + _class_name: 'FlaxEulerDiscreteScheduler', + prediction_type: 'epsilon', + rescale_zero_terminal_snr: False, + timestep_spacing: 'trailing' +} + +# Output directory +# Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/" +base_output_directory: "" + +# Hardware +hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu' +skip_jax_distributed_system: False + +# Parallelism +mesh_axes: ['data', 'fsdp', 'context', 'tensor'] + +# batch : batch dimension of data and activations +# hidden : +# embed : attention qkv dense layer hidden dim named as embed +# heads : attention head dim = num_heads * head_dim +# length : attention sequence length +# temb_in : dense.shape[0] of resnet dense before conv +# out_c : dense.shape[1] of resnet dense before conv +# out_channels : conv.shape[-1] activation +# keep_1 : conv.shape[0] weight +# keep_2 : conv.shape[1] weight +# conv_in : conv.shape[2] weight +# conv_out : conv.shape[-1] weight +logical_axis_rules: [ + ['batch', 'data'], + ['activation_batch', ['data','fsdp']], + ['activation_heads', 'tensor'], + ['activation_kv', 'tensor'], + ['mlp','tensor'], + ['embed','fsdp'], + ['heads', 'tensor'], + ['conv_batch', ['data','fsdp']], + ['out_channels', 'tensor'], + ['conv_out', 'fsdp'], + ] +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. +# By default, product of the DCN axes should equal number of slices +# and product of the ICI axes should equal number of devices per slice. +dcn_data_parallelism: 1 # recommended DCN axis to be auto-sharded +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 + +# Dataset +# Replace with dataset path or train_data_dir. One has to be set. +dataset_name: 'diffusers/pokemon-gpt4-captions' +train_split: 'train' +dataset_type: 'tfrecord' # Options: 'tfrecord', 'hf', 'tf', 'grain', 'synthetic' +cache_latents_text_encoder_outputs: True +dataset_save_location: '/tmp/pokemon-gpt4-captions_xl' +train_data_dir: '' +dataset_config_name: '' +jax_cache_dir: '' +hf_data_dir: '' +hf_train_files: '' +hf_access_token: '' +image_column: 'image' +caption_column: 'text' +resolution: 512 +center_crop: False +random_flip: False +tokenize_captions_num_proc: 4 +transform_images_num_proc: 4 +reuse_example_batch: False +enable_data_shuffling: True + +# checkpoint every number of samples, -1 means don't checkpoint. +checkpoint_every: -1 +# enables one replica to read the ckpt then broadcast to the rest +enable_single_replica_ckpt_restoring: False + +# Training loop +learning_rate: 1.e-5 +scale_lr: False +max_train_samples: -1 +# max_train_steps takes priority over num_train_epochs. +max_train_steps: 1500 +num_train_epochs: 1 +seed: 0 +output_dir: 'output/' +output_name: "flux2klein_generated_image.png" +per_device_batch_size: 1.0 + +warmup_steps_fraction: 0.1 +learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps. + +# AdamW optimizer parameters +adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients. +adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients. +adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root. +adam_weight_decay: 0 # AdamW Weight decay +opt_enable_grad_clipping: False +max_grad_value: 1.0 +opt_enable_grad_global_norm_clipping: False +max_grad_norm: 1.0 + +enable_profiler: False +skip_first_n_steps_for_profiler: 5 +profiler_steps: 10 +profiler: "" + +# Generation parameters +prompt: "a car jumping off of a cliff with a crowd cheering" +prompt_2: "A detailed vector illustration of a robotic hummingbird || A cinematic shot of a neon-lit cyberpunk street" +negative_prompt: "" +do_classifier_free_guidance: True +guidance_scale: 4.0 +guidance_rescale: 0.0 +num_inference_steps: 4 +num_reps: 1 +save_final_checkpoint: False + +# SDXL Lightning parameters +lightning_from_pt: True +lightning_repo: "" +lightning_ckpt: "" + +# LoRA parameters +lora_config: { + lora_model_name_or_path: [], + weight_name: [], + adapter_name: [], + scale: [], + from_pt: [] +} + +enable_mllog: False + +#controlnet +controlnet_model_name_or_path: 'diffusers/controlnet-canny-sdxl-1.0' +controlnet_from_pt: True +controlnet_conditioning_scale: 0.5 +controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png' +quantization: '' +quantization_local_shard_count: -1 +use_qwix_quantization: False +compile_topology_num_slices: -1 # Number of target slices, set to a positive integer. + +# ML Diagnostics settings +enable_ml_diagnostics: False +profiler_gcs_path: "" +enable_ondemand_xprof: False + +# Specific additions for generate_flux2klein execution +height: 1024 +width: 1024 +batch_size: 4 +interactive: False + +# Note: Architecture dimensions (depth, num_double_layers, num_attention_heads) are +# automatically inferred from pretrained_model_name_or_path (transformer/config.json). + diff --git a/src/maxdiffusion/base_flux2klein_9B.yml b/src/maxdiffusion/base_flux2klein_9B.yml new file mode 100644 index 000000000..669a1c29e --- /dev/null +++ b/src/maxdiffusion/base_flux2klein_9B.yml @@ -0,0 +1,276 @@ +# 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. + +# This sentinel is a reminder to choose a real run name. +run_name: 'flux2klein_9b_test_run' + +metrics_file: "" # for testing, local file that stores scalar metrics. If empty, no metrics are written. +# If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/ +write_metrics: True + +timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written. +write_timing_metrics: True + +gcs_metrics: False +# If true save config to GCS in {base_output_directory}/{run_name}/ +save_config_to_gcs: False +log_period: 100 + +pretrained_model_name_or_path: 'black-forest-labs/FLUX.2-klein-9B' +clip_model_name_or_path: 'ariG23498/clip-vit-large-patch14-text-flax' +t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax' + +# Flux params +flux_name: "flux2klein_9B" +scale_shift_order: "scale_shift" +use_latents: False +latents_path: "" +max_sequence_length: 512 +time_shift: True +base_shift: 0.5 +max_shift: 1.15 + + +unet_checkpoint: '' +revision: 'refs/pr/95' +# 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' + +# matmul and conv precision from https://jax.readthedocs.io/en/latest/jax.lax.html#jax.lax.Precision +# Options are "DEFAULT", "HIGH", "HIGHEST" +# fp32 activations and fp32 weights with HIGHEST will provide the best precision +# at the cost of time. +precision: "DEFAULT" + +# if False state is not jitted and instead replicate is called. This is good for debugging on single host +# It must be True for multi-host. +jit_initializers: True + +# Set true to load weights from pytorch +from_pt: True +split_head_dim: True +attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te +# If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens. +# Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster. +# However, when padding tokens are significant, this will lead to worse quality and should be set to True. +mask_padding_tokens: True +# 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 + +flash_block_sizes: {} +# GroupNorm groups +norm_num_groups: 32 + +# If train_new_flux, flux weights will be randomly initialized to train flux from scratch +# else they will be loaded from pretrained_model_name_or_path +train_new_flux: False + +# train text_encoder - Currently not supported for SDXL +train_text_encoder: False +text_encoder_learning_rate: 4.25e-6 + +# https://arxiv.org/pdf/2305.08891.pdf +snr_gamma: -1.0 + +timestep_bias: { + # a value of later will increase the frequence of the model's final training steps. + # none, earlier, later, range + strategy: "none", + # multiplier for bias, a value of 2.0 will double the weight of the bias, 0.5 will halve it. + multiplier: 1.0, + # when using strategy=range, the beginning (inclusive) timestep to bias. + begin: 0, + # when using strategy=range, the final step (inclusive) to bias. + end: 1000, + # portion of timesteps to bias. + # 0.5 will bias one half of the timesteps. Value of strategy determines + # whether the biased portions are in the earlier or later timesteps. + portion: 0.25 +} + +# Override parameters from checkpoints's scheduler. +diffusion_scheduler_config: { + _class_name: 'FlaxEulerDiscreteScheduler', + prediction_type: 'epsilon', + rescale_zero_terminal_snr: False, + timestep_spacing: 'trailing' +} + +# Output directory +# Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/" +base_output_directory: "" + +# Hardware +hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu' +skip_jax_distributed_system: False + +# Parallelism +mesh_axes: ['data', 'fsdp', 'context', 'tensor'] + +# batch : batch dimension of data and activations +# hidden : +# embed : attention qkv dense layer hidden dim named as embed +# heads : attention head dim = num_heads * head_dim +# length : attention sequence length +# temb_in : dense.shape[0] of resnet dense before conv +# out_c : dense.shape[1] of resnet dense before conv +# out_channels : conv.shape[-1] activation +# keep_1 : conv.shape[0] weight +# keep_2 : conv.shape[1] weight +# conv_in : conv.shape[2] weight +# conv_out : conv.shape[-1] weight +logical_axis_rules: [ + ['batch', 'data'], + ['activation_batch', ['data','fsdp']], + ['activation_heads', 'tensor'], + ['activation_kv', 'tensor'], + ['mlp','tensor'], + ['embed','fsdp'], + ['heads', 'tensor'], + ['conv_batch', ['data','fsdp']], + ['out_channels', 'tensor'], + ['conv_out', 'fsdp'], + ] +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. +# By default, product of the DCN axes should equal number of slices +# and product of the ICI axes should equal number of devices per slice. +dcn_data_parallelism: 1 # recommended DCN axis to be auto-sharded +dcn_fsdp_parallelism: -1 +dcn_context_parallelism: 1 +dcn_tensor_parallelism: 1 +ici_data_parallelism: 1 +ici_fsdp_parallelism: -1 # recommended ICI axis to be auto-sharded +ici_context_parallelism: 1 +ici_tensor_parallelism: 1 + +allow_split_physical_axes: False + +# Dataset +# Replace with dataset path or train_data_dir. One has to be set. +dataset_name: 'diffusers/pokemon-gpt4-captions' +train_split: 'train' +dataset_type: 'tfrecord' # Options: 'tfrecord', 'hf', 'tf', 'grain', 'synthetic' +cache_latents_text_encoder_outputs: True +dataset_save_location: '/tmp/pokemon-gpt4-captions_xl' +train_data_dir: '' +dataset_config_name: '' +jax_cache_dir: '' +hf_data_dir: '' +hf_train_files: '' +hf_access_token: '' +image_column: 'image' +caption_column: 'text' +resolution: 512 +center_crop: False +random_flip: False +tokenize_captions_num_proc: 4 +transform_images_num_proc: 4 +reuse_example_batch: False +enable_data_shuffling: True + +# checkpoint every number of samples, -1 means don't checkpoint. +checkpoint_every: -1 +# enables one replica to read the ckpt then broadcast to the rest +enable_single_replica_ckpt_restoring: False + +# Training loop +learning_rate: 1.e-5 +scale_lr: False +max_train_samples: -1 +# max_train_steps takes priority over num_train_epochs. +max_train_steps: 1500 +num_train_epochs: 1 +seed: 0 +output_dir: 'output/' +output_name: "flux2klein_generated_image.png" +per_device_batch_size: 1.0 + +warmup_steps_fraction: 0.1 +learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps. + +# AdamW optimizer parameters +adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients. +adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients. +adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root. +adam_weight_decay: 0 # AdamW Weight decay +opt_enable_grad_clipping: False +max_grad_value: 1.0 +opt_enable_grad_global_norm_clipping: False +max_grad_norm: 1.0 + +enable_profiler: False +skip_first_n_steps_for_profiler: 5 +profiler_steps: 10 +profiler: "" + +# Generation parameters +prompt: "A detailed vector illustration of a robotic hummingbird || A cinematic shot of a neon-lit cyberpunk street" +prompt_2: "A detailed vector illustration of a robotic hummingbird || A cinematic shot of a neon-lit cyberpunk street" +negative_prompt: "" +do_classifier_free_guidance: True +guidance_scale: 4.0 +guidance_rescale: 0.0 +num_inference_steps: 4 +num_reps: 1 +save_final_checkpoint: False + +# SDXL Lightning parameters +lightning_from_pt: True +lightning_repo: "" +lightning_ckpt: "" + +# LoRA parameters +lora_config: { + lora_model_name_or_path: [], + weight_name: [], + adapter_name: [], + scale: [], + from_pt: [] +} + +enable_mllog: False + +#controlnet +controlnet_model_name_or_path: 'diffusers/controlnet-canny-sdxl-1.0' +controlnet_from_pt: True +controlnet_conditioning_scale: 0.5 +controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png' +quantization: '' +quantization_local_shard_count: -1 +use_qwix_quantization: False +compile_topology_num_slices: -1 # Number of target slices, set to a positive integer. + +# ML Diagnostics settings +enable_ml_diagnostics: False +profiler_gcs_path: "" +enable_ondemand_xprof: False + +# Specific additions for generate_flux2klein execution +height: 1024 +width: 1024 +batch_size: 4 +interactive: False + +# Note: Architecture dimensions (depth, num_double_layers, num_attention_heads) are +# automatically inferred from pretrained_model_name_or_path (transformer/config.json). + diff --git a/src/maxdiffusion/configs/base_flux2klein.yml b/src/maxdiffusion/configs/base_flux2klein.yml index f2813c8fd..033e8a578 100644 --- a/src/maxdiffusion/configs/base_flux2klein.yml +++ b/src/maxdiffusion/configs/base_flux2klein.yml @@ -203,7 +203,7 @@ num_train_epochs: 1 seed: 0 output_dir: 'output/' output_name: "flux2klein_generated_image.png" -per_device_batch_size: 1 +per_device_batch_size: 1.0 warmup_steps_fraction: 0.1 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps. @@ -231,6 +231,7 @@ do_classifier_free_guidance: True guidance_scale: 4.0 guidance_rescale: 0.0 num_inference_steps: 4 +num_reps: 1 save_final_checkpoint: False # SDXL Lightning parameters diff --git a/src/maxdiffusion/configs/base_flux2klein_9B.yml b/src/maxdiffusion/configs/base_flux2klein_9B.yml index a6c670a69..669a1c29e 100644 --- a/src/maxdiffusion/configs/base_flux2klein_9B.yml +++ b/src/maxdiffusion/configs/base_flux2klein_9B.yml @@ -203,7 +203,7 @@ num_train_epochs: 1 seed: 0 output_dir: 'output/' output_name: "flux2klein_generated_image.png" -per_device_batch_size: 1 +per_device_batch_size: 1.0 warmup_steps_fraction: 0.1 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps. @@ -231,6 +231,7 @@ do_classifier_free_guidance: True guidance_scale: 4.0 guidance_rescale: 0.0 num_inference_steps: 4 +num_reps: 1 save_final_checkpoint: False # SDXL Lightning parameters diff --git a/src/maxdiffusion/flux2klein_pipeline.py b/src/maxdiffusion/flux2klein_pipeline.py new file mode 100644 index 000000000..092555934 --- /dev/null +++ b/src/maxdiffusion/flux2klein_pipeline.py @@ -0,0 +1,553 @@ +""" +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. +""" + +import os +import time +from typing import List, Union, Optional, Any +from PIL import Image + +import sys +import jax +import jax.numpy as jnp +from jax.experimental import multihost_utils +from jax.sharding import PartitionSpec as P +import numpy as np +from flax.linen import partitioning as nn_partitioning + +from maxdiffusion import max_logging +from maxdiffusion.max_utils import device_put_replicated +from ..pipeline_flax_utils import FlaxDiffusionPipeline +from ...models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel +from ...models.vae_flax import FlaxAutoencoderKL, FlaxDecoderOutput +from ...models.qwen3_flax import FlaxQwen3Model +from ...schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler, compute_empirical_mu + +from ...models.flux.util import ( + pack_latents, + prepare_latent_image_ids, + prepare_text_ids, +) + + +class FlaxFlux2KleinPipeline(FlaxDiffusionPipeline): + """ + Unified end-to-end inference pipeline for Flux.2-klein-4B and 9B models on JAX+TPU. + Supports dynamic parameter offloading to Host CPU to optimize HBM footprint. + """ + + def __init__( + self, + transformer: Flux2KleinTransformer2DModel, + vae: FlaxAutoencoderKL, + text_encoder: FlaxQwen3Model, + tokenizer, + scheduler: FlaxFlowMatchScheduler, + config, + mesh, + **kwargs, + ): + super().__init__() + self.register_modules( + transformer=transformer, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + ) + self._config = config + self.mesh = mesh + + if tokenizer is None: + tokenizer_path = getattr(config, "tokenizer_model_name_or_path", None) or getattr( + config, "pretrained_model_name_or_path", "" + ) + hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")) + repo_cache = os.path.join( + hf_home, + "hub", + f"models--{getattr(config, 'pretrained_model_name_or_path', '').replace('/', '--')}", + "snapshots", + ) + if os.path.exists(repo_cache) and os.listdir(repo_cache): + tokenizer_path = os.path.join(repo_cache, os.listdir(repo_cache)[0]) + + from transformers import Qwen2TokenizerFast + + try: + tokenizer = Qwen2TokenizerFast.from_pretrained(tokenizer_path, local_files_only=True) + except Exception: + tokenizer = Qwen2TokenizerFast.from_pretrained(tokenizer_path, subfolder="tokenizer", local_files_only=True) + + self.tokenizer = tokenizer + + # JIT compilation cache + self._jitted_qwen3_forward = None + self._jitted_transformer_step = None + self._jitted_vae_decode = None + + def _setup_jit_functions(self): + if self._jitted_qwen3_forward is not None: + return + + @jax.jit + def qwen3_forward(q_params, ids, mask): + return self.text_encoder.apply({"params": q_params}, input_ids=ids, attention_mask=mask) + + @jax.jit + def transformer_step(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timestep, guidance): + return self.transformer.apply( + {"params": t_params}, + hidden_states=latents, + img_ids=img_ids, + encoder_hidden_states=prompt_embeds, + txt_ids=txt_ids, + pooled_projections=vec, + timestep=timestep, + guidance=guidance, + ) + + @jax.jit(static_argnums=(4, 5), donate_argnums=(1,)) + def vae_decode(v_params, latents_packed, vae_bn_mean, vae_bn_std, height, width): + batch_size_val = latents_packed.shape[0] + h_latent = height // 8 + w_latent = width // 8 + + vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128) + vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128) + + latents_bn = latents_packed * vae_bn_std_seq + vae_bn_mean_seq + latents_unpacked = jnp.reshape(latents_bn, (batch_size_val, h_latent // 2, w_latent // 2, 32, 2, 2)) + latents_unpacked = jnp.transpose(latents_unpacked, (0, 3, 1, 4, 2, 5)) + latents_unpacked = jnp.reshape(latents_unpacked, (batch_size_val, 32, h_latent, w_latent)) + + res = self.vae.apply({"params": v_params}, latents=latents_unpacked, method=self.vae.decode) + return FlaxDecoderOutput(sample=res.sample) + + @jax.jit + def fused_denoise_loop(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance): + def scan_body(cur_latents, step_idx): + t_val = timesteps[step_idx] + t_vec = jnp.broadcast_to(t_val / 1000.0, (cur_latents.shape[0],)) + model_output = self.transformer.apply( + {"params": t_params}, + hidden_states=cur_latents, + img_ids=img_ids, + encoder_hidden_states=prompt_embeds, + txt_ids=txt_ids, + pooled_projections=vec, + timestep=t_vec, + guidance=guidance, + ) + sigma = sigmas[step_idx] + sigma_next = sigmas[step_idx + 1] + prev_sample = cur_latents + model_output.sample * (sigma_next - sigma) + return prev_sample, None + + steps = jnp.arange(timesteps.shape[0]) + final_latents, _ = jax.lax.scan(scan_body, latents, steps) + return final_latents + + self._jitted_qwen3_forward = qwen3_forward + self._jitted_transformer_step = transformer_step + self._jitted_fused_denoise_loop = fused_denoise_loop + self._jitted_vae_decode = vae_decode + + def _get_dynamic_batch_sharding(self): + """Dynamically infers the batch dimension sharding specification from self.mesh.""" + batch_axes = [axis for axis in ("data", "fsdp") if axis in self.mesh.axis_names and self.mesh.shape[axis] > 1] + spec = P(tuple(batch_axes)) if batch_axes else P(None) + return jax.sharding.NamedSharding(self.mesh, spec) + + def compile_aot_async( + self, params, vae_params, qwen3_params, vae_bn_mean, vae_bn_std, batch_size=1, height=1024, width=1024 + ): + """Triggers AOT compilation for Qwen3, Flux Transformer, and VAE concurrently using ThreadPoolExecutor.""" + self._setup_jit_functions() + max_logging.log("šŸš€ Pre-compiling XLA graphs for Qwen3, Flux Transformer, and VAE concurrently...") + from concurrent.futures import ThreadPoolExecutor + + seq_len_img = (height // 16) * (width // 16) + seq_len_txt = self._config.max_sequence_length + + dummy_ids = jnp.zeros((batch_size, seq_len_txt), dtype=jnp.int32) + dummy_mask = jnp.ones((batch_size, seq_len_txt), dtype=jnp.int32) + + dummy_latents = jnp.zeros((batch_size, seq_len_img, 128), dtype=jnp.float32) + dummy_img_ids = jnp.zeros((batch_size, seq_len_img, 4), dtype=jnp.int32) + dummy_prompt_embeds = jnp.zeros((batch_size, seq_len_txt, self.transformer.joint_attention_dim), dtype=jnp.bfloat16) + dummy_txt_ids = jnp.zeros((batch_size, seq_len_txt, 4), dtype=jnp.float32) + dummy_t_vec = jnp.zeros((batch_size,), dtype=jnp.float32) + + dummy_bn_mean = jnp.array(vae_bn_mean, dtype=jnp.float32) + dummy_bn_std = jnp.array(vae_bn_std, dtype=jnp.float32) + + data_sharding = self._get_dynamic_batch_sharding() + replicated_sharding = jax.sharding.NamedSharding(self.mesh, P()) + + def put_data_on_devices(x, sharding): + if isinstance(x, jax.Array) and hasattr(x, "sharding") and not x.sharding.is_fully_addressable: + return x + if hasattr(sharding, "is_fully_addressable") and sharding.is_fully_addressable: + return jax.device_put(x, sharding) + return device_put_replicated(x, sharding) + + dummy_ids = put_data_on_devices(dummy_ids, data_sharding) + dummy_mask = put_data_on_devices(dummy_mask, data_sharding) + dummy_latents = put_data_on_devices(dummy_latents, data_sharding) + dummy_img_ids = put_data_on_devices(dummy_img_ids, data_sharding) + dummy_prompt_embeds = put_data_on_devices(dummy_prompt_embeds, data_sharding) + dummy_txt_ids = put_data_on_devices(dummy_txt_ids, data_sharding) + dummy_t_vec = put_data_on_devices(dummy_t_vec, data_sharding) + dummy_bn_mean = put_data_on_devices(dummy_bn_mean, replicated_sharding) + dummy_bn_std = put_data_on_devices(dummy_bn_std, replicated_sharding) + + def compile_qwen3(): + t0 = time.perf_counter() + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + self._jitted_qwen3_forward.lower(qwen3_params, dummy_ids, dummy_mask).compile() + max_logging.log(f" -> [AOT COMPILED] Qwen3 Text Encoder in {time.perf_counter() - t0:.2f}s") + + num_steps = getattr(self._config, "num_inference_steps", 4) + dummy_timesteps = put_data_on_devices(jnp.zeros((num_steps,), dtype=jnp.float32), replicated_sharding) + dummy_sigmas = put_data_on_devices(jnp.zeros((num_steps + 1,), dtype=jnp.float32), replicated_sharding) + + def compile_transformer(): + t0 = time.perf_counter() + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + self._jitted_fused_denoise_loop.lower( + params, + dummy_latents, + dummy_img_ids, + dummy_prompt_embeds, + dummy_txt_ids, + None, + dummy_timesteps, + dummy_sigmas, + None, + ).compile() + max_logging.log(f" -> [AOT COMPILED] Fused Denoising Loop (Scan {num_steps} steps) in {time.perf_counter() - t0:.2f}s") + + def compile_vae(): + t0 = time.perf_counter() + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + self._jitted_vae_decode.lower(vae_params, dummy_latents, dummy_bn_mean, dummy_bn_std, height, width).compile() + max_logging.log(f" -> [AOT COMPILED] VAE Decoder in {time.perf_counter() - t0:.2f}s") + + t_start = time.perf_counter() + with ThreadPoolExecutor(max_workers=3) as executor: + futures = [ + executor.submit(compile_qwen3), + executor.submit(compile_transformer), + executor.submit(compile_vae), + ] + for future in futures: + future.result() + aot_duration = time.perf_counter() - t_start + max_logging.log(f"āš” [AOT CONCURRENT COMPILATION COMPLETE] Total AOT compile time: {aot_duration:.2f}s") + return aot_duration + + def _prepare_latents(self, config, batch_size, height, width): + num_channels_latents = 32 + latent_height = height // 8 + latent_width = width // 8 + latent_shape = (batch_size, num_channels_latents, latent_height, latent_width) + + seed_val = getattr(config, "seed", None) + if seed_val is None: + seed_val = int(time.time()) & 0x7FFFFFFF + max_logging.log( + f"Generating random gaussian noise in unpacked space (32 channels) with seed: {seed_val} and shape: {latent_shape}..." + ) + np.random.seed(seed_val) + latents_unpacked = np.random.randn(*latent_shape).astype(np.float32) + + # Pack/patchify noise exactly like PyTorch: + # (batch, 32, H/16, 2, W/16, 2) -> permute(0, 1, 3, 5, 2, 4) -> reshape(batch, 128, H/16, W/16) + B, C, H, W = latents_unpacked.shape + latents_packed = latents_unpacked.reshape(B, C, H // 2, 2, W // 2, 2) + latents_packed = np.transpose(latents_packed, (0, 1, 3, 5, 2, 4)) + latents_packed = latents_packed.reshape(B, 128, H // 2, W // 2) + + return latents_packed + + def __call__( + self, + prompt: Union[str, List[str]], + params, + vae_params, + qwen3_params, + vae_bn_mean, + vae_bn_std, + transformer_shardings, + vae_shardings, + qwen3_shardings, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 4, + batch_size: int = 1, + use_latents: bool = False, + latents: Optional[Any] = None, + measure_time: bool = False, + warmup: bool = False, + output_dir: str = "output/", + output_name: str = "flux2klein_generated_image.png", + ): + # 1. Setup JIT functions + self._setup_jit_functions() + + # 2. Setup prompts and inputs + if isinstance(prompt, str): + prompts = [prompt] * batch_size + else: + prompts = prompt + + seq_len_img = (height // 16) * (width // 16) + seq_len_txt = self._config.max_sequence_length + + # Load or generate latents + if use_latents and latents is not None: + latents_jax = jnp.array(latents) + if latents_jax.ndim == 4: + B, C, H, W = latents_jax.shape + if C == 32: + max_logging.log(" [PIPELINE] Unpacked 32-channel latents detected. Packing using pack_latents...") + latents_jax = pack_latents(latents_jax) + else: + latents_jax = jnp.transpose(jnp.reshape(latents_jax, (B, C, H * W)), (0, 2, 1)) + else: + latents_numpy = self._prepare_latents(self._config, batch_size, height, width) + B, C, H, W = latents_numpy.shape + latents_jax = jnp.transpose(jnp.reshape(latents_numpy, (B, C, H * W)), (0, 2, 1)) + + # RoPE position IDs + txt_ids_val = prepare_text_ids(batch_size, seq_len_txt) + img_ids_val = prepare_latent_image_ids(batch_size, height // 16, width // 16) + + # Scheduler + mu = compute_empirical_mu(seq_len_img, num_inference_steps) + scheduler_state = self.scheduler.create_state() + sigmas_custom = jnp.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps, dtype=jnp.float32) + scheduler_state = self.scheduler.set_timesteps_ltx2( + state=scheduler_state, + num_inference_steps=num_inference_steps, + shift=mu, + sigmas=sigmas_custom, + ) + + t_pipeline_start = time.perf_counter() + trace = {} + + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + proc_id = jax.process_index() + proc_cnt = jax.process_count() + host_prefix = f"[HOST {proc_id}/{proc_cnt}] " + + # Shard pipeline batch inputs across data axis ("data") for SPMD multi-host execution + data_sharding = jax.sharding.NamedSharding(self.mesh, P("data")) + + def put_data_on_devices(x, sharding): + if isinstance(x, jax.Array) and hasattr(x, "sharding") and not x.sharding.is_fully_addressable: + return x + if hasattr(sharding, "is_fully_addressable") and sharding.is_fully_addressable: + return jax.device_put(x, sharding) + return device_put_replicated(x, sharding) + + t0_qwen3_start = time.perf_counter() + trace["start_to_qwen3"] = t0_qwen3_start - t_pipeline_start + max_logging.log(f" -> [TIMING] Start to Qwen3: {trace['start_to_qwen3']:.4f} seconds ā±ļø") + + # --------------------------------------------------------------------- + # PHASE A: Encode Prompt (Qwen3) + # --------------------------------------------------------------------- + if prompts is None: + prompts = ["A dog running in a field with butterflies and tall grass"] + elif isinstance(prompts, str): + prompts = [prompts] + + max_logging.log(f"{host_prefix} [PHASE A] Encoding {len(prompts)} prompt(s) using JAX Qwen3 on TPU...") + + try: + # Tokenize using deterministic explicit template string (version-agnostic across transformers versions) + templated_texts = [ + f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" for p in prompts + ] + inputs = self.tokenizer( + templated_texts, return_tensors="np", padding="max_length", truncation=True, max_length=seq_len_txt + ) + prompt_ids = jnp.array(inputs["input_ids"]) + prompt_mask = jnp.array(inputs["attention_mask"]) + + # Run Text Encoding with sharded input arrays matching compile_aot_async + prompt_ids = put_data_on_devices(prompt_ids, data_sharding) + prompt_mask = put_data_on_devices(prompt_mask, data_sharding) + with jax.named_scope("qwen3_text_encoder"): + hidden_states, all_hidden_states = self._jitted_qwen3_forward(qwen3_params, prompt_ids, prompt_mask) + + # Stack layers 9, 18, 27 to form prompt embeddings + h_9 = all_hidden_states[9] + h_18 = all_hidden_states[18] + h_27 = all_hidden_states[27] + out = jnp.stack([h_9, h_18, h_27], axis=1) + # Transpose shape to [B, seq_len, 3*hidden_size] + prompt_embeds_jax = jnp.transpose(out, (0, 2, 1, 3)).reshape((batch_size, seq_len_txt, -1)) + prompt_embeds_jax.block_until_ready() + except Exception as e: + max_logging.log(f"āŒ {host_prefix} EXCEPTION IN PHASE A (QWEN3 ENCODING): {e}") + import traceback + + traceback.print_exc() + sys.stdout.flush() + raise e + + t0_qwen3_end = time.perf_counter() + trace["qwen3_encoding"] = t0_qwen3_end - t0_qwen3_start + trace["prompt_encoding"] = trace["qwen3_encoding"] + max_logging.log(f" -> [TIMING] Prompt Encoding (Qwen3): {trace['qwen3_encoding']:.4f} seconds ā±ļø") + + # Stage Sync 1: Phase A Complete + multihost_utils.sync_global_devices("phase_a_complete") + max_logging.log(f"{host_prefix} Passed Phase A Sync Barrier (phase_a_complete) successfully! āœ…") + + latents_jax = put_data_on_devices(latents_jax, data_sharding) + prompt_embeds_jax = put_data_on_devices(prompt_embeds_jax, data_sharding) + txt_ids_val = put_data_on_devices(txt_ids_val, data_sharding) + img_ids_val = put_data_on_devices(img_ids_val, data_sharding) + + max_logging.log( + f"{host_prefix} DIAGNOSTIC TENSORS BEFORE PHASE B:\n" + f" latents_jax: shape={latents_jax.shape}, dtype={latents_jax.dtype}, sharding={getattr(latents_jax, 'sharding', None)}\n" + f" prompt_embeds_jax: shape={prompt_embeds_jax.shape}, dtype={prompt_embeds_jax.dtype}, sharding={getattr(prompt_embeds_jax, 'sharding', None)}\n" + f" txt_ids_val: shape={txt_ids_val.shape}, dtype={txt_ids_val.dtype}, sharding={getattr(txt_ids_val, 'sharding', None)}\n" + f" img_ids_val: shape={img_ids_val.shape}, dtype={img_ids_val.dtype}, sharding={getattr(img_ids_val, 'sharding', None)}" + ) + + # Stage Sync 2: Pre-Phase B Start + multihost_utils.sync_global_devices("pre_phase_b_start") + max_logging.log(f"{host_prefix} Passed Pre-Phase B Sync Barrier (pre_phase_b_start) successfully! āœ…") + + t0_denoise_start = time.perf_counter() + trace["qwen3_to_denoise"] = t0_denoise_start - t0_qwen3_end + max_logging.log(f" -> [TIMING] Qwen3 to Denoising Overhead: {trace['qwen3_to_denoise']:.4f} seconds ā±ļø") + + # --------------------------------------------------------------------- + # PHASE B: Denoising Loop (Fused Flux Transformer Scan Loop) + # --------------------------------------------------------------------- + steps_to_run = num_inference_steps + max_logging.log( + f"{host_prefix} [PHASE B] Running fused {steps_to_run}-step E2E Denoising Loop Scan on a batch of {batch_size} images (warmup={warmup})..." + ) + + try: + guidance_vec_val = None + vec_val = None + replicated_sharding = jax.sharding.NamedSharding(self.mesh, P()) + timesteps_device = put_data_on_devices(scheduler_state.timesteps, replicated_sharding) + sigmas_device = put_data_on_devices(scheduler_state.sigmas, replicated_sharding) + + with jax.named_scope("fused_flux_denoise_loop"): + latents_jax = self._jitted_fused_denoise_loop( + params, + latents_jax, + img_ids_val, + prompt_embeds_jax, + txt_ids_val, + vec_val, + timesteps_device, + sigmas_device, + guidance_vec_val, + ) + latents_jax.block_until_ready() + + except Exception as e: + max_logging.log(f"āŒ {host_prefix} EXCEPTION IN DENOISE LOOP: {e}") + import traceback + + traceback.print_exc() + sys.stdout.flush() + raise e + + # Stage Sync 3: Phase B Complete + multihost_utils.sync_global_devices("phase_b_complete") + max_logging.log(f"{host_prefix} Passed Phase B Sync Barrier (phase_b_complete) successfully! āœ…") + + t0_denoise_end = time.perf_counter() + trace["denoise_loop"] = t0_denoise_end - t0_denoise_start + max_logging.log(f" -> [TIMING] Denoising Loop (Flux): {trace['denoise_loop']:.4f} seconds ā±ļø") + + # --------------------------------------------------------------------- + # PHASE C: Decode Latents (VAE Decoder) + # --------------------------------------------------------------------- + max_logging.log("[PHASE C] Decoding final latents to RGB image using JAX VAE decoder on TPU...") + + # Decode VAE latents to RGB pixels using fused JIT vae_decode + data_sharding = self._get_dynamic_batch_sharding() + replicated_sharding = jax.sharding.NamedSharding(self.mesh, P()) + latents_jax = put_data_on_devices(latents_jax, data_sharding) + vae_bn_mean_jax = put_data_on_devices(jnp.array(vae_bn_mean, dtype=jnp.float32), replicated_sharding) + vae_bn_std_jax = put_data_on_devices(jnp.array(vae_bn_std, dtype=jnp.float32), replicated_sharding) + + t0_vae_start = time.perf_counter() + trace["denoise_to_vae"] = t0_vae_start - t0_denoise_end + max_logging.log(f" -> [TIMING] Denoising to VAE Overhead: {trace['denoise_to_vae']:.4f} seconds ā±ļø") + + with jax.named_scope("vae_decoder"): + decoded_out = self._jitted_vae_decode(vae_params, latents_jax, vae_bn_mean_jax, vae_bn_std_jax, height, width) + images_rgb = decoded_out.sample + images_rgb.block_until_ready() + + t0_vae_end = time.perf_counter() + trace["vae_decode"] = t0_vae_end - t0_vae_start + max_logging.log(f" -> [TIMING] VAE Decoding: {trace['vae_decode']:.4f} seconds ā±ļø") + + # --------------------------------------------------------------------- + # POST-PROCESS: Format and Save Outputs + # --------------------------------------------------------------------- + max_logging.log("Postprocessing and saving generated images...") + saved_paths = [] + # Perform pixel scaling, clamping, and uint8 conversion directly on TPU hardware + images_uint8 = jnp.clip((images_rgb + 1.0) * 127.5, 0.0, 255.0).astype(jnp.uint8) + if jax.process_count() > 1: + images_numpy = multihost_utils.process_allgather(images_uint8, tiled=True) + else: + images_numpy = np.array(images_uint8) + + for b_idx in range(batch_size): + image_np = np.array(images_numpy[b_idx]) + # Transpose channel dimension if shape is (C, H, W) instead of (H, W, C) + if image_np.shape[0] == 3: + image_np = image_np.transpose(1, 2, 0) + + img = Image.fromarray(image_np) + + # Formulate output filename for this batch index + if batch_size > 1: + batch_output_name = output_name.replace(".png", f"_b{b_idx}.png") + else: + batch_output_name = output_name + + output_png_path = os.path.join(output_dir, batch_output_name) + img.save(output_png_path, format="PNG", compress_level=1) + max_logging.log(f" -> Saved image: {output_png_path} | Prompt: '{prompts[b_idx]}'") + saved_paths.append(output_png_path) + + t0_save_end = time.perf_counter() + trace["image_saving"] = t0_save_end - t0_vae_end + trace["e2e_pipeline_total"] = t0_save_end - t_pipeline_start + + max_logging.log(f" -> [TIMING] Image Saving: {trace['image_saving']:.4f} seconds ā±ļø") + max_logging.log(f" -> [TIMING] E2E Pipeline Total: {trace['e2e_pipeline_total']:.4f} seconds ā±ļø") + + return saved_paths, trace diff --git a/src/maxdiffusion/generate_flux2klein.py b/src/maxdiffusion/generate_flux2klein.py index 7956c850d..02b929c6f 100644 --- a/src/maxdiffusion/generate_flux2klein.py +++ b/src/maxdiffusion/generate_flux2klein.py @@ -79,8 +79,22 @@ def encode_prompt(prompt: str, snapshot_dir: str = None, repo_id: str = "black-f text_encoder_path = os.path.join(snapshot_dir, "text_encoder") tokenizer_path = os.path.join(snapshot_dir, "tokenizer") - if not os.path.exists(tokenizer_path): - tokenizer_path = text_encoder_path + + if not os.path.exists(os.path.join(text_encoder_path, "config.json")) or not os.path.exists(tokenizer_path): + try: + fb_dir = snapshot_download(repo_id=repo_id, local_files_only=True) + if not os.path.exists(os.path.join(text_encoder_path, "config.json")): + text_encoder_path = os.path.join(fb_dir, "text_encoder") + if not os.path.exists(tokenizer_path): + tokenizer_path = ( + os.path.join(fb_dir, "tokenizer") + if os.path.exists(os.path.join(fb_dir, "tokenizer")) + else os.path.join(fb_dir, "text_encoder") + ) + except Exception: + if not os.path.exists(tokenizer_path): + tokenizer_path = text_encoder_path + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) text_encoder = AutoModelForCausalLM.from_pretrained(text_encoder_path, torch_dtype=torch.float32) text_encoder.eval() @@ -134,15 +148,34 @@ def main(argv): from maxdiffusion.models.flux.util import ( load_and_convert_flux_klein_weights, load_and_convert_vae_weights, - cast_dict_to_bfloat16_inplace, ) from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline config = pyconfig.config os.makedirs(config.output_dir, exist_ok=True) + if hasattr(config, "per_device_batch_size") and config.per_device_batch_size > 0: + calculated_batch_size = int(config.per_device_batch_size * jax.device_count()) + assert calculated_batch_size >= 1, ( + f"Calculated global batch_size is {calculated_batch_size}, which is invalid (must be >= 1). " + f"per_device_batch_size={config.per_device_batch_size} multiplied by jax.device_count()={jax.device_count()} " + f"evaluated to {config.per_device_batch_size * jax.device_count()}, which truncates to 0. " + f"Please increase per_device_batch_size or specify an explicit batch_size in your configuration." + ) + if calculated_batch_size != config.batch_size: + max_logging.log( + f"ā„¹ļø Updating batch_size from {config.batch_size} to {calculated_batch_size} " + f"based on per_device_batch_size={config.per_device_batch_size} and device_count={jax.device_count()}." + ) + pyconfig._config.keys["batch_size"] = calculated_batch_size + # 2. Setup device mesh - if config.batch_size == 1 and config.ici_tensor_parallelism == 1 and jax.device_count() > 1: + if ( + config.batch_size == 1 + and config.ici_tensor_parallelism == 1 + and config.ici_context_parallelism == 1 + and jax.device_count() > 1 + ): max_logging.log( f"ā„¹ļø Auto-configuring Tensor Parallelism: ici_tensor_parallelism={jax.device_count()}, ici_fsdp_parallelism=1 for batch_size=1 on {jax.device_count()} TPU devices." ) @@ -174,8 +207,7 @@ def main(argv): # 3. Resolve weights repository snapshots repo_id = getattr(config, "pretrained_model_name_or_path", None) if not repo_id: - depth_val = getattr(config, "depth", None) - repo_id = "black-forest-labs/FLUX.2-klein-9B" if depth_val == 24 else "black-forest-labs/FLUX.2-klein-4B" + raise ValueError("pretrained_model_name_or_path must be specified in configuration YAML or CLI.") max_logging.log(f"Target model detected: {repo_id}") if os.path.exists(repo_id): @@ -184,8 +216,13 @@ def main(argv): else: from huggingface_hub import snapshot_download - max_logging.log(f"Resolving snapshot directory for model '{repo_id}' from HF Hub...") - snapshot_dir = snapshot_download(repo_id=repo_id) + rev = getattr(config, "revision", None) + if not rev or rev == "refs/pr/95": + rev = "main" + try: + snapshot_dir = snapshot_download(repo_id=repo_id, revision=rev, local_files_only=True) + except Exception: + snapshot_dir = snapshot_download(repo_id=repo_id, revision=rev) max_logging.log(f"Host {jax.process_index()} using HF snapshot directory: {snapshot_dir}") safetensors_path = os.path.join(snapshot_dir, "transformer") @@ -195,8 +232,7 @@ def main(argv): # 4. Load Qwen3 Config & Setup model layout from transformers import AutoConfig - max_logging.log(f"Loading Qwen3 config from text_encoder path: {text_encoder_path}...") - pt_config = AutoConfig.from_pretrained(text_encoder_path, local_files_only=True) + pt_config = AutoConfig.from_pretrained(text_encoder_path) qwen3_config = FlaxQwen3Config( vocab_size=pt_config.vocab_size, @@ -217,9 +253,24 @@ def main(argv): transformer_config_json = os.path.join(safetensors_path, "config.json") transformer_pt_cfg = {} + loaded_cfg = False if os.path.exists(transformer_config_json): - with open(transformer_config_json, "r") as f: - transformer_pt_cfg = json.load(f) + try: + with open(transformer_config_json, "r") as f: + transformer_pt_cfg = json.load(f) + loaded_cfg = True + except Exception as e: + max_logging.log(f"ā„¹ļø Could not parse {transformer_config_json}: {e}. Falling back to HF cache...") + + if not loaded_cfg and repo_id: + try: + from huggingface_hub import hf_hub_download + + cfg_file = hf_hub_download(repo_id=repo_id, filename="transformer/config.json", local_files_only=True) + with open(cfg_file, "r") as f: + transformer_pt_cfg = json.load(f) + except Exception as e: + max_logging.log(f"āš ļø Warning resolving transformer config fallback from HF cache: {e}") num_double_layers = getattr(config, "num_double_layers", -1) if num_double_layers is None or num_double_layers <= 0: @@ -342,6 +393,7 @@ def qwen3_init_fn(): def unbox_fn(x): return x.unbox() if isinstance(x, flax_spmd.LogicallyPartitioned) else x + t_sub0 = time.time() params = jax.tree_util.tree_map( unbox_fn, abstract_transformer_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned) ) @@ -357,17 +409,19 @@ def unbox_fn(x): ) qwen3_params = flax.core.unfreeze(qwen3_params) - params = load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, depth) - vae_params, vae_bn_mean, vae_bn_std = load_and_convert_vae_weights(vae_safetensors_path, vae_params) - qwen3_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params, qwen3_config) + max_logging.log(f" -> [SUB-TIMING 1/3] PyTree unboxing template setup: {time.time() - t_sub0:.2f}s") + t_sub1 = time.time() + + weight_dtype = jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32 - if config.weights_dtype == "bfloat16": - max_logging.log("Casting JAX parameters to bfloat16 in-place...") - cast_dict_to_bfloat16_inplace(params, exclude_keywords=("norm",)) - cast_dict_to_bfloat16_inplace(vae_params, exclude_keywords=("norm",)) - cast_dict_to_bfloat16_inplace(qwen3_params, exclude_keywords=("norm",)) - vae_bn_mean = vae_bn_mean.astype(jnp.bfloat16) - vae_bn_std = vae_bn_std.astype(jnp.bfloat16) + params = load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, depth, dtype=weight_dtype) + vae_params, vae_bn_mean, vae_bn_std = load_and_convert_vae_weights( + vae_safetensors_path, vae_params, dtype=weight_dtype + ) + qwen3_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params, qwen3_config) + max_logging.log( + f" -> [SUB-TIMING 2/3] Safetensors loading & key mapping (in target dtype): {time.time() - t_sub1:.4f}s" + ) params = flax.core.freeze(params) vae_params = flax.core.freeze(vae_params) @@ -376,6 +430,7 @@ def unbox_fn(x): max_logging.log("\n" + "=" * 80) max_logging.log("šŸš€ Pinning all parameters to TPU HBM permanently...") max_logging.log("=" * 80 + "\n") + t_sub3 = time.time() max_logging.log("Putting params on TPU HBM...") with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): try: @@ -394,12 +449,13 @@ def unbox_fn(x): vae_params = jax.tree_util.tree_map(max_utils.device_put_replicated, vae_params, vae_shardings) max_logging.log("Putting qwen3_params on TPU HBM...") qwen3_params = jax.tree_util.tree_map(max_utils.device_put_replicated, qwen3_params, qwen3_shardings) + max_logging.log(f" -> [SUB-TIMING 3/3] TPU HBM device_put placement: {time.time() - t_sub3:.4f}s") max_logging.log("All parameters placed on TPU HBM successfully!") gc.collect() jax.effects_barrier() load_time = time.time() - t_load_start - max_logging.log(f" -> [TIMING] Total Model Loading & Device Placement: {load_time:.2f} seconds ā±ļø\n") + max_logging.log(f" -> [TIMING] Total Model Loading & Device Placement: {load_time:.4f} seconds ā±ļø\n") # 9. Setup FlowMatch Scheduler scheduler = FlaxFlowMatchScheduler( @@ -426,16 +482,19 @@ def unbox_fn(x): mesh=mesh, ) - active_prompts = partition_prompts(config.prompt, config.batch_size) + prompt_str = getattr(config, "prompt", None) + if not prompt_str: + raise ValueError("Prompt must be specified in the configuration YAML or passed via CLI prompt='...'") + active_prompts = partition_prompts(prompt_str, config.batch_size) if getattr(config, "interactive", False): - print("\n" + "=" * 80) - print(" BATCHED INTERACTIVE GENERATION MODE ENABLED šŸŽ®") - print("The model has been fully loaded and compiled on the TPU.") - print(f"Batch size: {config.batch_size} parallel images.") - print("Enter prompts separated by '||' (e.g. A cute cat || A red car)") - print("Type 'exit' to quit.") - print("=" * 80) + max_logging.log("\n" + "=" * 80) + max_logging.log(" BATCHED INTERACTIVE GENERATION MODE ENABLED šŸŽ®") + max_logging.log("The model has been fully loaded and compiled on the TPU.") + max_logging.log(f"Batch size: {config.batch_size} parallel images.") + max_logging.log("Enter prompts separated by '||' (e.g. A cute cat || A red car)") + max_logging.log("Type 'exit' to quit.") + max_logging.log("=" * 80) image_idx = 1 while True: @@ -481,37 +540,23 @@ def unbox_fn(x): max_logging.log(f" -> Custom latents shape: {latents_to_use.shape} | sum: {latents_to_use.sum():.6f}") max_logging.log("\n" + "=" * 80) - max_logging.log("šŸš€ Running initial dry run (Warmup Pass) to compile XLA graphs...") + max_logging.log("šŸš€ Pre-compiling XLA graphs concurrently (AOT Compilation)...") max_logging.log("=" * 80) - _, warmup_trace = pipeline( - prompt=active_prompts, + aot_time = pipeline.compile_aot_async( params=params, vae_params=vae_params, qwen3_params=qwen3_params, vae_bn_mean=vae_bn_mean, vae_bn_std=vae_bn_std, - transformer_shardings=transformer_shardings, - vae_shardings=vae_shardings, - qwen3_shardings=qwen3_shardings, + batch_size=config.batch_size, height=config.height, width=config.width, - num_inference_steps=config.num_inference_steps, - batch_size=config.batch_size, - use_latents=use_latents_flag, - latents=latents_to_use, - output_dir=config.output_dir, - output_name="flux2klein_warmup.png", - ) - warmup_time = ( - warmup_trace.get("prompt_encoding", 0.0) - + warmup_trace.get("denoise_loop", 0.0) - + warmup_trace.get("vae_decode", 0.0) ) max_logging.log("\n" + "=" * 80) - max_logging.log("ā±ļø Running timed pass at full TPU speed...") + max_logging.log("šŸš€ Running initial dry run (Warmup Pass) to verify compiled graph execution...") max_logging.log("=" * 80) - _, main_trace = pipeline( + _, warmup_trace = pipeline( prompt=active_prompts, params=params, vae_params=vae_params, @@ -528,24 +573,101 @@ def unbox_fn(x): use_latents=use_latents_flag, latents=latents_to_use, output_dir=config.output_dir, - output_name=config.output_name, + output_name="flux2klein_warmup.png", + warmup=True, ) - main_time = ( - main_trace.get("prompt_encoding", 0.0) + main_trace.get("denoise_loop", 0.0) + main_trace.get("vae_decode", 0.0) + warmup_time = ( + warmup_trace.get("prompt_encoding", 0.0) + + warmup_trace.get("denoise_loop", 0.0) + + warmup_trace.get("vae_decode", 0.0) ) + num_reps = int(getattr(config, "num_reps", 1)) + max_logging.log("\n" + "=" * 80) + max_logging.log(f"ā±ļø Running timed pass at full TPU speed (num_reps={num_reps})...") + max_logging.log("=" * 80) + + main_traces = [] + main_times = [] + + for rep in range(num_reps): + rep_str = f" [Rep {rep+1}/{num_reps}]" if num_reps > 1 else "" + if rep > 0: + max_logging.log(f"ā±ļø Running timed pass{rep_str}...") + + if max_utils.profiler_enabled(config) and rep == 0: + max_logging.log(f"šŸš€ XProf / JAX Profiler active! Capturing trace into: {config.tensorboard_dir}") + with max_utils.Profiler(config, session_name="flux2klein_inference"): + _, trace_i = pipeline( + prompt=active_prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=use_latents_flag, + latents=latents_to_use, + output_dir=config.output_dir, + output_name=f"rep_{rep+1}_{config.output_name}" if num_reps > 1 else config.output_name, + ) + else: + _, trace_i = pipeline( + prompt=active_prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=use_latents_flag, + latents=latents_to_use, + output_dir=config.output_dir, + output_name=f"rep_{rep+1}_{config.output_name}" if num_reps > 1 else config.output_name, + ) + + tot_time_i = trace_i.get("prompt_encoding", 0.0) + trace_i.get("denoise_loop", 0.0) + trace_i.get("vae_decode", 0.0) + main_traces.append(trace_i) + main_times.append(tot_time_i) + if num_reps > 1: + max_logging.log( + f" -> Rep {rep+1}/{num_reps} Completed: Total={tot_time_i:.4f}s | Qwen3={trace_i.get('prompt_encoding', 0.0):.4f}s | Denoise={trace_i.get('denoise_loop', 0.0):.4f}s | VAE={trace_i.get('vae_decode', 0.0):.4f}s" + ) + + avg_main_time = sum(main_times) / num_reps + avg_prompt_enc = sum(tr.get("prompt_encoding", 0.0) for tr in main_traces) / num_reps + avg_denoise = sum(tr.get("denoise_loop", 0.0) for tr in main_traces) / num_reps + avg_vae_decode = sum(tr.get("vae_decode", 0.0) for tr in main_traces) / num_reps + + total_cold_start = load_time + aot_time + warmup_time + max_logging.log("\n" + "=" * 80) - max_logging.log("šŸ“Š FLUX.2-KLEIN LATENCY & TIMING BREAKDOWN (PURE MODEL INFERENCE)") + max_logging.log("šŸ“Š FLUX.2-KLEIN COMPLETE LATENCY & TIMING BREAKDOWN") max_logging.log("=" * 80) - max_logging.log(f"1) Total Model Loading & Placement Time: {load_time:.2f} seconds ā±ļø") - max_logging.log(f"2) Cold-Start / Warmup Pass (XLA Compilation): {warmup_time:.2f} seconds ā±ļø") - max_logging.log(f" - Qwen3 Encoding: {warmup_trace.get('prompt_encoding', 0.0):.2f}s") - max_logging.log(f" - Flux Denoising: {warmup_trace.get('denoise_loop', 0.0):.2f}s") - max_logging.log(f" - VAE Decoding: {warmup_trace.get('vae_decode', 0.0):.2f}s") - max_logging.log(f"3) Main Warmed-Up Pass (Pure Model Inference): {main_time:.2f} seconds ā±ļø") - max_logging.log(f" - Qwen3 Encoding: {main_trace.get('prompt_encoding', 0.0):.2f}s") - max_logging.log(f" - Flux Denoising: {main_trace.get('denoise_loop', 0.0):.2f}s") - max_logging.log(f" - VAE Decoding: {main_trace.get('vae_decode', 0.0):.2f}s") + max_logging.log(f"1) Model Loading & Placement Time: {load_time:.4f} seconds ā±ļø") + max_logging.log(f"2) Concurrent AOT XLA Compilation Time: {aot_time:.4f} seconds āš”") + max_logging.log(f"3) Warmup Pass Execution Time: {warmup_time:.4f} seconds ā±ļø") + max_logging.log(f" - Qwen3 Encoding: {warmup_trace.get('prompt_encoding', 0.0):.4f}s") + max_logging.log(f" - Flux Denoising: {warmup_trace.get('denoise_loop', 0.0):.4f}s") + max_logging.log(f" - VAE Decoding: {warmup_trace.get('vae_decode', 0.0):.4f}s") + max_logging.log(f"šŸ‘‰ TOTAL COLD-START TIME (Loading + AOT + Warmup): {total_cold_start:.4f} seconds šŸŽÆ") + rep_label = f" (Average across {num_reps} reps)" if num_reps > 1 else "" + max_logging.log(f"4) Main Warmed-Up Pass (Pure Inference Latency){rep_label}: {avg_main_time:.4f} seconds ā±ļø") + max_logging.log(f" - Qwen3 Encoding: {avg_prompt_enc:.4f}s") + max_logging.log(f" - Flux Denoising: {avg_denoise:.4f}s") + max_logging.log(f" - VAE Decoding: {avg_vae_decode:.4f}s") max_logging.log("=" * 80) max_logging.log("\n=======================================================") diff --git a/src/maxdiffusion/models/flux/transformers/generate_flux2klein.py b/src/maxdiffusion/models/flux/transformers/generate_flux2klein.py new file mode 100644 index 000000000..4deebca9f --- /dev/null +++ b/src/maxdiffusion/models/flux/transformers/generate_flux2klein.py @@ -0,0 +1,697 @@ +""" +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. +""" + +import gc +import os +import time +import sys +from typing import List + +from absl import app +import jax +import jax.numpy as jnp +import numpy as np +import flax +from flax import linen as nn +from flax.linen import partitioning as nn_partitioning +from jax.sharding import Mesh + +from maxdiffusion import pyconfig +from maxdiffusion import max_logging +from maxdiffusion import max_utils +from maxdiffusion.max_utils import create_device_mesh +from maxdiffusion.train_utils import transformer_engine_context + +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 +from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights +from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler + + +def partition_prompts(prompt_str: str, batch_size: int) -> List[str]: + """Splits a prompt string by '||' and replicates/truncates to fill the batch_size.""" + raw_prompts = [p.strip() for p in prompt_str.split("||") if p.strip()] + if not raw_prompts: + raw_prompts = ["A detailed vector illustration of a robotic hummingbird"] + + num_prompts = len(raw_prompts) + if num_prompts == 1: + return raw_prompts * batch_size + elif num_prompts <= batch_size: + reps = batch_size // num_prompts + active = [] + for p in raw_prompts: + active.extend([p] * reps) + if len(active) < batch_size: + active.extend([raw_prompts[-1]] * (batch_size - len(active))) + return active + else: + max_logging.log( + f"āš ļø Warning: Found {num_prompts} prompts, but batch_size is {batch_size}. Truncating to the first {batch_size}." + ) + return raw_prompts[:batch_size] + + +def encode_prompt(prompt: str, snapshot_dir: str = None, repo_id: str = "black-forest-labs/FLUX.2-klein-4B"): + """Encodes a prompt string into Qwen3 text embeddings using PyTorch text encoder on CPU.""" + import os + import torch + import gc + from transformers import AutoTokenizer, AutoModelForCausalLM + from huggingface_hub import snapshot_download + + if snapshot_dir is None: + snapshot_dir = snapshot_download(repo_id=repo_id) + + text_encoder_path = os.path.join(snapshot_dir, "text_encoder") + tokenizer_path = os.path.join(snapshot_dir, "tokenizer") + + if not os.path.exists(os.path.join(text_encoder_path, "config.json")) or not os.path.exists(tokenizer_path): + try: + fb_dir = snapshot_download(repo_id=repo_id, local_files_only=True) + if not os.path.exists(os.path.join(text_encoder_path, "config.json")): + text_encoder_path = os.path.join(fb_dir, "text_encoder") + if not os.path.exists(tokenizer_path): + tokenizer_path = ( + os.path.join(fb_dir, "tokenizer") + if os.path.exists(os.path.join(fb_dir, "tokenizer")) + else os.path.join(fb_dir, "text_encoder") + ) + except Exception: + if not os.path.exists(tokenizer_path): + tokenizer_path = text_encoder_path + + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + text_encoder = AutoModelForCausalLM.from_pretrained(text_encoder_path, torch_dtype=torch.float32) + text_encoder.eval() + + messages = [{"role": "user", "content": prompt}] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False) + inputs = tokenizer(text, padding="max_length", max_length=512, truncation=True, return_tensors="pt") + with torch.no_grad(): + outputs = text_encoder(inputs.input_ids, attention_mask=inputs.attention_mask, output_hidden_states=True) + out = torch.stack([outputs.hidden_states[k] for k in (9, 18, 27)], dim=1) + b, c, s, h = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(b, s, c * h) + + del text_encoder + gc.collect() + return prompt_embeds.cpu().numpy() + + +def main(argv): + # Enable shardy partitioner for TPU execution + jax.config.update("jax_use_shardy_partitioner", True) + + # 1. Load configurations + config_path = "src/maxdiffusion/configs/base_flux2klein.yml" + custom_overrides = [] + if len(argv) > 1: + if argv[1].endswith(".yml") or argv[1].endswith(".yaml"): + config_path = argv[1] + if len(argv) > 2: + custom_overrides = argv[2:] + else: + custom_overrides = argv[1:] + + max_logging.log(f"Initializing pyconfig with config: {config_path}") + default_args = [ + None, + config_path, + "run_name=flux2klein_generation", + "output_dir=output/", + ] + default_args.extend(custom_overrides) + + is_interactive = any(arg and "interactive=True" in arg.replace(" ", "") for arg in default_args) + if is_interactive: + max_logging.log("ā„¹ļø Interactive mode detected: overriding use_latents=False for dynamic inputs.") + default_args.append("use_latents=False") + + pyconfig.initialize(default_args) + + # Import modules after jax.distributed.initialize() has run via pyconfig.initialize() + from maxdiffusion.models.flux.util import ( + load_and_convert_flux_klein_weights, + load_and_convert_vae_weights, + ) + from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline + + config = pyconfig.config + os.makedirs(config.output_dir, exist_ok=True) + + if hasattr(config, "per_device_batch_size") and config.per_device_batch_size > 0: + calculated_batch_size = int(config.per_device_batch_size * jax.device_count()) + if calculated_batch_size != config.batch_size: + max_logging.log( + f"ā„¹ļø Updating batch_size from {config.batch_size} to {calculated_batch_size} " + f"based on per_device_batch_size={config.per_device_batch_size} and device_count={jax.device_count()}." + ) + pyconfig._config.keys["batch_size"] = calculated_batch_size + + # 2. Setup device mesh + if ( + config.batch_size == 1 + and config.ici_tensor_parallelism == 1 + and config.ici_context_parallelism == 1 + and jax.device_count() > 1 + ): + max_logging.log( + f"ā„¹ļø Auto-configuring Tensor Parallelism: ici_tensor_parallelism={jax.device_count()}, ici_fsdp_parallelism=1 for batch_size=1 on {jax.device_count()} TPU devices." + ) + pyconfig._config.keys["ici_tensor_parallelism"] = jax.device_count() + pyconfig._config.keys["ici_fsdp_parallelism"] = 1 + + max_logging.log("Setting up JAX device mesh...") + devices_array = create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + + # Check compatibility of batch dimension sharding + data_size = mesh.shape.get("data", 1) + fsdp_size = mesh.shape.get("fsdp", 1) + if config.batch_size % (data_size * fsdp_size) != 0: + max_logging.log( + f"āš ļø Warning: batch_size ({config.batch_size}) is not divisible by FSDP*Data mesh size ({fsdp_size * data_size})." + ) + max_logging.log( + " Automatically falling back to sharding batch dimension across 'data' axis only to prevent JAX SPMD errors." + ) + new_rules = [] + for rule in config.logical_axis_rules: + if rule[0] in ("activation_batch", "conv_batch"): + new_rules.append([rule[0], "data"]) + else: + new_rules.append(rule) + pyconfig._config.keys["logical_axis_rules"] = tuple(new_rules) + + # 3. Resolve weights repository snapshots + repo_id = getattr(config, "pretrained_model_name_or_path", None) + if not repo_id: + depth_val = getattr(config, "depth", None) + repo_id = "black-forest-labs/FLUX.2-klein-9B" if depth_val == 24 else "black-forest-labs/FLUX.2-klein-4B" + max_logging.log(f"Target model detected: {repo_id}") + + if os.path.exists(repo_id): + snapshot_dir = repo_id + max_logging.log(f"Using local model directory: {snapshot_dir}") + else: + from huggingface_hub import snapshot_download + + rev = getattr(config, "revision", None) + if not rev or rev == "refs/pr/95": + rev = "main" + try: + snapshot_dir = snapshot_download(repo_id=repo_id, revision=rev, local_files_only=True) + except Exception: + try: + snapshot_dir = snapshot_download(repo_id=repo_id, local_files_only=True) + except Exception: + snapshot_dir = snapshot_download(repo_id=repo_id) + + max_logging.log(f"Host {jax.process_index()} using HF snapshot directory: {snapshot_dir}") + safetensors_path = os.path.join(snapshot_dir, "transformer") + vae_safetensors_path = os.path.join(snapshot_dir, "vae", "diffusion_pytorch_model.safetensors") + text_encoder_path = os.path.join(snapshot_dir, "text_encoder") + + # 4. Load Qwen3 Config & Setup model layout + from transformers import AutoConfig + + try: + pt_config = AutoConfig.from_pretrained(text_encoder_path, local_files_only=True) + except Exception: + depth_val = getattr(config, "depth", 24) + hf_repo = "black-forest-labs/FLUX.2-klein-9B" if depth_val in (24, -1) else "black-forest-labs/FLUX.2-klein-4B" + max_logging.log(f"ā„¹ļø Config not found in {text_encoder_path}. Resolving from HF cache: {hf_repo}") + pt_config = AutoConfig.from_pretrained(hf_repo, subfolder="text_encoder", local_files_only=True) + + qwen3_config = FlaxQwen3Config( + vocab_size=pt_config.vocab_size, + hidden_size=pt_config.hidden_size, + intermediate_size=pt_config.intermediate_size, + num_hidden_layers=pt_config.num_hidden_layers, + num_attention_heads=pt_config.num_attention_heads, + num_key_value_heads=pt_config.num_key_value_heads, + max_position_embeddings=pt_config.max_position_embeddings, + rms_norm_eps=pt_config.rms_norm_eps, + rope_theta=pt_config.rope_theta, + dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + ) + qwen3_model = FlaxQwen3Model(qwen3_config) + + # Load Transformer HF config.json directly for model architecture parameters + import json + + transformer_config_json = os.path.join(safetensors_path, "config.json") + transformer_pt_cfg = {} + loaded_cfg = False + if os.path.exists(transformer_config_json): + try: + with open(transformer_config_json, "r") as f: + transformer_pt_cfg = json.load(f) + loaded_cfg = True + except Exception as e: + max_logging.log(f"ā„¹ļø Could not parse {transformer_config_json}: {e}. Falling back to HF cache...") + + if not loaded_cfg: + depth_val = getattr(config, "depth", 24) + hf_repo = "black-forest-labs/FLUX.2-klein-9B" if depth_val in (24, -1) else "black-forest-labs/FLUX.2-klein-4B" + try: + from huggingface_hub import hf_hub_download + + cfg_file = hf_hub_download(repo_id=hf_repo, filename="transformer/config.json", local_files_only=True) + with open(cfg_file, "r") as f: + transformer_pt_cfg = json.load(f) + except Exception as e: + max_logging.log(f"āš ļø Warning resolving transformer config fallback: {e}") + + num_double_layers = getattr(config, "num_double_layers", -1) + if num_double_layers is None or num_double_layers <= 0: + num_double_layers = transformer_pt_cfg.get("num_layers", 5) + + depth = getattr(config, "depth", -1) + if depth is None or depth <= 0: + depth = transformer_pt_cfg.get("num_single_layers", 20) + + num_attention_heads = getattr(config, "num_attention_heads", -1) + if num_attention_heads is None or num_attention_heads <= 0: + num_attention_heads = transformer_pt_cfg.get("num_attention_heads", 24) + + # 5. Instantiate JAX Flux2KleinTransformer2DModel + transformer = Flux2KleinTransformer2DModel( + in_channels=128, + num_layers=num_double_layers, + num_single_layers=depth, + attention_head_dim=128, + num_attention_heads=num_attention_heads, + joint_attention_dim=3 * pt_config.hidden_size, + pooled_projection_dim=768, + mlp_ratio=3.0, + qkv_bias=False, + joint_attention_bias=False, + x_embedder_bias=False, + proj_out_bias=False, + use_global_modulation=True, + use_swiglu=True, + axes_dims_rope=(32, 32, 32, 32), + theta=2000, + mesh=mesh, + dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + weights_dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + attention_kernel=config.attention, + ulysses_shards=getattr(config, "ulysses_shards", -1), + scale_shift_order=getattr(config, "scale_shift_order", "shift_scale"), + ) + + # 6. Instantiate JAX VAE + vae = FlaxAutoencoderKL( + in_channels=3, + out_channels=3, + down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"), + up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"), + block_out_channels=(128, 256, 512, 512), + layers_per_block=2, + act_fn="silu", + latent_channels=32, + norm_num_groups=32, + sample_size=512, + use_quant_conv=True, + use_post_quant_conv=True, + dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + ) + + # 7. Evaluate shapes & extract mesh shardings + max_logging.log("Evaluating model shapes and shardings...") + h_packed = config.height // 16 + w_packed = config.width // 16 + seq_len_img = h_packed * w_packed + seq_len_txt = config.max_sequence_length + + img_dummy = jnp.zeros((config.batch_size, seq_len_img, 128)) + img_ids_dummy = jnp.zeros((config.batch_size, seq_len_img, 4)) + txt_dummy = jnp.zeros((config.batch_size, seq_len_txt, 3 * pt_config.hidden_size)) + txt_ids_dummy = jnp.zeros((config.batch_size, seq_len_txt, 4)) + vec_dummy = jnp.zeros((config.batch_size, 768)) + t_vec_dummy = jnp.zeros((config.batch_size,)) + guidance_vec_dummy = jnp.zeros((config.batch_size,)) + dummy_img = jnp.zeros((config.batch_size, 3, 512, 512)) + dummy_ids = jnp.zeros((config.batch_size, seq_len_txt), dtype=jnp.int32) + dummy_mask = jnp.zeros((config.batch_size, seq_len_txt), dtype=jnp.int32) + + key = jax.random.PRNGKey(0) + key, vae_key, qwen_key = jax.random.split(key, 3) + + def transformer_init_fn(): + return transformer.init( + key, + hidden_states=img_dummy, + img_ids=img_ids_dummy, + encoder_hidden_states=txt_dummy, + txt_ids=txt_ids_dummy, + pooled_projections=vec_dummy, + timestep=t_vec_dummy, + guidance=guidance_vec_dummy, + ) + + def vae_init_fn(): + return vae.init(vae_key, dummy_img) + + def qwen3_init_fn(): + return qwen3_model.init(qwen_key, dummy_ids, dummy_mask) + + with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): + abstract_transformer_vars = jax.eval_shape(transformer_init_fn) + abstract_vae_vars = jax.eval_shape(vae_init_fn) + abstract_qwen3_vars = jax.eval_shape(qwen3_init_fn) + + logical_transformer_specs = nn.get_partition_spec(abstract_transformer_vars) + logical_vae_specs = nn.get_partition_spec(abstract_vae_vars) + logical_qwen3_specs = nn.get_partition_spec(abstract_qwen3_vars) + + transformer_mesh_shardings = nn.logical_to_mesh_sharding(logical_transformer_specs, mesh, config.logical_axis_rules) + vae_mesh_shardings = nn.logical_to_mesh_sharding(logical_vae_specs, mesh, config.logical_axis_rules) + qwen3_mesh_shardings = nn.logical_to_mesh_sharding(logical_qwen3_specs, mesh, config.logical_axis_rules) + + transformer_shardings = flax.core.freeze(transformer_mesh_shardings["params"]) + vae_shardings = flax.core.freeze(vae_mesh_shardings["params"]) + qwen3_shardings = flax.core.freeze(qwen3_mesh_shardings["params"]) + + # 8. Load weights on Host CPU + max_logging.log("Loading parameters on Host CPU...") + t_load_start = time.time() + cpu_device = jax.local_devices(backend="cpu")[0] + with jax.default_device(cpu_device): + with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): + import flax.linen.spmd as flax_spmd + + def unbox_fn(x): + return x.unbox() if isinstance(x, flax_spmd.LogicallyPartitioned) else x + + t_sub0 = time.time() + params = jax.tree_util.tree_map( + unbox_fn, abstract_transformer_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned) + ) + params = flax.core.unfreeze(params) + + vae_params = jax.tree_util.tree_map( + unbox_fn, abstract_vae_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned) + ) + vae_params = flax.core.unfreeze(vae_params) + + qwen3_params = jax.tree_util.tree_map( + unbox_fn, abstract_qwen3_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned) + ) + qwen3_params = flax.core.unfreeze(qwen3_params) + + max_logging.log(f" -> [SUB-TIMING 1/3] PyTree unboxing template setup: {time.time() - t_sub0:.2f}s") + t_sub1 = time.time() + + weight_dtype = jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32 + + params = load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, depth, dtype=weight_dtype) + vae_params, vae_bn_mean, vae_bn_std = load_and_convert_vae_weights( + vae_safetensors_path, vae_params, dtype=weight_dtype + ) + qwen3_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params, qwen3_config) + max_logging.log( + f" -> [SUB-TIMING 2/3] Safetensors loading & key mapping (in target dtype): {time.time() - t_sub1:.4f}s" + ) + + params = flax.core.freeze(params) + vae_params = flax.core.freeze(vae_params) + qwen3_params = flax.core.freeze(qwen3_params) + + max_logging.log("\n" + "=" * 80) + max_logging.log("šŸš€ Pinning all parameters to TPU HBM permanently...") + max_logging.log("=" * 80 + "\n") + t_sub3 = time.time() + max_logging.log("Putting params on TPU HBM...") + with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): + try: + params = jax.tree_util.tree_map(max_utils.device_put_replicated, params, transformer_shardings) + except Exception as err: + max_logging.log("\nāŒ jax.device_put(params, transformer_shardings) FAILED!") + flat_p = flax.traverse_util.flatten_dict(params) + flat_s = flax.traverse_util.flatten_dict(transformer_shardings) + k_p = set(flat_p.keys()) + k_s = set(flat_s.keys()) + max_logging.log(f"Keys in sharding spec but missing in params: {k_s - k_p}") + max_logging.log(f"Keys in params but missing in sharding spec: {k_p - k_s}") + sys.stdout.flush() + raise err + max_logging.log("Putting vae_params on TPU HBM...") + vae_params = jax.tree_util.tree_map(max_utils.device_put_replicated, vae_params, vae_shardings) + max_logging.log("Putting qwen3_params on TPU HBM...") + qwen3_params = jax.tree_util.tree_map(max_utils.device_put_replicated, qwen3_params, qwen3_shardings) + max_logging.log(f" -> [SUB-TIMING 3/3] TPU HBM device_put placement: {time.time() - t_sub3:.4f}s") + max_logging.log("All parameters placed on TPU HBM successfully!") + gc.collect() + jax.effects_barrier() + + load_time = time.time() - t_load_start + max_logging.log(f" -> [TIMING] Total Model Loading & Device Placement: {load_time:.4f} seconds ā±ļø\n") + + # 9. Setup FlowMatch Scheduler + scheduler = FlaxFlowMatchScheduler( + num_train_timesteps=1000, + shift=1.0, + sigma_max=1.0, + sigma_min=0.001, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False, + use_dynamic_shifting=True, + time_shift_type="exponential", + ) + + # 10. Instantiate and invoke FlaxFlux2KleinPipeline + max_logging.log("Instantiating JAX FlaxFlux2KleinPipeline...") + pipeline = FlaxFlux2KleinPipeline( + transformer=transformer, + vae=vae, + text_encoder=qwen3_model, + tokenizer=None, + scheduler=scheduler, + config=config, + mesh=mesh, + ) + + prompt_str = getattr(config, "prompt", "") or "A dog running in a field with butterflies and tall grass" + active_prompts = partition_prompts(prompt_str, config.batch_size) + + if getattr(config, "interactive", False): + max_logging.log("\n" + "=" * 80) + max_logging.log(" BATCHED INTERACTIVE GENERATION MODE ENABLED šŸŽ®") + max_logging.log("The model has been fully loaded and compiled on the TPU.") + max_logging.log(f"Batch size: {config.batch_size} parallel images.") + max_logging.log("Enter prompts separated by '||' (e.g. A cute cat || A red car)") + max_logging.log("Type 'exit' to quit.") + max_logging.log("=" * 80) + + image_idx = 1 + while True: + try: + user_input = input("\nEnter prompt(s): ") + except (KeyboardInterrupt, EOFError): + break + if user_input.strip().lower() in ("exit", "quit"): + break + if not user_input.strip(): + continue + + prompts = partition_prompts(user_input, config.batch_size) + output_file = f"generated_{image_idx:03d}.png" + + pipeline( + prompt=prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=False, + output_dir=config.output_dir, + output_name=output_file, + ) + image_idx += 1 + else: + # Run one-shot generation + latents_to_use = None + use_latents_flag = False + if getattr(config, "latents_path", ""): + max_logging.log(f"Loading custom starting noise latents from: {config.latents_path}...") + latents_to_use = np.load(config.latents_path) + use_latents_flag = True + max_logging.log(f" -> Custom latents shape: {latents_to_use.shape} | sum: {latents_to_use.sum():.6f}") + + max_logging.log("\n" + "=" * 80) + max_logging.log("šŸš€ Pre-compiling XLA graphs concurrently (AOT Compilation)...") + max_logging.log("=" * 80) + aot_time = pipeline.compile_aot_async( + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + batch_size=config.batch_size, + height=config.height, + width=config.width, + ) + + max_logging.log("\n" + "=" * 80) + max_logging.log("šŸš€ Running initial dry run (Warmup Pass) to verify compiled graph execution...") + max_logging.log("=" * 80) + _, warmup_trace = pipeline( + prompt=active_prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=use_latents_flag, + latents=latents_to_use, + output_dir=config.output_dir, + output_name="flux2klein_warmup.png", + warmup=True, + ) + warmup_time = ( + warmup_trace.get("prompt_encoding", 0.0) + + warmup_trace.get("denoise_loop", 0.0) + + warmup_trace.get("vae_decode", 0.0) + ) + + num_reps = int(getattr(config, "num_reps", 1)) + max_logging.log("\n" + "=" * 80) + max_logging.log(f"ā±ļø Running timed pass at full TPU speed (num_reps={num_reps})...") + max_logging.log("=" * 80) + + main_traces = [] + main_times = [] + + for rep in range(num_reps): + rep_str = f" [Rep {rep+1}/{num_reps}]" if num_reps > 1 else "" + if rep > 0: + max_logging.log(f"ā±ļø Running timed pass{rep_str}...") + + if max_utils.profiler_enabled(config) and rep == 0: + max_logging.log(f"šŸš€ XProf / JAX Profiler active! Capturing trace into: {config.tensorboard_dir}") + with max_utils.Profiler(config, session_name="flux2klein_inference"): + _, trace_i = pipeline( + prompt=active_prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=use_latents_flag, + latents=latents_to_use, + output_dir=config.output_dir, + output_name=f"rep_{rep+1}_{config.output_name}" if num_reps > 1 else config.output_name, + ) + else: + _, trace_i = pipeline( + prompt=active_prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=use_latents_flag, + latents=latents_to_use, + output_dir=config.output_dir, + output_name=f"rep_{rep+1}_{config.output_name}" if num_reps > 1 else config.output_name, + ) + + tot_time_i = trace_i.get( + "e2e_pipeline_total", + trace_i.get("prompt_encoding", 0.0) + trace_i.get("denoise_loop", 0.0) + trace_i.get("vae_decode", 0.0), + ) + main_traces.append(trace_i) + main_times.append(tot_time_i) + if num_reps > 1: + max_logging.log( + f" -> Rep {rep+1}/{num_reps} Completed: Total={tot_time_i:.4f}s | Qwen3={trace_i.get('qwen3_encoding', 0.0):.4f}s | Denoise={trace_i.get('denoise_loop', 0.0):.4f}s | VAE={trace_i.get('vae_decode', 0.0):.4f}s" + ) + + avg_main_time = sum(main_times) / num_reps + avg_start_to_qwen3 = sum(tr.get("start_to_qwen3", 0.0) for tr in main_traces) / num_reps + avg_prompt_enc = sum(tr.get("qwen3_encoding", tr.get("prompt_encoding", 0.0)) for tr in main_traces) / num_reps + avg_qwen3_to_denoise = sum(tr.get("qwen3_to_denoise", 0.0) for tr in main_traces) / num_reps + avg_denoise = sum(tr.get("denoise_loop", 0.0) for tr in main_traces) / num_reps + avg_denoise_to_vae = sum(tr.get("denoise_to_vae", 0.0) for tr in main_traces) / num_reps + avg_vae_decode = sum(tr.get("vae_decode", 0.0) for tr in main_traces) / num_reps + avg_image_saving = sum(tr.get("image_saving", 0.0) for tr in main_traces) / num_reps + + total_cold_start = load_time + aot_time + warmup_time + + max_logging.log("\n" + "=" * 80) + max_logging.log("šŸ“Š FLUX.2-KLEIN COMPLETE LATENCY & TIMING BREAKDOWN") + max_logging.log("=" * 80) + max_logging.log(f"1) Model Loading & Placement Time: {load_time:.4f} seconds ā±ļø") + max_logging.log(f"2) Concurrent AOT XLA Compilation Time: {aot_time:.4f} seconds āš”") + max_logging.log(f"3) Warmup Pass Execution Time: {warmup_time:.4f} seconds ā±ļø") + max_logging.log(f" - Qwen3 Encoding: {warmup_trace.get('prompt_encoding', 0.0):.4f}s") + max_logging.log(f" - Flux Denoising: {warmup_trace.get('denoise_loop', 0.0):.4f}s") + max_logging.log(f" - VAE Decoding: {warmup_trace.get('vae_decode', 0.0):.4f}s") + max_logging.log(f"šŸ‘‰ TOTAL COLD-START TIME (Loading + AOT + Warmup): {total_cold_start:.4f} seconds šŸŽÆ") + rep_label = f" (Average across {num_reps} reps)" if num_reps > 1 else "" + max_logging.log(f"4) Main Warmed-Up Pass (Pure Inference Latency){rep_label}: {avg_main_time:.4f} seconds ā±ļø") + max_logging.log(f" - 1. Start -> Qwen3: {avg_start_to_qwen3*1000:.2f} ms ({avg_start_to_qwen3:.4f}s)") + max_logging.log(f" - 2. Qwen3 Encoding: {avg_prompt_enc*1000:.2f} ms ({avg_prompt_enc:.4f}s)") + max_logging.log(f" - 3. Qwen3 -> Denoising: {avg_qwen3_to_denoise*1000:.2f} ms ({avg_qwen3_to_denoise:.4f}s)") + max_logging.log(f" - 4. Flux Denoising Loop: {avg_denoise*1000:.2f} ms ({avg_denoise:.4f}s)") + max_logging.log(f" - 5. Denoising -> VAE: {avg_denoise_to_vae*1000:.2f} ms ({avg_denoise_to_vae:.4f}s)") + max_logging.log(f" - 6. VAE Decoding: {avg_vae_decode*1000:.2f} ms ({avg_vae_decode:.4f}s)") + max_logging.log(f" - 7. Image Saving: {avg_image_saving*1000:.2f} ms ({avg_image_saving:.4f}s)") + max_logging.log(f" - šŸ‘‰ TOTAL E2E PIPELINE: {avg_main_time*1000:.2f} ms ({avg_main_time:.4f}s)") + max_logging.log("=" * 80) + + max_logging.log("\n=======================================================") + max_logging.log(f"SUCCESS! Batched generation complete for {config.batch_size} images! šŸŽØšŸŽ‰") + max_logging.log("=======================================================\n") + + +if __name__ == "__main__": + with transformer_engine_context(): + app.run(main) diff --git a/src/maxdiffusion/models/flux/util.py b/src/maxdiffusion/models/flux/util.py index 952519776..f7567e8cb 100644 --- a/src/maxdiffusion/models/flux/util.py +++ b/src/maxdiffusion/models/flux/util.py @@ -300,17 +300,17 @@ def unpack_latents(latents, batch_size, num_channels_latents, height, width): Unpacks packed sequence of shape (batch_size, (height//16)*(width//16), channels*4) back to the unpacked spatial grid shape (batch_size, channels, height//8, width//8). """ - import numpy as np + import jax.numpy as jnp h_latent = height // 8 w_latent = width // 8 # 1. Reshape to split spatial grid and packed channel blocks - latents = np.reshape(latents, (batch_size, h_latent // 2, w_latent // 2, num_channels_latents, 2, 2)) + latents = jnp.reshape(latents, (batch_size, h_latent // 2, w_latent // 2, num_channels_latents, 2, 2)) # 2. Permute dimensions back to unpacked order - latents = np.transpose(latents, (0, 3, 1, 4, 2, 5)) + latents = jnp.transpose(latents, (0, 3, 1, 4, 2, 5)) # 3. Flatten back to 4D unpacked latent shape - latents = np.reshape(latents, (batch_size, num_channels_latents, h_latent, w_latent)) + latents = jnp.reshape(latents, (batch_size, num_channels_latents, h_latent, w_latent)) return latents @@ -398,11 +398,12 @@ def cast_dict_to_bfloat16_inplace(d, device=None, exclude_keywords=None, parent_ is_excluded = exclude_keywords and any(kw.lower() in current_key.lower() for kw in exclude_keywords) target_dtype = jnp.float32 if is_excluded else jnp.bfloat16 - d[k] = v.astype(target_dtype) - if hasattr(d[k], "block_until_ready"): - d[k].block_until_ready() - del v - gc.collect() + if v.dtype != target_dtype: + d[k] = v.astype(target_dtype) + if hasattr(d[k], "block_until_ready"): + d[k].block_until_ready() + del v + gc.collect() # ----------------------------------------------------------------------------- @@ -410,7 +411,9 @@ def cast_dict_to_bfloat16_inplace(d, device=None, exclude_keywords=None, parent_ # ----------------------------------------------------------------------------- -def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, num_single_layers): +def load_and_convert_flux_klein_weights( + safetensors_path, params, num_double_layers, num_single_layers, dtype=None, pt_state_dict=None +): """ Loads weights from safetensors via zero-copy safetensors.numpy and converts them to JAX parameter dictionary. Supports dynamic layer counts (double and single stream blocks) and sharded safetensors directories. @@ -422,28 +425,30 @@ def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_lay import os import gc - pt_state_dict = {} - if os.path.isdir(safetensors_path): - shards = glob.glob(os.path.join(safetensors_path, "*.safetensors")) - max_logging.log(f"Loading sharded weights from directory: {safetensors_path} (Found {len(shards)} shards)...") - for shard in sorted(shards): - max_logging.log(f"Loading shard: {shard}...") - pt_state_dict.update(load_file(shard)) - else: - max_logging.log(f"Loading weights from: {safetensors_path}") - pt_state_dict = load_file(safetensors_path) + if pt_state_dict is None: + pt_state_dict = {} + if os.path.isdir(safetensors_path): + shards = glob.glob(os.path.join(safetensors_path, "*.safetensors")) + max_logging.log(f"Loading sharded weights from directory: {safetensors_path} (Found {len(shards)} shards)...") + for shard in sorted(shards): + max_logging.log(f"Loading shard: {shard}...") + pt_state_dict.update(load_file(shard)) + else: + max_logging.log(f"Loading weights from: {safetensors_path}") + pt_state_dict = load_file(safetensors_path) max_logging.log("Mapping weights to JAX parameters...") expected_pytree = jax.tree_util.tree_map(lambda leaf: leaf, params) first_leaf = jax.tree_util.tree_leaves(params)[0] - target_dtype = first_leaf.dtype + target_dtype = dtype if dtype is not None else first_leaf.dtype - def convert_and_transpose_tensor(tensor, transpose=False): + def convert_and_transpose_tensor(tensor, transpose=False, is_norm=False): if transpose and len(tensor.shape) == 2: tensor = tensor.T - return jnp.array(tensor, dtype=target_dtype) + leaf_dtype = jnp.float32 if is_norm else target_dtype + return jnp.array(tensor, dtype=leaf_dtype) # Global layers params["context_embedder"]["kernel"] = convert_and_transpose_tensor( @@ -562,21 +567,28 @@ def convert_and_transpose_tensor(tensor, transpose=False): return params -def load_and_convert_vae_weights(safetensors_path, jax_params): +def load_and_convert_vae_weights(safetensors_path, jax_params, dtype=None, pt_state_dict=None): """Loads VAE weights from safetensors via zero-copy safetensors.numpy, maps them to JAX, and extracts BN stats.""" from safetensors.numpy import load_file import flax import jax.numpy as jnp - max_logging.log(f"Loading VAE weights from: {safetensors_path}") - pt_state_dict = load_file(safetensors_path) - - def get_pytorch_weight_tensor(key): - return pt_state_dict[key] + if pt_state_dict is None: + max_logging.log(f"Loading VAE weights from: {safetensors_path}") + pt_state_dict = load_file(safetensors_path) # Unfreeze JAX params so we can load the weights jax_params = flax.core.unfreeze(jax_params) + first_leaf = jax.tree_util.tree_leaves(jax_params)[0] + target_dtype = dtype if dtype is not None else first_leaf.dtype + + def get_pytorch_weight_tensor(key, dtype_val=target_dtype): + tensor = pt_state_dict[key] + is_norm = any(kw in key.lower() for kw in ("norm", "layernorm", "rmsnorm", "groupnorm")) + leaf_dtype = jnp.float32 if is_norm else dtype_val + return jnp.array(tensor, dtype=leaf_dtype) + # Map weights max_logging.log("Mapping VAE decoder weights to JAX parameters...") diff --git a/src/maxdiffusion/models/generate_flux2klein.py b/src/maxdiffusion/models/generate_flux2klein.py new file mode 100644 index 000000000..2e869aeb0 --- /dev/null +++ b/src/maxdiffusion/models/generate_flux2klein.py @@ -0,0 +1,712 @@ +""" +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. +""" + +import gc +import os +import time +import sys +from typing import List + +from absl import app +import jax +import jax.numpy as jnp +import numpy as np +import flax +from flax import linen as nn +from flax.linen import partitioning as nn_partitioning +from jax.sharding import Mesh + +from maxdiffusion import pyconfig +from maxdiffusion import max_logging +from maxdiffusion import max_utils +from maxdiffusion.max_utils import create_device_mesh +from maxdiffusion.train_utils import transformer_engine_context + +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 +from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights +from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler + + +def partition_prompts(prompt_str: str, batch_size: int) -> List[str]: + """Splits a prompt string by '||' and replicates/truncates to fill the batch_size.""" + raw_prompts = [p.strip() for p in prompt_str.split("||") if p.strip()] + if not raw_prompts: + raw_prompts = ["A detailed vector illustration of a robotic hummingbird"] + + num_prompts = len(raw_prompts) + if num_prompts == 1: + return raw_prompts * batch_size + elif num_prompts <= batch_size: + reps = batch_size // num_prompts + active = [] + for p in raw_prompts: + active.extend([p] * reps) + if len(active) < batch_size: + active.extend([raw_prompts[-1]] * (batch_size - len(active))) + return active + else: + max_logging.log( + f"āš ļø Warning: Found {num_prompts} prompts, but batch_size is {batch_size}. Truncating to the first {batch_size}." + ) + return raw_prompts[:batch_size] + + +def encode_prompt(prompt: str, snapshot_dir: str = None, repo_id: str = "black-forest-labs/FLUX.2-klein-4B"): + """Encodes a prompt string into Qwen3 text embeddings using PyTorch text encoder on CPU.""" + import os + import torch + import gc + from transformers import AutoTokenizer, AutoModelForCausalLM + from huggingface_hub import snapshot_download + + if snapshot_dir is None: + snapshot_dir = snapshot_download(repo_id=repo_id) + + text_encoder_path = os.path.join(snapshot_dir, "text_encoder") + tokenizer_path = os.path.join(snapshot_dir, "tokenizer") + + if not os.path.exists(os.path.join(text_encoder_path, "config.json")) or not os.path.exists(tokenizer_path): + try: + fb_dir = snapshot_download(repo_id=repo_id, local_files_only=True) + if not os.path.exists(os.path.join(text_encoder_path, "config.json")): + text_encoder_path = os.path.join(fb_dir, "text_encoder") + if not os.path.exists(tokenizer_path): + tokenizer_path = ( + os.path.join(fb_dir, "tokenizer") + if os.path.exists(os.path.join(fb_dir, "tokenizer")) + else os.path.join(fb_dir, "text_encoder") + ) + except Exception: + if not os.path.exists(tokenizer_path): + tokenizer_path = text_encoder_path + + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + text_encoder = AutoModelForCausalLM.from_pretrained(text_encoder_path, torch_dtype=torch.float32) + text_encoder.eval() + + messages = [{"role": "user", "content": prompt}] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False) + inputs = tokenizer(text, padding="max_length", max_length=512, truncation=True, return_tensors="pt") + with torch.no_grad(): + outputs = text_encoder(inputs.input_ids, attention_mask=inputs.attention_mask, output_hidden_states=True) + out = torch.stack([outputs.hidden_states[k] for k in (9, 18, 27)], dim=1) + b, c, s, h = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(b, s, c * h) + + del text_encoder + gc.collect() + return prompt_embeds.cpu().numpy() + + +def main(argv): + # Enable shardy partitioner for TPU execution + jax.config.update("jax_use_shardy_partitioner", True) + + # 1. Load configurations + config_path = "src/maxdiffusion/configs/base_flux2klein.yml" + custom_overrides = [] + if len(argv) > 1: + if argv[1].endswith(".yml") or argv[1].endswith(".yaml"): + config_path = argv[1] + if len(argv) > 2: + custom_overrides = argv[2:] + else: + custom_overrides = argv[1:] + + max_logging.log(f"Initializing pyconfig with config: {config_path}") + default_args = [ + None, + config_path, + "run_name=flux2klein_generation", + "output_dir=output/", + ] + default_args.extend(custom_overrides) + + is_interactive = any(arg and "interactive=True" in arg.replace(" ", "") for arg in default_args) + if is_interactive: + max_logging.log("ā„¹ļø Interactive mode detected: overriding use_latents=False for dynamic inputs.") + default_args.append("use_latents=False") + + pyconfig.initialize(default_args) + + # Import modules after jax.distributed.initialize() has run via pyconfig.initialize() + from maxdiffusion.models.flux.util import ( + load_and_convert_flux_klein_weights, + load_and_convert_vae_weights, + ) + from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline + + config = pyconfig.config + os.makedirs(config.output_dir, exist_ok=True) + + if hasattr(config, "per_device_batch_size") and config.per_device_batch_size > 0: + calculated_batch_size = int(config.per_device_batch_size * jax.device_count()) + if calculated_batch_size != config.batch_size: + max_logging.log( + f"ā„¹ļø Updating batch_size from {config.batch_size} to {calculated_batch_size} " + f"based on per_device_batch_size={config.per_device_batch_size} and device_count={jax.device_count()}." + ) + pyconfig._config.keys["batch_size"] = calculated_batch_size + + # 2. Setup device mesh + if ( + config.batch_size == 1 + and config.ici_tensor_parallelism == 1 + and config.ici_context_parallelism == 1 + and jax.device_count() > 1 + ): + max_logging.log( + f"ā„¹ļø Auto-configuring Tensor Parallelism: ici_tensor_parallelism={jax.device_count()}, ici_fsdp_parallelism=1 for batch_size=1 on {jax.device_count()} TPU devices." + ) + pyconfig._config.keys["ici_tensor_parallelism"] = jax.device_count() + pyconfig._config.keys["ici_fsdp_parallelism"] = 1 + + max_logging.log("Setting up JAX device mesh...") + devices_array = create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + + # Check compatibility of batch dimension sharding + data_size = mesh.shape.get("data", 1) + fsdp_size = mesh.shape.get("fsdp", 1) + if config.batch_size % (data_size * fsdp_size) != 0: + max_logging.log( + f"āš ļø Warning: batch_size ({config.batch_size}) is not divisible by FSDP*Data mesh size ({fsdp_size * data_size})." + ) + max_logging.log( + " Automatically falling back to sharding batch dimension across 'data' axis only to prevent JAX SPMD errors." + ) + new_rules = [] + for rule in config.logical_axis_rules: + if rule[0] in ("activation_batch", "conv_batch"): + new_rules.append([rule[0], "data"]) + else: + new_rules.append(rule) + pyconfig._config.keys["logical_axis_rules"] = tuple(new_rules) + + # 3. Resolve weights repository snapshots + repo_id = getattr(config, "pretrained_model_name_or_path", None) + if not repo_id: + depth_val = getattr(config, "depth", None) + repo_id = "black-forest-labs/FLUX.2-klein-9B" if depth_val == 24 else "black-forest-labs/FLUX.2-klein-4B" + max_logging.log(f"Target model detected: {repo_id}") + + if os.path.exists(repo_id): + snapshot_dir = repo_id + max_logging.log(f"Using local model directory: {snapshot_dir}") + else: + from huggingface_hub import snapshot_download + + rev = getattr(config, "revision", None) + if not rev or rev == "refs/pr/95": + rev = "main" + try: + snapshot_dir = snapshot_download(repo_id=repo_id, revision=rev, local_files_only=True) + except Exception: + try: + snapshot_dir = snapshot_download(repo_id=repo_id, local_files_only=True) + except Exception: + snapshot_dir = snapshot_download(repo_id=repo_id) + + max_logging.log(f"Host {jax.process_index()} using HF snapshot directory: {snapshot_dir}") + safetensors_path = os.path.join(snapshot_dir, "transformer") + vae_safetensors_path = os.path.join(snapshot_dir, "vae", "diffusion_pytorch_model.safetensors") + text_encoder_path = os.path.join(snapshot_dir, "text_encoder") + + # 4. Load Qwen3 Config & Setup model layout + from transformers import AutoConfig + + try: + pt_config = AutoConfig.from_pretrained(text_encoder_path, local_files_only=True) + except Exception: + depth_val = getattr(config, "depth", 24) + hf_repo = "black-forest-labs/FLUX.2-klein-9B" if depth_val in (24, -1) else "black-forest-labs/FLUX.2-klein-4B" + max_logging.log(f"ā„¹ļø Config not found in {text_encoder_path}. Resolving from HF cache: {hf_repo}") + pt_config = AutoConfig.from_pretrained(hf_repo, subfolder="text_encoder", local_files_only=True) + + text_encoder_attn = getattr(config, "text_encoder_attention", "dot_product") + text_encoder_block_sizes = getattr(config, "text_encoder_flash_block_sizes", None) + if isinstance(text_encoder_block_sizes, str) and text_encoder_block_sizes: + import ast + + try: + text_encoder_block_sizes = ast.literal_eval(text_encoder_block_sizes) + except Exception: + pass + ulysses_shards_val = getattr(config, "ulysses_shards", -1) + + qwen3_config = FlaxQwen3Config( + vocab_size=pt_config.vocab_size, + hidden_size=pt_config.hidden_size, + intermediate_size=pt_config.intermediate_size, + num_hidden_layers=pt_config.num_hidden_layers, + num_attention_heads=pt_config.num_attention_heads, + num_key_value_heads=pt_config.num_key_value_heads, + max_position_embeddings=pt_config.max_position_embeddings, + rms_norm_eps=pt_config.rms_norm_eps, + rope_theta=pt_config.rope_theta, + dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + mesh=mesh, + attention_kernel=text_encoder_attn, + ulysses_shards=ulysses_shards_val, + flash_block_sizes=text_encoder_block_sizes, + ) + qwen3_model = FlaxQwen3Model(qwen3_config) + + # Load Transformer HF config.json directly for model architecture parameters + import json + + transformer_config_json = os.path.join(safetensors_path, "config.json") + transformer_pt_cfg = {} + loaded_cfg = False + if os.path.exists(transformer_config_json): + try: + with open(transformer_config_json, "r") as f: + transformer_pt_cfg = json.load(f) + loaded_cfg = True + except Exception as e: + max_logging.log(f"ā„¹ļø Could not parse {transformer_config_json}: {e}. Falling back to HF cache...") + + if not loaded_cfg: + depth_val = getattr(config, "depth", 24) + hf_repo = "black-forest-labs/FLUX.2-klein-9B" if depth_val in (24, -1) else "black-forest-labs/FLUX.2-klein-4B" + try: + from huggingface_hub import hf_hub_download + + cfg_file = hf_hub_download(repo_id=hf_repo, filename="transformer/config.json", local_files_only=True) + with open(cfg_file, "r") as f: + transformer_pt_cfg = json.load(f) + except Exception as e: + max_logging.log(f"āš ļø Warning resolving transformer config fallback: {e}") + + num_double_layers = getattr(config, "num_double_layers", -1) + if num_double_layers is None or num_double_layers <= 0: + num_double_layers = transformer_pt_cfg.get("num_layers", 5) + + depth = getattr(config, "depth", -1) + if depth is None or depth <= 0: + depth = transformer_pt_cfg.get("num_single_layers", 20) + + num_attention_heads = getattr(config, "num_attention_heads", -1) + if num_attention_heads is None or num_attention_heads <= 0: + num_attention_heads = transformer_pt_cfg.get("num_attention_heads", 24) + + # 5. Instantiate JAX Flux2KleinTransformer2DModel + transformer = Flux2KleinTransformer2DModel( + in_channels=128, + num_layers=num_double_layers, + num_single_layers=depth, + attention_head_dim=128, + num_attention_heads=num_attention_heads, + joint_attention_dim=3 * pt_config.hidden_size, + pooled_projection_dim=768, + mlp_ratio=3.0, + qkv_bias=False, + joint_attention_bias=False, + x_embedder_bias=False, + proj_out_bias=False, + use_global_modulation=True, + use_swiglu=True, + axes_dims_rope=(32, 32, 32, 32), + theta=2000, + mesh=mesh, + dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + weights_dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + attention_kernel=config.attention, + ulysses_shards=getattr(config, "ulysses_shards", -1), + scale_shift_order=getattr(config, "scale_shift_order", "shift_scale"), + ) + + # 6. Instantiate JAX VAE + vae = FlaxAutoencoderKL( + in_channels=3, + out_channels=3, + down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"), + up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"), + block_out_channels=(128, 256, 512, 512), + layers_per_block=2, + act_fn="silu", + latent_channels=32, + norm_num_groups=32, + sample_size=512, + use_quant_conv=True, + use_post_quant_conv=True, + dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + ) + + # 7. Evaluate shapes & extract mesh shardings + max_logging.log("Evaluating model shapes and shardings...") + h_packed = config.height // 16 + w_packed = config.width // 16 + seq_len_img = h_packed * w_packed + seq_len_txt = config.max_sequence_length + + img_dummy = jnp.zeros((config.batch_size, seq_len_img, 128)) + img_ids_dummy = jnp.zeros((config.batch_size, seq_len_img, 4)) + txt_dummy = jnp.zeros((config.batch_size, seq_len_txt, 3 * pt_config.hidden_size)) + txt_ids_dummy = jnp.zeros((config.batch_size, seq_len_txt, 4)) + vec_dummy = jnp.zeros((config.batch_size, 768)) + t_vec_dummy = jnp.zeros((config.batch_size,)) + guidance_vec_dummy = jnp.zeros((config.batch_size,)) + dummy_img = jnp.zeros((config.batch_size, 3, 512, 512)) + dummy_ids = jnp.zeros((config.batch_size, seq_len_txt), dtype=jnp.int32) + dummy_mask = jnp.zeros((config.batch_size, seq_len_txt), dtype=jnp.int32) + + key = jax.random.PRNGKey(0) + key, vae_key, qwen_key = jax.random.split(key, 3) + + def transformer_init_fn(): + return transformer.init( + key, + hidden_states=img_dummy, + img_ids=img_ids_dummy, + encoder_hidden_states=txt_dummy, + txt_ids=txt_ids_dummy, + pooled_projections=vec_dummy, + timestep=t_vec_dummy, + guidance=guidance_vec_dummy, + ) + + def vae_init_fn(): + return vae.init(vae_key, dummy_img) + + def qwen3_init_fn(): + return qwen3_model.init(qwen_key, dummy_ids, dummy_mask) + + with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): + abstract_transformer_vars = jax.eval_shape(transformer_init_fn) + abstract_vae_vars = jax.eval_shape(vae_init_fn) + abstract_qwen3_vars = jax.eval_shape(qwen3_init_fn) + + logical_transformer_specs = nn.get_partition_spec(abstract_transformer_vars) + logical_vae_specs = nn.get_partition_spec(abstract_vae_vars) + logical_qwen3_specs = nn.get_partition_spec(abstract_qwen3_vars) + + transformer_mesh_shardings = nn.logical_to_mesh_sharding(logical_transformer_specs, mesh, config.logical_axis_rules) + vae_mesh_shardings = nn.logical_to_mesh_sharding(logical_vae_specs, mesh, config.logical_axis_rules) + qwen3_mesh_shardings = nn.logical_to_mesh_sharding(logical_qwen3_specs, mesh, config.logical_axis_rules) + + transformer_shardings = flax.core.freeze(transformer_mesh_shardings["params"]) + vae_shardings = flax.core.freeze(vae_mesh_shardings["params"]) + qwen3_shardings = flax.core.freeze(qwen3_mesh_shardings["params"]) + + # 8. Load weights on Host CPU + max_logging.log("Loading parameters on Host CPU...") + t_load_start = time.time() + cpu_device = jax.local_devices(backend="cpu")[0] + with jax.default_device(cpu_device): + with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): + import flax.linen.spmd as flax_spmd + + def unbox_fn(x): + return x.unbox() if isinstance(x, flax_spmd.LogicallyPartitioned) else x + + t_sub0 = time.time() + params = jax.tree_util.tree_map( + unbox_fn, abstract_transformer_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned) + ) + params = flax.core.unfreeze(params) + + vae_params = jax.tree_util.tree_map( + unbox_fn, abstract_vae_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned) + ) + vae_params = flax.core.unfreeze(vae_params) + + qwen3_params = jax.tree_util.tree_map( + unbox_fn, abstract_qwen3_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned) + ) + qwen3_params = flax.core.unfreeze(qwen3_params) + + max_logging.log(f" -> [SUB-TIMING 1/3] PyTree unboxing template setup: {time.time() - t_sub0:.2f}s") + t_sub1 = time.time() + + weight_dtype = jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32 + + params = load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, depth, dtype=weight_dtype) + vae_params, vae_bn_mean, vae_bn_std = load_and_convert_vae_weights( + vae_safetensors_path, vae_params, dtype=weight_dtype + ) + qwen3_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params, qwen3_config) + max_logging.log( + f" -> [SUB-TIMING 2/3] Safetensors loading & key mapping (in target dtype): {time.time() - t_sub1:.4f}s" + ) + + params = flax.core.freeze(params) + vae_params = flax.core.freeze(vae_params) + qwen3_params = flax.core.freeze(qwen3_params) + + max_logging.log("\n" + "=" * 80) + max_logging.log("šŸš€ Pinning all parameters to TPU HBM permanently...") + max_logging.log("=" * 80 + "\n") + t_sub3 = time.time() + max_logging.log("Putting params on TPU HBM...") + with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): + try: + params = jax.tree_util.tree_map(max_utils.device_put_replicated, params, transformer_shardings) + except Exception as err: + max_logging.log("\nāŒ jax.device_put(params, transformer_shardings) FAILED!") + flat_p = flax.traverse_util.flatten_dict(params) + flat_s = flax.traverse_util.flatten_dict(transformer_shardings) + k_p = set(flat_p.keys()) + k_s = set(flat_s.keys()) + max_logging.log(f"Keys in sharding spec but missing in params: {k_s - k_p}") + max_logging.log(f"Keys in params but missing in sharding spec: {k_p - k_s}") + sys.stdout.flush() + raise err + max_logging.log("Putting vae_params on TPU HBM...") + vae_params = jax.tree_util.tree_map(max_utils.device_put_replicated, vae_params, vae_shardings) + max_logging.log("Putting qwen3_params on TPU HBM...") + qwen3_params = jax.tree_util.tree_map(max_utils.device_put_replicated, qwen3_params, qwen3_shardings) + max_logging.log(f" -> [SUB-TIMING 3/3] TPU HBM device_put placement: {time.time() - t_sub3:.4f}s") + max_logging.log("All parameters placed on TPU HBM successfully!") + gc.collect() + jax.effects_barrier() + + load_time = time.time() - t_load_start + max_logging.log(f" -> [TIMING] Total Model Loading & Device Placement: {load_time:.4f} seconds ā±ļø\n") + + # 9. Setup FlowMatch Scheduler + scheduler = FlaxFlowMatchScheduler( + num_train_timesteps=1000, + shift=1.0, + sigma_max=1.0, + sigma_min=0.001, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False, + use_dynamic_shifting=True, + time_shift_type="exponential", + ) + + # 10. Instantiate and invoke FlaxFlux2KleinPipeline + max_logging.log("Instantiating JAX FlaxFlux2KleinPipeline...") + pipeline = FlaxFlux2KleinPipeline( + transformer=transformer, + vae=vae, + text_encoder=qwen3_model, + tokenizer=None, + scheduler=scheduler, + config=config, + mesh=mesh, + ) + + prompt_str = getattr(config, "prompt", "") or "A dog running in a field with butterflies and tall grass" + active_prompts = partition_prompts(prompt_str, config.batch_size) + + if getattr(config, "interactive", False): + max_logging.log("\n" + "=" * 80) + max_logging.log(" BATCHED INTERACTIVE GENERATION MODE ENABLED šŸŽ®") + max_logging.log("The model has been fully loaded and compiled on the TPU.") + max_logging.log(f"Batch size: {config.batch_size} parallel images.") + max_logging.log("Enter prompts separated by '||' (e.g. A cute cat || A red car)") + max_logging.log("Type 'exit' to quit.") + max_logging.log("=" * 80) + + image_idx = 1 + while True: + try: + user_input = input("\nEnter prompt(s): ") + except (KeyboardInterrupt, EOFError): + break + if user_input.strip().lower() in ("exit", "quit"): + break + if not user_input.strip(): + continue + + prompts = partition_prompts(user_input, config.batch_size) + output_file = f"generated_{image_idx:03d}.png" + + pipeline( + prompt=prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=False, + output_dir=config.output_dir, + output_name=output_file, + ) + image_idx += 1 + else: + # Run one-shot generation + latents_to_use = None + use_latents_flag = False + if getattr(config, "latents_path", ""): + max_logging.log(f"Loading custom starting noise latents from: {config.latents_path}...") + latents_to_use = np.load(config.latents_path) + use_latents_flag = True + max_logging.log(f" -> Custom latents shape: {latents_to_use.shape} | sum: {latents_to_use.sum():.6f}") + + max_logging.log("\n" + "=" * 80) + max_logging.log("šŸš€ Pre-compiling XLA graphs concurrently (AOT Compilation)...") + max_logging.log("=" * 80) + aot_time = pipeline.compile_aot_async( + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + batch_size=config.batch_size, + height=config.height, + width=config.width, + ) + + max_logging.log("\n" + "=" * 80) + max_logging.log("šŸš€ Running initial dry run (Warmup Pass) to verify compiled graph execution...") + max_logging.log("=" * 80) + _, warmup_trace = pipeline( + prompt=active_prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=use_latents_flag, + latents=latents_to_use, + output_dir=config.output_dir, + output_name="flux2klein_warmup.png", + warmup=True, + ) + warmup_time = ( + warmup_trace.get("prompt_encoding", 0.0) + + warmup_trace.get("denoise_loop", 0.0) + + warmup_trace.get("vae_decode", 0.0) + ) + + num_reps = int(getattr(config, "num_reps", 1)) + max_logging.log("\n" + "=" * 80) + max_logging.log(f"ā±ļø Running timed pass at full TPU speed (num_reps={num_reps})...") + max_logging.log("=" * 80) + + main_traces = [] + main_times = [] + + for rep in range(num_reps): + rep_str = f" [Rep {rep+1}/{num_reps}]" if num_reps > 1 else "" + if rep > 0: + max_logging.log(f"ā±ļø Running timed pass{rep_str}...") + + if max_utils.profiler_enabled(config) and rep == 0: + max_logging.log(f"šŸš€ XProf / JAX Profiler active! Capturing trace into: {config.tensorboard_dir}") + with max_utils.Profiler(config, session_name="flux2klein_inference"): + _, trace_i = pipeline( + prompt=active_prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=use_latents_flag, + latents=latents_to_use, + output_dir=config.output_dir, + output_name=f"rep_{rep+1}_{config.output_name}" if num_reps > 1 else config.output_name, + ) + else: + _, trace_i = pipeline( + prompt=active_prompts, + params=params, + vae_params=vae_params, + qwen3_params=qwen3_params, + vae_bn_mean=vae_bn_mean, + vae_bn_std=vae_bn_std, + transformer_shardings=transformer_shardings, + vae_shardings=vae_shardings, + qwen3_shardings=qwen3_shardings, + height=config.height, + width=config.width, + num_inference_steps=config.num_inference_steps, + batch_size=config.batch_size, + use_latents=use_latents_flag, + latents=latents_to_use, + output_dir=config.output_dir, + output_name=f"rep_{rep+1}_{config.output_name}" if num_reps > 1 else config.output_name, + ) + + tot_time_i = trace_i.get( + "e2e_pipeline_total", + trace_i.get("prompt_encoding", 0.0) + trace_i.get("denoise_loop", 0.0) + trace_i.get("vae_decode", 0.0), + ) + main_traces.append(trace_i) + main_times.append(tot_time_i) + if num_reps > 1: + max_logging.log( + f" -> Rep {rep+1}/{num_reps} Completed: Total={tot_time_i:.4f}s | Qwen3={trace_i.get('qwen3_encoding', 0.0):.4f}s | Denoise={trace_i.get('denoise_loop', 0.0):.4f}s | VAE={trace_i.get('vae_decode', 0.0):.4f}s" + ) + + avg_main_time = sum(main_times) / num_reps + avg_start_to_qwen3 = sum(tr.get("start_to_qwen3", 0.0) for tr in main_traces) / num_reps + avg_prompt_enc = sum(tr.get("qwen3_encoding", tr.get("prompt_encoding", 0.0)) for tr in main_traces) / num_reps + avg_qwen3_to_denoise = sum(tr.get("qwen3_to_denoise", 0.0) for tr in main_traces) / num_reps + avg_denoise = sum(tr.get("denoise_loop", 0.0) for tr in main_traces) / num_reps + avg_denoise_to_vae = sum(tr.get("denoise_to_vae", 0.0) for tr in main_traces) / num_reps + avg_vae_decode = sum(tr.get("vae_decode", 0.0) for tr in main_traces) / num_reps + avg_image_saving = sum(tr.get("image_saving", 0.0) for tr in main_traces) / num_reps + + total_cold_start = load_time + aot_time + warmup_time + + max_logging.log("\n" + "=" * 80) + max_logging.log("šŸ“Š FLUX.2-KLEIN COMPLETE LATENCY & TIMING BREAKDOWN") + max_logging.log("=" * 80) + max_logging.log(f"1) Model Loading & Placement Time: {load_time:.4f} seconds ā±ļø") + max_logging.log(f"2) Concurrent AOT XLA Compilation Time: {aot_time:.4f} seconds āš”") + max_logging.log(f"3) Warmup Pass Execution Time: {warmup_time:.4f} seconds ā±ļø") + max_logging.log(f" - Qwen3 Encoding: {warmup_trace.get('prompt_encoding', 0.0):.4f}s") + max_logging.log(f" - Flux Denoising: {warmup_trace.get('denoise_loop', 0.0):.4f}s") + max_logging.log(f" - VAE Decoding: {warmup_trace.get('vae_decode', 0.0):.4f}s") + max_logging.log(f"šŸ‘‰ TOTAL COLD-START TIME (Loading + AOT + Warmup): {total_cold_start:.4f} seconds šŸŽÆ") + rep_label = f" (Average across {num_reps} reps)" if num_reps > 1 else "" + max_logging.log(f"4) Main Warmed-Up Pass (Pure Inference Latency){rep_label}: {avg_main_time:.4f} seconds ā±ļø") + max_logging.log(f" - 1. Start -> Qwen3: {avg_start_to_qwen3*1000:.2f} ms ({avg_start_to_qwen3:.4f}s)") + max_logging.log(f" - 2. Qwen3 Encoding: {avg_prompt_enc*1000:.2f} ms ({avg_prompt_enc:.4f}s)") + max_logging.log(f" - 3. Qwen3 -> Denoising: {avg_qwen3_to_denoise*1000:.2f} ms ({avg_qwen3_to_denoise:.4f}s)") + max_logging.log(f" - 4. Flux Denoising Loop: {avg_denoise*1000:.2f} ms ({avg_denoise:.4f}s)") + max_logging.log(f" - 5. Denoising -> VAE: {avg_denoise_to_vae*1000:.2f} ms ({avg_denoise_to_vae:.4f}s)") + max_logging.log(f" - 6. VAE Decoding: {avg_vae_decode*1000:.2f} ms ({avg_vae_decode:.4f}s)") + max_logging.log(f" - 7. Image Saving: {avg_image_saving*1000:.2f} ms ({avg_image_saving:.4f}s)") + max_logging.log(f" - šŸ‘‰ TOTAL E2E PIPELINE: {avg_main_time*1000:.2f} ms ({avg_main_time:.4f}s)") + max_logging.log("=" * 80) + + max_logging.log("\n=======================================================") + max_logging.log(f"SUCCESS! Batched generation complete for {config.batch_size} images! šŸŽØšŸŽ‰") + max_logging.log("=======================================================\n") + + +if __name__ == "__main__": + with transformer_engine_context(): + app.run(main) diff --git a/src/maxdiffusion/models/resnet_flax.py b/src/maxdiffusion/models/resnet_flax.py index 79ddcb30e..8371a4432 100644 --- a/src/maxdiffusion/models/resnet_flax.py +++ b/src/maxdiffusion/models/resnet_flax.py @@ -57,9 +57,8 @@ def setup(self): @nn.compact def __call__(self, hidden_states): batch, height, width, channels = hidden_states.shape - hidden_states = jax.image.resize( - hidden_states, shape=(batch, height * 2, width * 2, channels), method="nearest", precision=self.precision - ) + hidden_states = jnp.broadcast_to(hidden_states[:, :, None, :, None, :], (batch, height, 2, width, 2, channels)) + hidden_states = jnp.reshape(hidden_states, (batch, height * 2, width * 2, channels)) hidden_states = nn.with_logical_constraint(hidden_states, ("conv_batch", "height", "keep_2", "out_channels")) diff --git a/src/maxdiffusion/models/vae_flax.py b/src/maxdiffusion/models/vae_flax.py index 72adcbe79..af13327bf 100644 --- a/src/maxdiffusion/models/vae_flax.py +++ b/src/maxdiffusion/models/vae_flax.py @@ -87,11 +87,8 @@ def setup(self): def __call__(self, hidden_states): batch, height, width, channels = hidden_states.shape - hidden_states = jax.image.resize( - hidden_states, - shape=(batch, height * 2, width * 2, channels), - method="nearest", - ) + hidden_states = jnp.broadcast_to(hidden_states[:, :, None, :, None, :], (batch, height, 2, width, 2, channels)) + hidden_states = jnp.reshape(hidden_states, (batch, height * 2, width * 2, channels)) hidden_states = self.conv(hidden_states) return hidden_states diff --git a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py index 634ec8d9e..df2a8d172 100644 --- a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py +++ b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py @@ -31,13 +31,12 @@ from maxdiffusion.max_utils import device_put_replicated from ..pipeline_flax_utils import FlaxDiffusionPipeline from ...models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel -from ...models.vae_flax import FlaxAutoencoderKL +from ...models.vae_flax import FlaxAutoencoderKL, FlaxDecoderOutput from ...models.qwen3_flax import FlaxQwen3Model from ...schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler, compute_empirical_mu from ...models.flux.util import ( pack_latents, - unpack_latents, prepare_latent_image_ids, prepare_text_ids, ) @@ -70,6 +69,7 @@ def __init__( ) self._config = config self.mesh = mesh + self.tokenizer = tokenizer # JIT compilation cache self._jitted_qwen3_forward = None @@ -97,14 +97,111 @@ def transformer_step(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, ti guidance=guidance, ) - @jax.jit - def vae_decode(v_params, latents_unpatched): - return self.vae.apply({"params": v_params}, latents=latents_unpatched, method=self.vae.decode) + @jax.jit(static_argnums=(4, 5), donate_argnums=(1,)) + def vae_decode(v_params, latents_packed, vae_bn_mean, vae_bn_std, height, width): + def decode_single(single_latent): + vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128) + vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128) + latents_bn = single_latent.reshape(1, -1, 128) * vae_bn_std_seq + vae_bn_mean_seq + + h_latent = height // 8 + w_latent = width // 8 + latents_unpacked = jnp.reshape(latents_bn, (1, h_latent // 2, w_latent // 2, 32, 2, 2)) + latents_unpacked = jnp.transpose(latents_unpacked, (0, 3, 1, 4, 2, 5)) + latents_unpacked = jnp.reshape(latents_unpacked, (1, 32, h_latent, w_latent)) + + res = self.vae.apply({"params": v_params}, latents=latents_unpacked, method=self.vae.decode) + return res.sample[0] + + images = jax.vmap(decode_single)(latents_packed) + return FlaxDecoderOutput(sample=images) self._jitted_qwen3_forward = qwen3_forward self._jitted_transformer_step = transformer_step self._jitted_vae_decode = vae_decode + def _get_dynamic_batch_sharding(self): + """Dynamically infers the batch dimension sharding specification from self.mesh.""" + batch_axes = [axis for axis in ("data", "fsdp") if axis in self.mesh.axis_names and self.mesh.shape[axis] > 1] + spec = P(tuple(batch_axes)) if batch_axes else P(None) + return jax.sharding.NamedSharding(self.mesh, spec) + + def compile_aot_async( + self, params, vae_params, qwen3_params, vae_bn_mean, vae_bn_std, batch_size=1, height=1024, width=1024 + ): + """Triggers AOT compilation for Qwen3, Flux Transformer, and VAE concurrently using ThreadPoolExecutor.""" + self._setup_jit_functions() + max_logging.log("šŸš€ Pre-compiling XLA graphs for Qwen3, Flux Transformer, and VAE concurrently...") + from concurrent.futures import ThreadPoolExecutor + + seq_len_img = (height // 16) * (width // 16) + seq_len_txt = self._config.max_sequence_length + + dummy_ids = jnp.zeros((batch_size, seq_len_txt), dtype=jnp.int32) + dummy_mask = jnp.ones((batch_size, seq_len_txt), dtype=jnp.int32) + + dummy_latents = jnp.zeros((batch_size, seq_len_img, 128), dtype=jnp.float32) + dummy_img_ids = jnp.zeros((batch_size, seq_len_img, 4), dtype=jnp.int32) + dummy_prompt_embeds = jnp.zeros((batch_size, seq_len_txt, self.transformer.joint_attention_dim), dtype=jnp.bfloat16) + dummy_txt_ids = jnp.zeros((batch_size, seq_len_txt, 4), dtype=jnp.float32) + dummy_t_vec = jnp.zeros((batch_size,), dtype=jnp.float32) + + dummy_bn_mean = jnp.array(vae_bn_mean, dtype=jnp.float32) + dummy_bn_std = jnp.array(vae_bn_std, dtype=jnp.float32) + + data_sharding = self._get_dynamic_batch_sharding() + replicated_sharding = jax.sharding.NamedSharding(self.mesh, P()) + + def put_data_on_devices(x, sharding): + if isinstance(x, jax.Array) and hasattr(x, "sharding") and not x.sharding.is_fully_addressable: + return x + if hasattr(sharding, "is_fully_addressable") and sharding.is_fully_addressable: + return jax.device_put(x, sharding) + return device_put_replicated(x, sharding) + + dummy_ids = put_data_on_devices(dummy_ids, data_sharding) + dummy_mask = put_data_on_devices(dummy_mask, data_sharding) + dummy_latents = put_data_on_devices(dummy_latents, data_sharding) + dummy_img_ids = put_data_on_devices(dummy_img_ids, data_sharding) + dummy_prompt_embeds = put_data_on_devices(dummy_prompt_embeds, data_sharding) + dummy_txt_ids = put_data_on_devices(dummy_txt_ids, data_sharding) + dummy_t_vec = put_data_on_devices(dummy_t_vec, data_sharding) + dummy_bn_mean = put_data_on_devices(dummy_bn_mean, replicated_sharding) + dummy_bn_std = put_data_on_devices(dummy_bn_std, replicated_sharding) + + def compile_qwen3(): + t0 = time.perf_counter() + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + self._jitted_qwen3_forward.lower(qwen3_params, dummy_ids, dummy_mask).compile() + max_logging.log(f" -> [AOT COMPILED] Qwen3 Text Encoder in {time.perf_counter() - t0:.2f}s") + + def compile_transformer(): + t0 = time.perf_counter() + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + self._jitted_transformer_step.lower( + params, dummy_latents, dummy_img_ids, dummy_prompt_embeds, dummy_txt_ids, None, dummy_t_vec, None + ).compile() + max_logging.log(f" -> [AOT COMPILED] Flux Transformer Step in {time.perf_counter() - t0:.2f}s") + + def compile_vae(): + t0 = time.perf_counter() + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + self._jitted_vae_decode.lower(vae_params, dummy_latents, dummy_bn_mean, dummy_bn_std, height, width).compile() + max_logging.log(f" -> [AOT COMPILED] VAE Decoder in {time.perf_counter() - t0:.2f}s") + + t_start = time.perf_counter() + with ThreadPoolExecutor(max_workers=3) as executor: + futures = [ + executor.submit(compile_qwen3), + executor.submit(compile_transformer), + executor.submit(compile_vae), + ] + for future in futures: + future.result() + aot_duration = time.perf_counter() - t_start + max_logging.log(f"āš” [AOT CONCURRENT COMPILATION COMPLETE] Total AOT compile time: {aot_duration:.2f}s") + return aot_duration + def _prepare_latents(self, config, batch_size, height, width): num_channels_latents = 32 latent_height = height // 8 @@ -147,6 +244,7 @@ def __call__( use_latents: bool = False, latents: Optional[Any] = None, measure_time: bool = False, + warmup: bool = False, output_dir: str = "output/", output_name: str = "flux2klein_generated_image.png", ): @@ -199,18 +297,37 @@ def __call__( proc_cnt = jax.process_count() host_prefix = f"[HOST {proc_id}/{proc_cnt}] " + # Shard pipeline batch inputs across data axis ("data") for SPMD multi-host execution + data_sharding = jax.sharding.NamedSharding(self.mesh, P("data")) + + def put_data_on_devices(x, sharding): + if isinstance(x, jax.Array) and hasattr(x, "sharding") and not x.sharding.is_fully_addressable: + return x + if hasattr(sharding, "is_fully_addressable") and sharding.is_fully_addressable: + return jax.device_put(x, sharding) + return device_put_replicated(x, sharding) + # --------------------------------------------------------------------- # PHASE A: Encode Prompt (Qwen3) # --------------------------------------------------------------------- - print(f"{host_prefix} [PHASE A] Encoding {len(prompts)} prompt(s) using JAX Qwen3 on TPU...", flush=True) + if not prompts: + raise ValueError("Prompt must be provided to FlaxFlux2KleinPipeline") + if isinstance(prompts, str): + prompts = [prompts] + + max_logging.log(f"{host_prefix} [PHASE A] Encoding {len(prompts)} prompt(s) using JAX Qwen3 on TPU...") t0 = time.perf_counter() try: - # Resolve tokenizer path from config - tokenizer_path = self._config.tokenizer_model_name_or_path + tokenizer_path = getattr(self._config, "tokenizer_model_name_or_path", None) or getattr( + self._config, "pretrained_model_name_or_path", "" + ) hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")) repo_cache = os.path.join( - hf_home, "hub", f"models--{self._config.pretrained_model_name_or_path.replace('/', '--')}", "snapshots" + hf_home, + "hub", + f"models--{getattr(self._config, 'pretrained_model_name_or_path', '').replace('/', '--')}", + "snapshots", ) if os.path.exists(repo_cache) and os.listdir(repo_cache): tokenizer_path = os.path.join(repo_cache, os.listdir(repo_cache)[0]) @@ -232,8 +349,11 @@ def __call__( prompt_ids = jnp.array(inputs["input_ids"]) prompt_mask = jnp.array(inputs["attention_mask"]) - # Run Text Encoding - hidden_states, all_hidden_states = self._jitted_qwen3_forward(qwen3_params, prompt_ids, prompt_mask) + # Run Text Encoding with sharded input arrays matching compile_aot_async + prompt_ids = put_data_on_devices(prompt_ids, data_sharding) + prompt_mask = put_data_on_devices(prompt_mask, data_sharding) + with jax.named_scope("qwen3_text_encoder"): + hidden_states, all_hidden_states = self._jitted_qwen3_forward(qwen3_params, prompt_ids, prompt_mask) # Stack layers 9, 18, 27 to form prompt embeddings h_9 = all_hidden_states[9] @@ -244,7 +364,7 @@ def __call__( prompt_embeds_jax = jnp.transpose(out, (0, 2, 1, 3)).reshape((batch_size, seq_len_txt, -1)) prompt_embeds_jax.block_until_ready() except Exception as e: - print(f"āŒ {host_prefix} EXCEPTION IN PHASE A (QWEN3 ENCODING): {e}", flush=True) + max_logging.log(f"āŒ {host_prefix} EXCEPTION IN PHASE A (QWEN3 ENCODING): {e}") import traceback traceback.print_exc() @@ -260,56 +380,49 @@ def __call__( # Stage Sync 1: Phase A Complete multihost_utils.sync_global_devices("phase_a_complete") - print(f"{host_prefix} Passed Phase A Sync Barrier (phase_a_complete) successfully! āœ…", flush=True) - - # Shard pipeline batch inputs across data axis ("data") for SPMD multi-host execution - data_sharding = jax.sharding.NamedSharding(self.mesh, P("data")) - - def put_data_on_devices(x, sharding): - if isinstance(x, jax.Array) and hasattr(x, "sharding") and not x.sharding.is_fully_addressable: - return x - if hasattr(sharding, "is_fully_addressable") and sharding.is_fully_addressable: - return jax.device_put(x, sharding) - return device_put_replicated(x, sharding) + max_logging.log(f"{host_prefix} Passed Phase A Sync Barrier (phase_a_complete) successfully! āœ…") latents_jax = put_data_on_devices(latents_jax, data_sharding) prompt_embeds_jax = put_data_on_devices(prompt_embeds_jax, data_sharding) txt_ids_val = put_data_on_devices(txt_ids_val, data_sharding) img_ids_val = put_data_on_devices(img_ids_val, data_sharding) - print( + max_logging.log( f"{host_prefix} DIAGNOSTIC TENSORS BEFORE PHASE B:\n" f" latents_jax: shape={latents_jax.shape}, dtype={latents_jax.dtype}, sharding={getattr(latents_jax, 'sharding', None)}\n" f" prompt_embeds_jax: shape={prompt_embeds_jax.shape}, dtype={prompt_embeds_jax.dtype}, sharding={getattr(prompt_embeds_jax, 'sharding', None)}\n" f" txt_ids_val: shape={txt_ids_val.shape}, dtype={txt_ids_val.dtype}, sharding={getattr(txt_ids_val, 'sharding', None)}\n" - f" img_ids_val: shape={img_ids_val.shape}, dtype={img_ids_val.dtype}, sharding={getattr(img_ids_val, 'sharding', None)}", - flush=True, + f" img_ids_val: shape={img_ids_val.shape}, dtype={img_ids_val.dtype}, sharding={getattr(img_ids_val, 'sharding', None)}" ) # Stage Sync 2: Pre-Phase B Start multihost_utils.sync_global_devices("pre_phase_b_start") - print(f"{host_prefix} Passed Pre-Phase B Sync Barrier (pre_phase_b_start) successfully! āœ…", flush=True) + max_logging.log(f"{host_prefix} Passed Pre-Phase B Sync Barrier (pre_phase_b_start) successfully! āœ…") # --------------------------------------------------------------------- # PHASE B: Denoising Loop (Flux Transformer - Standalone Step JIT) # --------------------------------------------------------------------- - print( - f"{host_prefix} [PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images...", - flush=True, + steps_to_run = 1 if warmup else num_inference_steps + max_logging.log( + f"{host_prefix} [PHASE B] Running {steps_to_run}-step E2E Denoising Loop on a batch of {batch_size} images (warmup={warmup})..." ) t0 = time.perf_counter() try: guidance_vec_val = None vec_val = None + active_latents_sharding = getattr(latents_jax, "sharding", data_sharding) - for step_idx in range(num_inference_steps): + for step_idx in range(steps_to_run): + t_step_start = time.perf_counter() timestep = scheduler_state.timesteps[step_idx] t_vec = jnp.full((batch_size,), timestep / 1000.0, dtype=latents_jax.dtype) + t_vec = put_data_on_devices(t_vec, data_sharding) - model_output = self._jitted_transformer_step( - params, latents_jax, img_ids_val, prompt_embeds_jax, txt_ids_val, vec_val, t_vec, guidance_vec_val - ) + with jax.named_scope(f"flux_transformer_step_{step_idx+1}"): + model_output = self._jitted_transformer_step( + params, latents_jax, img_ids_val, prompt_embeds_jax, txt_ids_val, vec_val, t_vec, guidance_vec_val + ) prev_sample, _ = self.scheduler.step( state=scheduler_state, @@ -318,11 +431,15 @@ def put_data_on_devices(x, sharding): sample=latents_jax, return_dict=False, ) - latents_jax = prev_sample + latents_jax = put_data_on_devices(prev_sample, active_latents_sharding) + latents_jax.block_until_ready() + t_step_duration = time.perf_counter() - t_step_start + max_logging.log( + f"{host_prefix} -> Step {step_idx+1}/{steps_to_run} complete in {t_step_duration:.4f}s | latents_sharding={getattr(latents_jax, 'sharding', None)}" + ) - latents_jax.block_until_ready() except Exception as e: - print(f"āŒ {host_prefix} EXCEPTION IN DENOISE LOOP: {e}", flush=True) + max_logging.log(f"āŒ {host_prefix} EXCEPTION IN DENOISE LOOP: {e}") import traceback traceback.print_exc() @@ -331,7 +448,7 @@ def put_data_on_devices(x, sharding): # Stage Sync 3: Phase B Complete multihost_utils.sync_global_devices("phase_b_complete") - print(f"{host_prefix} Passed Phase B Sync Barrier (phase_b_complete) successfully! āœ…", flush=True) + max_logging.log(f"{host_prefix} Passed Phase B Sync Barrier (phase_b_complete) successfully! āœ…") trace["denoise_loop"] = time.perf_counter() - t0 max_logging.log(f" -> [TIMING] Denoising Loop (Flux): {trace['denoise_loop']:.4f} seconds ā±ļø") @@ -342,17 +459,14 @@ def put_data_on_devices(x, sharding): max_logging.log("[PHASE C] Decoding final latents to RGB image using JAX VAE decoder on TPU...") t0 = time.perf_counter() - # Apply Channel-wise Batch Normalization Scaling in packed sequence format (denormalize) - vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128) - vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128) - latents_bn = latents_jax * vae_bn_std_seq + vae_bn_mean_seq - - # Unpack packed latents back to spatial grid - latents_unpacked = unpack_latents(latents_bn, batch_size, 32, height, width) - - # Decode VAE latents to RGB pixels - decoded_out = self._jitted_vae_decode(vae_params, latents_unpacked) - # VAE output is in decoded_out.sample + # Decode VAE latents to RGB pixels using fused JIT vae_decode + data_sharding = self._get_dynamic_batch_sharding() + replicated_sharding = jax.sharding.NamedSharding(self.mesh, P()) + latents_jax = put_data_on_devices(latents_jax, data_sharding) + vae_bn_mean_jax = put_data_on_devices(jnp.array(vae_bn_mean, dtype=jnp.float32), replicated_sharding) + vae_bn_std_jax = put_data_on_devices(jnp.array(vae_bn_std, dtype=jnp.float32), replicated_sharding) + with jax.named_scope("vae_decoder"): + decoded_out = self._jitted_vae_decode(vae_params, latents_jax, vae_bn_mean_jax, vae_bn_std_jax, height, width) images_rgb = decoded_out.sample images_rgb.block_until_ready() diff --git a/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py b/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py index 24362d35d..3f1ae6186 100644 --- a/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py +++ b/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py @@ -17,6 +17,7 @@ import os import unittest import pytest +import jax import numpy as np from PIL import Image @@ -58,6 +59,7 @@ def test_flux2klein_4b_smoke(self): f"prompt={PROMPT}", "height=512", "width=512", + f"per_device_batch_size={1.0 / jax.device_count()}", "batch_size=1", "seed=42", "ici_fsdp_parallelism=-1", @@ -101,6 +103,7 @@ def test_flux2klein_9b_smoke(self): f"prompt={PROMPT}", "height=512", "width=512", + f"per_device_batch_size={1.0 / jax.device_count()}", "batch_size=1", "seed=42", "ici_fsdp_parallelism=-1", @@ -117,7 +120,7 @@ def test_flux2klein_9b_smoke(self): self.assertEqual(base_image.shape, test_image.shape) ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255) print(f"\n[SMOKE TEST 9B] SSIM Score: {ssim_compare:.6f}") - self.assertGreaterEqual(ssim_compare, 0.80) + self.assertGreaterEqual(ssim_compare, 0.8) if __name__ == "__main__": diff --git a/src/maxdiffusion/tests/images/ref_flux2klein_9b.png b/src/maxdiffusion/tests/images/ref_flux2klein_9b.png index 594464a8f..c27f959ac 100644 Binary files a/src/maxdiffusion/tests/images/ref_flux2klein_9b.png and b/src/maxdiffusion/tests/images/ref_flux2klein_9b.png differ