diff --git a/debugging_ulysses_ring_custom.sh b/debugging_ulysses_ring_custom.sh new file mode 100755 index 000000000..4598abfd5 --- /dev/null +++ b/debugging_ulysses_ring_custom.sh @@ -0,0 +1,38 @@ +export TPU_VISIBLE_CHIPS=0 +#!/bin/bash +set -e + +# Activate environment +if [ -f "/mnt/workspace/maxdiffusion_venv/bin/activate" ]; then + source /mnt/workspace/maxdiffusion_venv/bin/activate +fi + +export LIBTPU_INIT_ARGS='--xla_tpu_dvfs_p_state=7 --xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_dot_strength_reduction=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true --xla_tpu_enable_llo_profiling=true --xla_enable_mxu_trace=true --xla_enable_transpose_trace=true --xla_enable_local_dma_trace=true --xla_tpu_emit_tracing_vwaits=true --xla_jf_debug_level=1' + +if [ -f "/home/ameypasarkar_google_com/cloud-devkit/tpu_python" ]; then + PYTHON_BIN="/home/ameypasarkar_google_com/cloud-devkit/tpu_python" +elif [ -f "/home/ameypasarkar/cloud-devkit/tpu_python" ]; then + PYTHON_BIN="/home/ameypasarkar/cloud-devkit/tpu_python" +else + PYTHON_BIN="python3" +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +TPU_VISIBLE_DEVICES=0 $PYTHON_BIN src/maxdiffusion/generate_flux2klein.py \ + src/maxdiffusion/configs/base_flux2klein.yml \ + output_name=flux2klein_ring_custom.png \ + attention=ulysses_ring_custom \ + ulysses_attention_chunks=1 \ + text_encoder_attention=dot_product \ + ici_context_parallelism=2 \ + ulysses_shards=2 \ + per_device_batch_size=0.5 \ + height=1024 \ + width=1024 \ + num_inference_steps=4 \ + prompt='a dog running in a field' \ + num_reps=5 \ + mask_padding_tokens=False \ + flash_block_sizes='{"block_q": 4608, "block_kv": 1024, "block_kv_compute": 1024}' diff --git a/prxpixel_cfg_1.png b/prxpixel_cfg_1.png new file mode 100644 index 000000000..8edfbd673 Binary files /dev/null and b/prxpixel_cfg_1.png differ diff --git a/prxpixel_cfg_2.png b/prxpixel_cfg_2.png new file mode 100644 index 000000000..110fcdf37 Binary files /dev/null and b/prxpixel_cfg_2.png differ diff --git a/prxpixel_cfg_2_no_np.png b/prxpixel_cfg_2_no_np.png new file mode 100644 index 000000000..d4905a3c1 Binary files /dev/null and b/prxpixel_cfg_2_no_np.png differ diff --git a/prxpixel_cfg_4.5.png b/prxpixel_cfg_4.5.png new file mode 100644 index 000000000..7aa75dfa8 Binary files /dev/null and b/prxpixel_cfg_4.5.png differ diff --git a/prxpixel_sample_512.png b/prxpixel_sample_512.png new file mode 100644 index 000000000..fcba8b759 Binary files /dev/null and b/prxpixel_sample_512.png differ diff --git a/run_component_profiles.sh b/run_component_profiles.sh new file mode 100755 index 000000000..70f3f1842 --- /dev/null +++ b/run_component_profiles.sh @@ -0,0 +1,98 @@ +#!/bin/bash +set -e + +BLOCK_Q=${1:-4608} +BLOCK_KV=${2:-1024} +PROFILE_TARGET=${3:-"all"} +LOG_FILE=${4:-"output/component_profiles.log"} +ATTENTION_KERNEL=${5:-"ulysses_ring_custom_fixed_m"} +TENSORBOARD_DIR=${6:-"/mnt/workspace/maxdiffusion/tensorboard_traces"} +NUM_STEPS=${7:-1} + +mkdir -p output +mkdir -p "${TENSORBOARD_DIR}" + +export LIBTPU_INIT_ARGS='--xla_tpu_dvfs_p_state=7 \ +--xla_tpu_spmd_rng_bit_generator_unsafe=true \ +--xla_tpu_enable_dot_strength_reduction=true \ +--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \ +--xla_enable_async_collective_permute=true \ +--xla_tpu_data_parallel_opt_different_sized_ops=true \ +--xla_tpu_enable_async_collective_fusion=true \ +--xla_tpu_enable_async_collective_fusion_multiple_steps=true \ +--xla_tpu_overlap_compute_collective_tc=true \ +--xla_enable_async_all_gather=true \ +--xla_tpu_scoped_vmem_limit_kib=65536 \ +--xla_tpu_enable_async_all_to_all=true \ +--xla_tpu_enable_all_experimental_scheduler_features=true \ +--xla_tpu_enable_scheduler_memory_pressure_tracking=true \ +--xla_tpu_host_transfer_overlap_limit=24 \ +--xla_tpu_aggressive_opt_barrier_removal=ENABLED \ +--xla_lhs_prioritize_async_depth_over_stall=ENABLED \ +--xla_should_allow_loop_variant_parameter_in_chain=ENABLED \ +--xla_should_add_loop_invariant_op_in_chain=ENABLED \ +--xla_tpu_enable_ici_ag_pipelining=true \ +--xla_max_concurrent_host_send_recv=100 \ +--xla_tpu_scheduler_percent_shared_memory_limit=100 \ +--xla_latency_hiding_scheduler_rerun=2 \ +--xla_tpu_use_minor_sharding_for_major_trivial_input=true \ +--xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 \ +--xla_tpu_enable_latency_hiding_scheduler=true \ +--xla_tpu_enable_ag_backward_pipelining=true \ +--xla_tpu_enable_megacore_fusion=true \ +--xla_tpu_megacore_fusion_allow_ags=true \ +--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \ +--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \ +--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \ +--xla_tpu_enable_sparse_core_collective_aggregator=true \ +--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \ +--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \ +--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \ +--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \ +--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \ +--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \ +--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \ +--xla_tpu_enable_concurrent_sparse_core_offloading=true \ +--xla_tpu_assign_all_reduce_scatter_layout=true \ +--xla_tpu_enable_llo_profiling=true \ +--xla_enable_mxu_trace=true \ +--xla_enable_transpose_trace=true \ +--xla_enable_local_dma_trace=true \ +--xla_tpu_emit_tracing_vwaits=true \ +--xla_jf_debug_level=1' + +echo "================================================================================" +echo "🚀 Running FLUX.2-Klein Component Profiling" +echo " BLOCK_Q: ${BLOCK_Q}" +echo " BLOCK_KV: ${BLOCK_KV}" +echo " PROFILE_TARGET: ${PROFILE_TARGET}" +echo " ATTENTION_KERNEL: ${ATTENTION_KERNEL}" +echo " TENSORBOARD_DIR: ${TENSORBOARD_DIR}" +echo "================================================================================" + +TPU_PYTHON="/home/ameypasarkar_google_com/cloud-devkit/tpu_python" +if [ ! -f "${TPU_PYTHON}" ]; then + TPU_PYTHON="python3" +fi + +TPU_VISIBLE_DEVICES=0 ${TPU_PYTHON} src/maxdiffusion/generate_flux2klein_component_profiles.py \ + src/maxdiffusion/configs/base_flux2klein_9B.yml \ + pretrained_model_name_or_path=/mnt/hyperdisk_weights/hub/flux2klein-9b \ + attention="${ATTENTION_KERNEL}" \ + ulysses_attention_chunks=1 \ + text_encoder_attention=flash \ + text_encoder_flash_block_sizes='{"block_q": 512, "block_kv": 512, "block_kv_compute": 512}' \ + text_encoder_max_layer=27 \ + ici_context_parallelism=2 \ + ulysses_shards=2 \ + per_device_batch_size=0.5 \ + height=1024 \ + width=1024 \ + num_inference_steps="${NUM_STEPS}" \ + prompt='a dog running in a field' \ + num_reps=1 \ + mask_padding_tokens=False \ + profile_target="${PROFILE_TARGET}" \ + tensorboard_dir="${TENSORBOARD_DIR}" \ + flash_block_sizes="{\"block_q\": ${BLOCK_Q}, \"block_kv\": ${BLOCK_KV}, \"block_kv_compute\": ${BLOCK_KV}}" \ + 2>&1 | tee "${LOG_FILE}" diff --git a/run_e2e_image_edit_parity.sh b/run_e2e_image_edit_parity.sh new file mode 100755 index 000000000..aa31bc8ec --- /dev/null +++ b/run_e2e_image_edit_parity.sh @@ -0,0 +1,32 @@ +#!/bin/bash +set -e + +# Activate virtual environment +if [ -f "/mnt/workspace/maxdiffusion_venv/bin/activate" ]; then + source /mnt/workspace/maxdiffusion_venv/bin/activate +elif [ -f "/mnt/data/mdiff_venv/bin/activate" ]; then + source /mnt/data/mdiff_venv/bin/activate +fi + +# Detect CDK tpu_python wrapper +if [ -f "/home/ameypasarkar_google_com/cloud-devkit/tpu_python" ]; then + PYTHON_BIN="/home/ameypasarkar_google_com/cloud-devkit/tpu_python" +elif [ -f "/home/ameypasarkar/cloud-devkit/tpu_python" ]; then + PYTHON_BIN="/home/ameypasarkar/cloud-devkit/tpu_python" +else + PYTHON_BIN="python3" +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +# Point to workspace HF cache containing FLUX.2-Klein 4B weights +export HF_HOME="${HF_HOME:-/mnt/workspace/hf_cache}" + +echo "================================================================================" +echo "🚀 Launching FLUX.2-Klein Multi-Image Editing E2E Parity Test on TPU v7 via CDK" +echo " Python Runner: $PYTHON_BIN" +echo " HF_HOME: $HF_HOME" +echo "================================================================================" + +$PYTHON_BIN -m pytest src/maxdiffusion/tests/edit_flux2klein_e2e_test.py -k test_e2e_image_edit_parity_vs_diffusers -s -v "$@" diff --git a/run_image_edit.sh b/run_image_edit.sh new file mode 100755 index 000000000..0a0ddbaa7 --- /dev/null +++ b/run_image_edit.sh @@ -0,0 +1,46 @@ +#!/bin/bash +set -e + +# Activate environment +if [ -f "/mnt/workspace/maxdiffusion_venv/bin/activate" ]; then + source /mnt/workspace/maxdiffusion_venv/bin/activate +fi + +# ============================================================================== +# Advanced XLA / TPU Runtime Optimization Flags +# ============================================================================== +export LIBTPU_INIT_ARGS='--xla_tpu_dvfs_p_state=7 --xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_dot_strength_reduction=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true --xla_tpu_enable_llo_profiling=true --xla_enable_mxu_trace=true --xla_enable_transpose_trace=true --xla_enable_local_dma_trace=true --xla_tpu_emit_tracing_vwaits=true --xla_jf_debug_level=1' + +# Detect python binary (cloud-devkit tpu_python or system python3) +if [ -f "/home/ameypasarkar/cloud-devkit/tpu_python" ]; then + PYTHON_BIN="/home/ameypasarkar/cloud-devkit/tpu_python" +elif [ -f "/home/ameypasarkar_google_com/cloud-devkit/tpu_python" ]; then + PYTHON_BIN="/home/ameypasarkar_google_com/cloud-devkit/tpu_python" +else + PYTHON_BIN="python3" +fi + +# Detect workspace root +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +TPU_VISIBLE_DEVICES=0 $PYTHON_BIN src/maxdiffusion/generate_flux2klein.py \ + src/maxdiffusion/configs/base_flux2klein_9B.yml \ + pretrained_model_name_or_path=/mnt/hyperdisk_weights/hub/flux2klein-9b \ + prompt="zoom in on this picture and show a grand entrance to a big ballroom from the front door" \ + image_paths="['${SCRIPT_DIR}/src/maxdiffusion/test_output/test_sdxl.png']" \ + output_dir="${SCRIPT_DIR}/src/maxdiffusion/test_output/" \ + output_name=edited_ballroom_9b.png \ + height=1024 \ + width=1024 \ + num_inference_steps=4 \ + num_reps=5 \ + per_device_batch_size=0.5 \ + ici_context_parallelism=2 \ + ulysses_shards=2 \ + ulysses_attention_chunks=1 \ + attention=ulysses_ring_custom_fixed_m \ + text_encoder_attention=dot_product \ + mask_padding_tokens=False \ + flash_block_sizes='{"block_q": 1024, "block_kv": 1024, "block_kv_compute": 1024}' \ + seed=42 diff --git a/run_prxpixel_cfg_sweep.py b/run_prxpixel_cfg_sweep.py new file mode 100644 index 000000000..0ad4dc017 --- /dev/null +++ b/run_prxpixel_cfg_sweep.py @@ -0,0 +1,139 @@ +""" +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. +""" + +"""Runs PRXPixel CFG Guidance Scale Sweep (3 Prompts x 4 CFG Levels) on TPU.""" + +import os +import sys +import time +import jax +import jax.numpy as jnp +from PIL import Image, ImageDraw + +from maxdiffusion.checkpointing.prx_pixel_checkpointer import PRXPixelCheckpointer + +SNAPSHOT_DIR = os.path.expanduser("~/.cache/huggingface/hub/models--Photoroom--prxpixel-t2i/snapshots/bcd5e63f072257a220c5d0ba039c97657398b1c2") +OUTPUT_DIR = "/mnt/data/maxdiffusion/prxpixel_testoutputs" + +PROMPTS = [ + ("sports_car", "A sleek modern sports car on a winding mountain road at dusk, cinematic lighting"), + ("punk_musician", "a portrait of a punk musician with green mohawk and nose ring, leather jacket, harsh flash photography"), + ("male_model", "an adult male model wearing a short-sleeve t-shirt and jeans under studio lighting with white background"), +] + +CFG_SCALES = [1.0, 2.0, 3.0, 4.0] +STEPS = 28 +HEIGHT = 512 +WIDTH = 512 +SEED = 42 + + +def create_comparison_strip(images, cfg_scales, title, out_path): + """Creates a 4-column side-by-side comparison image with header banners.""" + w, h = images[0].size + header_h = 50 + padding = 10 + total_w = w * len(images) + padding * (len(images) - 1) + total_h = h + header_h + + combined = Image.new("RGB", (total_w, total_h), (240, 240, 240)) + draw = ImageDraw.Draw(combined) + + for i, (img, cfg) in enumerate(zip(images, cfg_scales)): + x_offset = i * (w + padding) + combined.paste(img, (x_offset, header_h)) + text = f"CFG = {cfg:.1f}" + draw.text((x_offset + w // 2 - 35, 18), text, fill=(0, 0, 0)) + + combined.save(out_path) + print(f" đŸ–ŧī¸ Saved comparison strip: {out_path}", flush=True) + + +def main(): + os.makedirs(OUTPUT_DIR, exist_ok=True) + jax.config.update("jax_default_matmul_precision", "highest") + + print("=" * 90, flush=True) + print("🚀 Starting PRXPixel CFG Guidance Scale Sweep on TPU", flush=True) + print(f" Prompts: {len(PROMPTS)} | CFG Scales: {CFG_SCALES} | Steps: {STEPS} | Resolution: {HEIGHT}x{WIDTH}", flush=True) + print(f" Output Directory: {OUTPUT_DIR}", flush=True) + print(f" Devices: {jax.devices()}", flush=True) + print("=" * 90, flush=True) + + t0 = time.perf_counter() + pipeline = PRXPixelCheckpointer.load_pipeline(SNAPSHOT_DIR, dtype=jnp.bfloat16) + load_time = time.perf_counter() - t0 + print(f"⏱ī¸ Model Load Time: {load_time:.2f}s\n", flush=True) + + # Warmup run to trigger JIT compilation + print("đŸ”Ĩ JIT Warming up pipeline...", flush=True) + t_w0 = time.perf_counter() + _ = pipeline( + prompt="warmup test prompt", + height=HEIGHT, + width=WIDTH, + num_inference_steps=STEPS, + guidance_scale=4.0, + generator=jax.random.PRNGKey(0), + ) + print(f"⏱ī¸ Warmup JIT Time: {time.perf_counter() - t_w0:.2f}s\n", flush=True) + + for prompt_idx, (tag, prompt_text) in enumerate(PROMPTS, 1): + print("=" * 90, flush=True) + print(f"📸 [{prompt_idx}/{len(PROMPTS)}] Prompt: \"{prompt_text}\"", flush=True) + print("=" * 90, flush=True) + + prompt_images = [] + # Deterministic initial noise per prompt + initial_noise = 2.0 * jax.random.normal( + jax.random.PRNGKey(SEED), + shape=(1, 3, HEIGHT, WIDTH), + dtype=jnp.bfloat16, + ) + + for cfg in CFG_SCALES: + t_start = time.perf_counter() + img = pipeline( + prompt=prompt_text, + negative_prompt="", + height=HEIGHT, + width=WIDTH, + num_inference_steps=STEPS, + guidance_scale=cfg, + latents=initial_noise, + output_type="pil", + )[0] + gen_time = time.perf_counter() - t_start + + img_filename = f"{tag}_cfg{int(cfg)}.png" + img_path = os.path.join(OUTPUT_DIR, img_filename) + img.save(img_path) + prompt_images.append(img) + print(f" ✅ CFG {cfg:.1f} generated in {gen_time:.2f}s -> {img_filename}", flush=True) + + # Generate 4-column comparison strip + strip_filename = f"{tag}_cfg_comparison_1_to_4.png" + strip_path = os.path.join(OUTPUT_DIR, strip_filename) + create_comparison_strip(prompt_images, CFG_SCALES, prompt_text, strip_path) + + print("\n" + "=" * 90, flush=True) + print("🎉 All 12 CFG sweep images & comparison strips successfully generated!", flush=True) + print(f"📁 Files saved in: {OUTPUT_DIR}", flush=True) + print("=" * 90, flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/maxdiffusion/checkpointing/base_stable_diffusion_checkpointer.py b/src/maxdiffusion/checkpointing/base_stable_diffusion_checkpointer.py index 26851b1ff..e37d4f60d 100644 --- a/src/maxdiffusion/checkpointing/base_stable_diffusion_checkpointer.py +++ b/src/maxdiffusion/checkpointing/base_stable_diffusion_checkpointer.py @@ -32,7 +32,13 @@ max_logging, ) -from transformers import (CLIPTokenizer, FlaxCLIPTextModel, CLIPTextConfig, FlaxCLIPTextModelWithProjection) +from transformers import CLIPTokenizer +try: + from transformers import FlaxCLIPTextModel, CLIPTextConfig, FlaxCLIPTextModelWithProjection +except ImportError: + FlaxCLIPTextModel = None + CLIPTextConfig = None + FlaxCLIPTextModelWithProjection = None from maxdiffusion.checkpointing.checkpointing_utils import ( create_orbax_checkpoint_manager, diff --git a/src/maxdiffusion/checkpointing/flux_checkpointer.py b/src/maxdiffusion/checkpointing/flux_checkpointer.py index 70b54d08d..1a69be675 100644 --- a/src/maxdiffusion/checkpointing/flux_checkpointer.py +++ b/src/maxdiffusion/checkpointing/flux_checkpointer.py @@ -30,7 +30,12 @@ from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel from ..pipelines.flux.flux_pipeline import FluxPipeline -from transformers import (CLIPTokenizer, FlaxCLIPTextModel, FlaxT5EncoderModel, AutoTokenizer) +from transformers import CLIPTokenizer, AutoTokenizer +try: + from transformers import FlaxCLIPTextModel, FlaxT5EncoderModel +except ImportError: + FlaxCLIPTextModel = None + FlaxT5EncoderModel = None from maxdiffusion.checkpointing.checkpointing_utils import (create_orbax_checkpoint_manager) from maxdiffusion.models.flux.util import load_flow_model diff --git a/src/maxdiffusion/configs/base_flux2klein.yml b/src/maxdiffusion/configs/base_flux2klein.yml index f1a7e0fa4..16a8a5cfa 100644 --- a/src/maxdiffusion/configs/base_flux2klein.yml +++ b/src/maxdiffusion/configs/base_flux2klein.yml @@ -40,6 +40,8 @@ max_sequence_length: 512 time_shift: True base_shift: 0.5 max_shift: 1.15 +image_paths: [] +use_base2_exp: True unet_checkpoint: '' diff --git a/src/maxdiffusion/configs/base_flux2klein_9B.yml b/src/maxdiffusion/configs/base_flux2klein_9B.yml index a3a0afeac..c6dc2689b 100644 --- a/src/maxdiffusion/configs/base_flux2klein_9B.yml +++ b/src/maxdiffusion/configs/base_flux2klein_9B.yml @@ -40,6 +40,8 @@ max_sequence_length: 512 time_shift: True base_shift: 0.5 max_shift: 1.15 +image_paths: [] +use_base2_exp: True unet_checkpoint: '' diff --git a/src/maxdiffusion/generate_flux.py b/src/maxdiffusion/generate_flux.py index 8d2153c65..18f0b9aeb 100644 --- a/src/maxdiffusion/generate_flux.py +++ b/src/maxdiffusion/generate_flux.py @@ -28,8 +28,13 @@ import flax.linen as nn from chex import Array from einops import rearrange -from flax.linen import partitioning as nn_partitioning -from transformers import (CLIPTokenizer, FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel, AutoTokenizer) +from transformers import CLIPTokenizer, AutoTokenizer +try: + from transformers import FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel +except ImportError: + FlaxCLIPTextModel = None + T5EncoderModel = None + FlaxT5EncoderModel = None from maxdiffusion import FlaxAutoencoderKL, pyconfig, max_logging, max_utils from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel diff --git a/src/maxdiffusion/generate_flux2klein.py b/src/maxdiffusion/generate_flux2klein.py index b1427f937..5d148bd64 100644 --- a/src/maxdiffusion/generate_flux2klein.py +++ b/src/maxdiffusion/generate_flux2klein.py @@ -20,6 +20,7 @@ import sys from typing import List +from PIL import Image, UnidentifiedImageError from absl import app import jax import jax.numpy as jnp @@ -267,10 +268,10 @@ def main(argv): 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, + num_key_value_heads=getattr(pt_config, "num_key_value_heads", pt_config.num_attention_heads), + max_position_embeddings=getattr(pt_config, "max_position_embeddings", 32768), + rms_norm_eps=getattr(pt_config, "rms_norm_eps", 1e-6), + rope_theta=getattr(pt_config, "rope_theta", getattr(pt_config, "rope_base", 1000000.0)), dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, attention_kernel=getattr(config, "text_encoder_attention", "flash"), flash_block_sizes=te_bs, @@ -321,6 +322,7 @@ def main(argv): scale_shift_order=getattr(config, "scale_shift_order", "scale_shift"), ulysses_shards=getattr(config, "ulysses_shards", -1), ulysses_attention_chunks=getattr(config, "ulysses_attention_chunks", 1), + use_base2_exp=getattr(config, "use_base2_exp", True), ) # 6. Instantiate JAX VAE @@ -464,6 +466,34 @@ def unbox_fn(x): raise ValueError("Prompt must be specified in the configuration YAML or passed via CLI prompt='...'") active_prompts = partition_prompts(prompt_str, config.batch_size) + # Parse reference image paths for multi-image editing if provided + images = None + image_paths = getattr(config, "image_paths", None) + if image_paths is not None: + if isinstance(image_paths, str) and image_paths.strip(): + import ast + + try: + image_paths = ast.literal_eval(image_paths) + except Exception: + image_paths = [p.strip() for p in image_paths.split(",") if p.strip()] + if isinstance(image_paths, (list, tuple)) and len(image_paths) > 0: + max_logging.log(f" -> Loading {len(image_paths)} reference image(s) for multi-image editing...") + images = [] + for p in image_paths: + try: + if not os.path.exists(p): + raise FileNotFoundError(f"Reference image file not found: {p}") + with Image.open(p) as img_raw: + img = img_raw.convert("RGB").resize((config.width, config.height), Image.Resampling.BICUBIC) + images.append(img) + except (UnidentifiedImageError, OSError, FileNotFoundError) as e: + max_logging.log(f"❌ Error loading reference image '{p}': {e}") + raise ValueError(f"Failed to load reference image '{p}': {e}") from e + except Exception as e: + max_logging.log(f"❌ Unexpected error loading reference image '{p}': {e}") + raise ValueError(f"Failed to load reference image '{p}': {e}") from e + if getattr(config, "interactive", False): max_logging.log("\n" + "=" * 80) max_logging.log(" BATCHED INTERACTIVE GENERATION MODE ENABLED 🎮") @@ -501,6 +531,7 @@ def unbox_fn(x): width=config.width, num_inference_steps=config.num_inference_steps, batch_size=config.batch_size, + images=images, use_latents=False, output_dir=config.output_dir, output_name=output_file, @@ -528,6 +559,7 @@ def unbox_fn(x): batch_size=config.batch_size, height=config.height, width=config.width, + images=images, ) max_logging.log("\n" + "=" * 80) @@ -547,6 +579,7 @@ def unbox_fn(x): width=config.width, num_inference_steps=config.num_inference_steps, batch_size=config.batch_size, + images=images, use_latents=use_latents_flag, latents=latents_to_use, output_dir=config.output_dir, @@ -554,7 +587,8 @@ def unbox_fn(x): warmup=True, ) warmup_time = ( - warmup_trace.get("prompt_encoding", 0.0) + warmup_trace.get("vae_encode", 0.0) + + warmup_trace.get("prompt_encoding", 0.0) + warmup_trace.get("denoise_loop", 0.0) + warmup_trace.get("vae_decode", 0.0) ) @@ -589,6 +623,7 @@ def unbox_fn(x): width=config.width, num_inference_steps=config.num_inference_steps, batch_size=config.batch_size, + images=images, use_latents=use_latents_flag, latents=latents_to_use, output_dir=config.output_dir, @@ -609,6 +644,7 @@ def unbox_fn(x): width=config.width, num_inference_steps=config.num_inference_steps, batch_size=config.batch_size, + images=images, use_latents=use_latents_flag, latents=latents_to_use, output_dir=config.output_dir, @@ -617,16 +653,22 @@ def unbox_fn(x): 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), + trace_i.get("vae_encode", 0.0) + + 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: + vae_enc_str = f" | VAE_Enc={trace_i.get('vae_encode', 0.0):.4f}s" if trace_i.get("vae_encode", 0.0) > 0 else "" 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" + f" -> Rep {rep+1}/{num_reps} Completed: Total={tot_time_i:.4f}s{vae_enc_str} | Qwen3={trace_i.get('qwen3_encoding', 0.0):.4f}s | Denoise={trace_i.get('denoise_loop', 0.0):.4f}s | VAE_Dec={trace_i.get('vae_decode', 0.0):.4f}s" ) avg_main_time = sum(main_times) / num_reps + avg_vae_encode = sum(tr.get("vae_encode", 0.0) for tr in main_traces) / num_reps + avg_vae_to_qwen3 = sum(tr.get("vae_encode_to_qwen3", 0.0) for tr in main_traces) / 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 @@ -643,19 +685,38 @@ def unbox_fn(x): 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 ⏱ī¸") + if warmup_trace.get("vae_encode", 0.0) > 0: + max_logging.log(f" - VAE Encoding: {warmup_trace.get('vae_encode', 0.0):.4f}s") 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)") + step_num = 1 + if avg_vae_encode > 0: + max_logging.log(f" - {step_num}. VAE Image Encoding: {avg_vae_encode*1000:.2f} ms ({avg_vae_encode:.4f}s)") + step_num += 1 + max_logging.log(f" - {step_num}. VAE -> Qwen3: {avg_vae_to_qwen3*1000:.2f} ms ({avg_vae_to_qwen3:.4f}s)") + step_num += 1 + else: + max_logging.log( + f" - {step_num}. Start -> Qwen3: {avg_start_to_qwen3*1000:.2f} ms ({avg_start_to_qwen3:.4f}s)" + ) + step_num += 1 + max_logging.log(f" - {step_num}. Qwen3 Encoding: {avg_prompt_enc*1000:.2f} ms ({avg_prompt_enc:.4f}s)") + step_num += 1 + max_logging.log( + f" - {step_num}. Qwen3 -> Denoising: {avg_qwen3_to_denoise*1000:.2f} ms ({avg_qwen3_to_denoise:.4f}s)" + ) + step_num += 1 + max_logging.log(f" - {step_num}. Flux Denoising Loop: {avg_denoise*1000:.2f} ms ({avg_denoise:.4f}s)") + step_num += 1 + max_logging.log(f" - {step_num}. Denoising -> VAE: {avg_denoise_to_vae*1000:.2f} ms ({avg_denoise_to_vae:.4f}s)") + step_num += 1 + max_logging.log(f" - {step_num}. VAE Decoding: {avg_vae_decode*1000:.2f} ms ({avg_vae_decode:.4f}s)") + step_num += 1 + max_logging.log(f" - {step_num}. 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) diff --git a/src/maxdiffusion/generate_flux2klein_component_profiles.py b/src/maxdiffusion/generate_flux2klein_component_profiles.py new file mode 100644 index 000000000..286bb1e51 --- /dev/null +++ b/src/maxdiffusion/generate_flux2klein_component_profiles.py @@ -0,0 +1,622 @@ +""" +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 PIL import Image, UnidentifiedImageError +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.vae.autoencoder_kl_flux2_nnx import ( + NNXAutoencoderKLFlux2, + load_and_convert_flux2klein_nnx_vae_weights, +) +from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model +from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights +from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler +from maxdiffusion.pipelines.flux.flux2klein_profile_pipeline import FlaxFlux2KleinProfilePipeline + + +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 main(argv): + # Enable shardy partitioner for TPU execution + jax.config.update("jax_use_shardy_partitioner", True) + + # Extract profiling-specific CLI overrides before pyconfig validation + profile_target = os.environ.get("PROFILE_TARGET", "all") + tensorboard_dir = os.environ.get("TENSORBOARD_DIR", "/mnt/workspace/maxdiffusion/tensorboard_traces") + + # 1. Load configurations + config_path = "src/maxdiffusion/configs/base_flux2klein.yml" + custom_overrides = [] + if len(argv) > 1: + raw_args = argv[1:] + if raw_args[0].endswith(".yml") or raw_args[0].endswith(".yaml"): + config_path = raw_args[0] + raw_args = raw_args[1:] + + for arg in raw_args: + if arg and arg.startswith("profile_target="): + profile_target = arg.split("=", 1)[1] + elif arg and arg.startswith("tensorboard_dir="): + tensorboard_dir = arg.split("=", 1)[1] + else: + custom_overrides.append(arg) + + max_logging.log(f"Initializing pyconfig with config: {config_path}") + default_args = [ + None, + config_path, + "run_name=flux2klein_component_profiling", + "output_dir=output/", + ] + default_args.extend(custom_overrides) + + pyconfig.initialize(default_args) + + config = pyconfig.config + os.makedirs(config.output_dir, exist_ok=True) + + num_devices_to_use = getattr(config, "num_devices", None) + if num_devices_to_use is not None and num_devices_to_use > 0: + active_devices = jax.devices()[:num_devices_to_use] + else: + active_devices = jax.devices() + active_device_count = len(active_devices) + + if hasattr(config, "per_device_batch_size") and config.per_device_batch_size > 0: + calculated_batch_size = int(config.per_device_batch_size * active_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 active_device_count={active_device_count} " + f"evaluated to {config.per_device_batch_size * active_device_count}, which truncates to 0." + ) + 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 active_device_count={active_device_count}." + ) + pyconfig._config.keys["batch_size"] = calculated_batch_size + + # 2. Setup device mesh + custom_parallelism_set = any( + any(arg.startswith(f"{k}=") for arg in sys.argv) + for k in [ + "ici_data_parallelism", + "ici_fsdp_parallelism", + "ici_context_parallelism", + "ici_tensor_parallelism", + ] + ) + + if not custom_parallelism_set and active_device_count > 1: + max_logging.log( + f"ℹī¸ Defaulting to Tensor Parallelism: ici_tensor_parallelism={active_device_count} on {active_device_count} TPU devices." + ) + pyconfig._config.keys["ici_tensor_parallelism"] = active_device_count + pyconfig._config.keys["ici_data_parallelism"] = 1 + pyconfig._config.keys["ici_fsdp_parallelism"] = 1 + pyconfig._config.keys["ici_context_parallelism"] = 1 + + max_logging.log("Setting up JAX device mesh...") + devices_array = create_device_mesh(config, devices=active_devices) + 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: + raise ValueError("pretrained_model_name_or_path must be specified in configuration YAML or CLI.") + + use_kv = getattr(config, "use_kv", False) + if use_kv: + if repo_id in ("black-forest-labs/FLUX.2-klein-4B", "black-forest-labs/FLUX.2-klein-4b"): + max_logging.log("⚠ī¸ Warning: KV cache not supported for 4B model, ignoring use_kv=True.") + pyconfig._config.keys["use_kv"] = False + elif repo_id in ("black-forest-labs/FLUX.2-klein-9B", "black-forest-labs/FLUX.2-klein-9b"): + repo_id = "black-forest-labs/FLUX.2-klein-9b-kv" + pyconfig._config.keys["pretrained_model_name_or_path"] = repo_id + max_logging.log(f"ℹī¸ use_kv=True: switched pretrained_model_name_or_path to KV model variant: {repo_id}") + + 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: + 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") + 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 + from maxdiffusion.max_utils import get_flash_block_sizes + from flax import nnx + from maxdiffusion.models.flux.transformers.transformer_flux_flax import NNXFlux2KleinTransformer2DModel + from maxdiffusion.models.flux.util import load_and_convert_flux_klein_nnx_weights + + pt_config = AutoConfig.from_pretrained(text_encoder_path) + + te_bs = get_flash_block_sizes( + type( + "Config", + (), + { + "flash_block_sizes": getattr(config, "text_encoder_flash_block_sizes", {}) or {}, + "attention": getattr(config, "text_encoder_attention", "flash"), + }, + )() + ) + + 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=getattr(pt_config, "num_key_value_heads", pt_config.num_attention_heads), + max_position_embeddings=getattr(pt_config, "max_position_embeddings", 32768), + rms_norm_eps=getattr(pt_config, "rms_norm_eps", 1e-6), + rope_theta=getattr(pt_config, "rope_theta", getattr(pt_config, "rope_base", 1000000.0)), + dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + attention_kernel=getattr(config, "text_encoder_attention", "flash"), + flash_block_sizes=te_bs, + mesh=mesh, + ulysses_shards=getattr(config, "ulysses_shards", -1), + ulysses_attention_chunks=getattr(config, "ulysses_attention_chunks", 1), + max_layer_to_run=getattr(config, "text_encoder_max_layer", 27), + is_causal=getattr(config, "text_encoder_is_causal", True), + ) + qwen3_model = FlaxQwen3Model(qwen3_config) + + # Load Transformer config for layer counts if present + transformer_pt_cfg = {} + transformer_config_json = os.path.join(safetensors_path, "config.json") + if os.path.exists(transformer_config_json): + try: + import json + + with open(transformer_config_json, "r") as f: + transformer_pt_cfg = json.load(f) + except Exception: + pass + + num_double_layers = getattr(config, "num_double_layers", None) or transformer_pt_cfg.get("num_layers", 5) + depth = getattr(config, "depth", None) or transformer_pt_cfg.get("num_single_layers", 20) + num_attention_heads = getattr(config, "num_attention_heads", None) or transformer_pt_cfg.get("num_attention_heads", 24) + + # 5. Instantiate JAX NNXFlux2KleinTransformer2DModel + transformer = NNXFlux2KleinTransformer2DModel( + rngs=nnx.Rngs(0), + 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, + guidance_embeds=True, + axes_dim=(32, 32, 32, 32), + theta=2000.0, + mlp_ratio=3.0, + attention_kernel=config.attention, + flash_min_seq_length=512, + flash_block_sizes=get_flash_block_sizes(config), + 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, + scale_shift_order=getattr(config, "scale_shift_order", "scale_shift"), + ulysses_shards=getattr(config, "ulysses_shards", -1), + ulysses_attention_chunks=getattr(config, "ulysses_attention_chunks", 1), + use_base2_exp=getattr(config, "use_base2_exp", True), + ) + + # 6. Instantiate JAX NNX VAE + vae = NNXAutoencoderKLFlux2( + in_channels=3, + out_channels=3, + latent_channels=32, + block_out_channels=(128, 256, 512, 512), + layers_per_block=2, + norm_num_groups=32, + dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32, + param_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...") + seq_len_txt = config.max_sequence_length + 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) + qwen_key = jax.random.fold_in(key, 1) + + abstract_state = nnx.state(transformer, nnx.Param) + abstract_vae_state = nnx.state(vae, nnx.Param) + + def qwen3_init_fn(): + return qwen3_model.init(qwen_key, dummy_ids, dummy_mask) + + with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): + logical_transformer_specs = nnx.get_partition_spec(abstract_state) + logical_vae_specs = nnx.get_partition_spec(abstract_vae_state) + abstract_qwen3_vars = jax.eval_shape(qwen3_init_fn) + 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) + + vae_shardings = vae_mesh_shardings + qwen3_shardings = flax.core.freeze(qwen3_mesh_shardings["params"]) + transformer_shardings = transformer_mesh_shardings + + # 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() + 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_nnx_weights( + safetensors_path, abstract_state, num_double_layers, depth, dtype=weight_dtype + ) + vae_bn_mean, vae_bn_std = load_and_convert_flux2klein_nnx_vae_weights(vae_safetensors_path, vae, dtype=weight_dtype) + vae_params = nnx.state(vae, nnx.Param) + 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" + ) + + 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): + params = jax.tree_util.tree_map(max_utils.device_put_replicated, params, transformer_shardings) + 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) + nnx.update(vae, vae_params) + 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 FlaxFlux2KleinProfilePipeline + max_logging.log("Instantiating JAX FlaxFlux2KleinProfilePipeline...") + pipeline = FlaxFlux2KleinProfilePipeline( + transformer=transformer, + vae=vae, + text_encoder=qwen3_model, + tokenizer=None, + scheduler=scheduler, + config=config, + mesh=mesh, + ) + + 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) + + # Parse reference image paths for multi-image editing if provided + images = None + image_paths = getattr(config, "image_paths", None) + if image_paths is not None: + if isinstance(image_paths, str) and image_paths.strip(): + import ast + + try: + image_paths = ast.literal_eval(image_paths) + except Exception: + image_paths = [p.strip() for p in image_paths.split(",") if p.strip()] + if isinstance(image_paths, (list, tuple)) and len(image_paths) > 0: + max_logging.log(f" -> Loading {len(image_paths)} reference image(s) for multi-image editing...") + images = [] + for p in image_paths: + try: + if not os.path.exists(p): + raise FileNotFoundError(f"Reference image file not found: {p}") + with Image.open(p) as img_raw: + img = img_raw.convert("RGB") + images.append(img) + except (UnidentifiedImageError, OSError, FileNotFoundError) as e: + max_logging.log(f"❌ Error loading reference image '{p}': {e}") + raise ValueError(f"Failed to load reference image '{p}': {e}") from e + except Exception as e: + max_logging.log(f"❌ Unexpected error loading reference image '{p}': {e}") + raise ValueError(f"Failed to load reference image '{p}': {e}") from e + + 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, + images=images, + use_kv=getattr(config, "use_kv", False), + ) + + 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, + images=images, + use_latents=use_latents_flag, + latents=latents_to_use, + use_kv=getattr(config, "use_kv", False), + output_dir=config.output_dir, + output_name="flux2klein_warmup.png", + warmup=True, + profile_target=None, + ) + warmup_time = ( + warmup_trace.get("vae_encode", 0.0) + + 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"🚀 Capturing isolated XProf traces (profile_target={profile_target}) into: {tensorboard_dir}") + 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}...") + + current_profile_target = profile_target if rep == 0 else None + + _, 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, + images=images, + use_latents=use_latents_flag, + latents=latents_to_use, + use_kv=getattr(config, "use_kv", False), + output_dir=config.output_dir, + output_name=f"rep_{rep+1}_{config.output_name}" if num_reps > 1 else config.output_name, + profile_target=current_profile_target, + tensorboard_dir=tensorboard_dir, + ) + + tot_time_i = trace_i.get( + "e2e_pipeline_total", + trace_i.get("vae_encode", 0.0) + + 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: + vae_enc_str = f" | VAE_Enc={trace_i.get('vae_encode', 0.0):.4f}s" if trace_i.get("vae_encode", 0.0) > 0 else "" + max_logging.log( + f" -> Rep {rep+1}/{num_reps} Completed: Total={tot_time_i:.4f}s{vae_enc_str} | Qwen3={trace_i.get('qwen3_encoding', 0.0):.4f}s | Denoise={trace_i.get('denoise_loop', 0.0):.4f}s | VAE_Dec={trace_i.get('vae_decode', 0.0):.4f}s" + ) + + avg_main_time = sum(main_times) / num_reps + avg_vae_encode = sum(tr.get("vae_encode", 0.0) for tr in main_traces) / num_reps + avg_vae_to_qwen3 = sum(tr.get("vae_encode_to_qwen3", 0.0) for tr in main_traces) / 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 ⏱ī¸") + if warmup_trace.get("vae_encode", 0.0) > 0: + max_logging.log(f" - VAE Encoding: {warmup_trace.get('vae_encode', 0.0):.4f}s") + 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 ⏱ī¸") + step_num = 1 + if avg_vae_encode > 0: + max_logging.log(f" - {step_num}. VAE Image Encoding: {avg_vae_encode*1000:.2f} ms ({avg_vae_encode:.4f}s)") + step_num += 1 + max_logging.log(f" - {step_num}. VAE -> Qwen3: {avg_vae_to_qwen3*1000:.2f} ms ({avg_vae_to_qwen3:.4f}s)") + step_num += 1 + else: + max_logging.log( + f" - {step_num}. Start -> Qwen3: {avg_start_to_qwen3*1000:.2f} ms ({avg_start_to_qwen3:.4f}s)" + ) + step_num += 1 + max_logging.log(f" - {step_num}. Qwen3 Encoding: {avg_prompt_enc*1000:.2f} ms ({avg_prompt_enc:.4f}s)") + step_num += 1 + max_logging.log( + f" - {step_num}. Qwen3 -> Denoising: {avg_qwen3_to_denoise*1000:.2f} ms ({avg_qwen3_to_denoise:.4f}s)" + ) + step_num += 1 + max_logging.log(f" - {step_num}. Flux Denoising Loop: {avg_denoise*1000:.2f} ms ({avg_denoise:.4f}s)") + step_num += 1 + max_logging.log(f" - {step_num}. Denoising -> VAE: {avg_denoise_to_vae*1000:.2f} ms ({avg_denoise_to_vae:.4f}s)") + step_num += 1 + max_logging.log(f" - {step_num}. VAE Decoding: {avg_vae_decode*1000:.2f} ms ({avg_vae_decode:.4f}s)") + step_num += 1 + max_logging.log(f" - {step_num}. 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(f"📁 Isolated XProf Traces written to: {tensorboard_dir}") + 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/generate_flux_multi_res.py b/src/maxdiffusion/generate_flux_multi_res.py index c74533d42..3b3f87e13 100644 --- a/src/maxdiffusion/generate_flux_multi_res.py +++ b/src/maxdiffusion/generate_flux_multi_res.py @@ -27,8 +27,13 @@ import jax.numpy as jnp import flax.linen as nn from chex import Array -from flax.linen import partitioning as nn_partitioning -from transformers import (CLIPTokenizer, FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel, AutoTokenizer) +from transformers import CLIPTokenizer, AutoTokenizer +try: + from transformers import FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel +except ImportError: + FlaxCLIPTextModel = None + T5EncoderModel = None + FlaxT5EncoderModel = None from maxdiffusion import FlaxAutoencoderKL, pyconfig, max_logging, max_utils from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel diff --git a/src/maxdiffusion/models/embeddings_flax.py b/src/maxdiffusion/models/embeddings_flax.py index 61e7956ce..17ac2b07a 100644 --- a/src/maxdiffusion/models/embeddings_flax.py +++ b/src/maxdiffusion/models/embeddings_flax.py @@ -615,7 +615,7 @@ def __init__( weights_dtype=weights_dtype, ) - if pooled_projection_dim > 0: + if pooled_projection_dim is not None and pooled_projection_dim > 0: self.pooled_embedder = NNXPixArtAlphaTextProjection( rngs=rngs, in_features=pooled_projection_dim, @@ -643,7 +643,7 @@ def __call__( else: time_guidance_emb = timestep_emb - if pooled_projection is not None and self.pooled_projection_dim > 0: + if pooled_projection is not None and self.pooled_projection_dim is not None and self.pooled_projection_dim > 0: pooled_projections = self.pooled_embedder(pooled_projection) conditioning = time_guidance_emb + pooled_projections else: diff --git a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py index 42bfca5d3..3183f462c 100644 --- a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py +++ b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py @@ -1333,6 +1333,7 @@ def __init__( qkv_bias: bool = False, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + use_base2_exp: bool = False, ): self.heads = heads self.dim_head = dim_head @@ -1352,6 +1353,7 @@ def __init__( split_head_dim=False, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, ) kernel_axes = ("embed", "heads") @@ -1504,6 +1506,7 @@ def __init__( weights_dtype: jnp.dtype = jnp.float32, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + use_base2_exp: bool = False, ): self.num_attention_heads = num_attention_heads self.attention_head_dim = attention_head_dim @@ -1522,6 +1525,7 @@ def __init__( split_head_dim=False, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, ) self.query_norm = nnx.RMSNorm( num_features=attention_head_dim, @@ -1560,6 +1564,7 @@ def __init__( qkv_bias: bool = False, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + use_base2_exp: bool = False, ): self.dim = dim self.num_heads = num_attention_heads @@ -1616,6 +1621,7 @@ def __init__( qkv_bias=qkv_bias, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, ) self.ff = NNXFlaxSwiGluFeedForward( @@ -1703,6 +1709,7 @@ def __init__( weights_dtype: jnp.dtype = jnp.float32, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + use_base2_exp: bool = False, ): self.dim = dim self.num_attention_heads = num_attention_heads @@ -1753,6 +1760,7 @@ def __init__( weights_dtype=weights_dtype, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, ) def __call__( @@ -1834,6 +1842,7 @@ def __init__( scale_shift_order: str = "scale_shift", ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + use_base2_exp: bool = False, ): self.in_channels = in_channels self.out_channels = in_channels @@ -1914,6 +1923,7 @@ def __init__( weights_dtype=weights_dtype, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, ) for _ in range(num_layers) ] @@ -1935,6 +1945,7 @@ def __init__( weights_dtype=weights_dtype, ulysses_shards=ulysses_shards, ulysses_attention_chunks=ulysses_attention_chunks, + use_base2_exp=use_base2_exp, ) for _ in range(num_single_layers) ] diff --git a/src/maxdiffusion/models/flux/util.py b/src/maxdiffusion/models/flux/util.py index aa43609a1..5eb25572c 100644 --- a/src/maxdiffusion/models/flux/util.py +++ b/src/maxdiffusion/models/flux/util.py @@ -734,6 +734,74 @@ def set_val(var, tensor): return nnx.from_flat_state(flat_state) +def patchify_latents(latents): + """Patchifies latents: (B, C, H, W) -> (B, C*4, H//2, W//2).""" + import jax.numpy as jnp + + batch_size, num_channels, height, width = latents.shape + latents = latents.reshape((batch_size, num_channels, height // 2, 2, width // 2, 2)) + latents = jnp.transpose(latents, (0, 1, 3, 5, 2, 4)) + latents = latents.reshape((batch_size, num_channels * 4, height // 2, width // 2)) + return latents + + +def prepare_multi_image_ids(image_latents_list, scale=10): + """Generates 4D position IDs (T, H, W, L) for a sequence of reference image latents. + + For the k-th image, T = scale * (k + 1). + image_latents_list: list of arrays with shape (1, C, H, W) or (C, H, W). + Returns: array of shape (1, total_tokens, 4). + """ + import jax.numpy as jnp + + all_ids = [] + for idx, latent in enumerate(image_latents_list): + if latent.ndim == 4: + latent = latent[0] + _, h, w = latent.shape + t_val = scale * (idx + 1) + t = jnp.full((h * w, 1), t_val, dtype=jnp.int32) + h_grid, w_grid = jnp.meshgrid(jnp.arange(h, dtype=jnp.int32), jnp.arange(w, dtype=jnp.int32), indexing="ij") + h_coords = h_grid.reshape(-1, 1) + w_coords = w_grid.reshape(-1, 1) + l_coords = jnp.zeros((h * w, 1), dtype=jnp.int32) + coords = jnp.concatenate([t, h_coords, w_coords, l_coords], axis=-1) + all_ids.append(coords) + combined = jnp.concatenate(all_ids, axis=0) + return jnp.expand_dims(combined, axis=0) + + +def prepare_image_latents(vae, images, bn_mean, bn_std, scale=10): + """Encodes, patchifies, normalizes, packs, and generates 4D RoPE IDs for a list of reference images. + + images: list of arrays of shape (1, 3, H_k, W_k) or (3, H_k, W_k) in range [-1, 1]. + Returns: + image_latents_concat: shape (1, total_ref_tokens, 128) + image_latent_ids: shape (1, total_ref_tokens, 4) + """ + import jax.numpy as jnp + from einops import rearrange + + norm_latents = [] + for img in images: + if img.ndim == 3: + img = jnp.expand_dims(img, axis=0) + raw_latents = vae.encode(img) # (1, 32, H/8, W/8) + patchified = patchify_latents(raw_latents) # (1, 128, H/16, W/16) + normalized = (patchified - bn_mean) / bn_std + norm_latents.append(normalized) + + image_latent_ids = prepare_multi_image_ids(norm_latents, scale=scale) + + packed_latents = [] + for latent in norm_latents: + packed = rearrange(latent, "b c h w -> b (h w) c") + packed_latents.append(packed) + + image_latents_concat = jnp.concatenate(packed_latents, axis=1) + return image_latents_concat, image_latent_ids + + 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 @@ -756,62 +824,49 @@ def get_pytorch_weight_tensor(key, dtype_val=target_dtype): 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...") + # 1. Map VAE Encoder Weights + if "encoder" in jax_params: + max_logging.log("Mapping VAE encoder weights to JAX parameters...") + enc_jax = jax_params["encoder"] + + if "encoder.conv_in.weight" in pt_state_dict: + enc_jax["conv_in"]["kernel"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_in.weight").transpose(2, 3, 1, 0)) + enc_jax["conv_in"]["bias"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_in.bias")) + + for b_idx in range(4): + down_block_pt = f"encoder.down_blocks.{b_idx}" + down_block_jax = enc_jax[f"down_blocks_{b_idx}"] + + for r_idx in range(2): + res_pt = f"{down_block_pt}.resnets.{r_idx}" + res_jax = down_block_jax[f"resnets_{r_idx}"] + + res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.weight")) + res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.bias")) + res_jax["conv1"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.weight").transpose(2, 3, 1, 0)) + res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.bias")) + + res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.weight")) + res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.bias")) + res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.weight").transpose(2, 3, 1, 0)) + res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.bias")) + + shortcut_key = f"{res_pt}.conv_shortcut.weight" + if shortcut_key in pt_state_dict: + res_jax["conv_shortcut"]["kernel"] = jnp.array(get_pytorch_weight_tensor(shortcut_key).transpose(2, 3, 1, 0)) + res_jax["conv_shortcut"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv_shortcut.bias")) + + if b_idx < 3: + downsampler_pt = f"{down_block_pt}.downsamplers.0" + downsampler_jax = down_block_jax["downsamplers_0"] + downsampler_jax["conv"]["kernel"] = jnp.array( + get_pytorch_weight_tensor(f"{downsampler_pt}.conv.weight").transpose(2, 3, 1, 0) + ) + downsampler_jax["conv"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{downsampler_pt}.conv.bias")) - # post_quant_conv - jax_params["post_quant_conv"]["kernel"] = jnp.array( - get_pytorch_weight_tensor("post_quant_conv.weight").transpose(2, 3, 1, 0) - ) - jax_params["post_quant_conv"]["bias"] = jnp.array(get_pytorch_weight_tensor("post_quant_conv.bias")) - - # decoder.conv_in - jax_params["decoder"]["conv_in"]["kernel"] = jnp.array( - get_pytorch_weight_tensor("decoder.conv_in.weight").transpose(2, 3, 1, 0) - ) - jax_params["decoder"]["conv_in"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.bias")) - - # decoder.mid_block - # resnets - for idx in [0, 1]: - res_jax = jax_params["decoder"]["mid_block"][f"resnets_{idx}"] - res_pt_prefix = f"decoder.mid_block.resnets.{idx}" - - res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.weight")) - res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.bias")) - res_jax["conv1"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.weight").transpose(2, 3, 1, 0)) - res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.bias")) - - res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.weight")) - res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.bias")) - res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.weight").transpose(2, 3, 1, 0)) - res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.bias")) - - # attentions - attn_pt_prefix = "decoder.mid_block.attentions.0" - attn_jax = jax_params["decoder"]["mid_block"]["attentions_0"] - - attn_jax["group_norm"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.weight")) - attn_jax["group_norm"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.bias")) - - attn_jax["query"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.weight").T) - attn_jax["query"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.bias")) - attn_jax["key"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.weight").T) - attn_jax["key"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.bias")) - attn_jax["value"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.weight").T) - attn_jax["value"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.bias")) - - attn_jax["proj_attn"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.weight").T) - attn_jax["proj_attn"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.bias")) - - # decoder.up_blocks - for b_idx in range(4): - up_block_jax = jax_params["decoder"][f"up_blocks_{b_idx}"] - up_block_pt = f"decoder.up_blocks.{b_idx}" - - for r_idx in range(3): - res_jax = up_block_jax[f"resnets_{r_idx}"] - res_pt = f"{up_block_pt}.resnets.{r_idx}" + for r_idx in range(2): + res_pt = f"encoder.mid_block.resnets.{r_idx}" + res_jax = enc_jax["mid_block"][f"resnets_{r_idx}"] res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.weight")) res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.bias")) @@ -823,27 +878,112 @@ def get_pytorch_weight_tensor(key, dtype_val=target_dtype): res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.weight").transpose(2, 3, 1, 0)) res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.bias")) - shortcut_key = f"{res_pt}.conv_shortcut.weight" - if shortcut_key in pt_state_dict: - res_jax["conv_shortcut"]["kernel"] = jnp.array(get_pytorch_weight_tensor(shortcut_key).transpose(2, 3, 1, 0)) - res_jax["conv_shortcut"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv_shortcut.bias")) + attn_enc_pt = "encoder.mid_block.attentions.0" + attn_enc_jax = enc_jax["mid_block"]["attentions_0"] + attn_enc_jax["group_norm"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.group_norm.weight")) + attn_enc_jax["group_norm"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.group_norm.bias")) + attn_enc_jax["query"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_q.weight").T) + attn_enc_jax["query"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_q.bias")) + attn_enc_jax["key"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_k.weight").T) + attn_enc_jax["key"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_k.bias")) + attn_enc_jax["value"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_v.weight").T) + attn_enc_jax["value"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_v.bias")) + attn_enc_jax["proj_attn"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_out.0.weight").T) + attn_enc_jax["proj_attn"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_out.0.bias")) + + enc_jax["conv_norm_out"]["scale"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_norm_out.weight")) + enc_jax["conv_norm_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_norm_out.bias")) + enc_jax["conv_out"]["kernel"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_out.weight").transpose(2, 3, 1, 0)) + enc_jax["conv_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_out.bias")) + + if "quant_conv" in jax_params and "quant_conv.weight" in pt_state_dict: + jax_params["quant_conv"]["kernel"] = jnp.array(get_pytorch_weight_tensor("quant_conv.weight").transpose(2, 3, 1, 0)) + jax_params["quant_conv"]["bias"] = jnp.array(get_pytorch_weight_tensor("quant_conv.bias")) + + # 2. Map VAE Decoder Weights + max_logging.log("Mapping VAE decoder weights to JAX parameters...") + + if "post_quant_conv" in jax_params and "post_quant_conv.weight" in pt_state_dict: + jax_params["post_quant_conv"]["kernel"] = jnp.array( + get_pytorch_weight_tensor("post_quant_conv.weight").transpose(2, 3, 1, 0) + ) + jax_params["post_quant_conv"]["bias"] = jnp.array(get_pytorch_weight_tensor("post_quant_conv.bias")) + + if "decoder" in jax_params: + dec_jax = jax_params["decoder"] + dec_jax["conv_in"]["kernel"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.weight").transpose(2, 3, 1, 0)) + dec_jax["conv_in"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.bias")) - if b_idx < 3: - upsampler_jax = up_block_jax["upsamplers_0"] - upsampler_pt = f"{up_block_pt}.upsamplers.0" + for idx in [0, 1]: + res_jax = dec_jax["mid_block"][f"resnets_{idx}"] + res_pt_prefix = f"decoder.mid_block.resnets.{idx}" - upsampler_jax["conv"]["kernel"] = jnp.array( - get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0) + res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.weight")) + res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.bias")) + res_jax["conv1"]["kernel"] = jnp.array( + get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.weight").transpose(2, 3, 1, 0) ) - upsampler_jax["conv"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{upsampler_pt}.conv.bias")) + res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.bias")) - # decoder.conv_norm_out & conv_out - jax_params["decoder"]["conv_norm_out"]["scale"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.weight")) - jax_params["decoder"]["conv_norm_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.bias")) - jax_params["decoder"]["conv_out"]["kernel"] = jnp.array( - get_pytorch_weight_tensor("decoder.conv_out.weight").transpose(2, 3, 1, 0) - ) - jax_params["decoder"]["conv_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.bias")) + res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.weight")) + res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.bias")) + res_jax["conv2"]["kernel"] = jnp.array( + get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.weight").transpose(2, 3, 1, 0) + ) + res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.bias")) + + attn_pt_prefix = "decoder.mid_block.attentions.0" + attn_jax = dec_jax["mid_block"]["attentions_0"] + + attn_jax["group_norm"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.weight")) + attn_jax["group_norm"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.bias")) + + attn_jax["query"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.weight").T) + attn_jax["query"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.bias")) + attn_jax["key"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.weight").T) + attn_jax["key"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.bias")) + attn_jax["value"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.weight").T) + attn_jax["value"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.bias")) + + attn_jax["proj_attn"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.weight").T) + attn_jax["proj_attn"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.bias")) + + for b_idx in range(4): + up_block_jax = dec_jax[f"up_blocks_{b_idx}"] + up_block_pt = f"decoder.up_blocks.{b_idx}" + + for r_idx in range(3): + res_jax = up_block_jax[f"resnets_{r_idx}"] + res_pt = f"{up_block_pt}.resnets.{r_idx}" + + res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.weight")) + res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.bias")) + res_jax["conv1"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.weight").transpose(2, 3, 1, 0)) + res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.bias")) + + res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.weight")) + res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.bias")) + res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.weight").transpose(2, 3, 1, 0)) + res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.bias")) + + shortcut_key = f"{res_pt}.conv_shortcut.weight" + if shortcut_key in pt_state_dict: + res_jax["conv_shortcut"]["kernel"] = jnp.array(get_pytorch_weight_tensor(shortcut_key).transpose(2, 3, 1, 0)) + res_jax["conv_shortcut"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv_shortcut.bias")) + + if b_idx < 3: + upsampler_jax = up_block_jax["upsamplers_0"] + upsampler_pt = f"{up_block_pt}.upsamplers.0" + + upsampler_jax["conv"]["kernel"] = jnp.array( + get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0) + ) + upsampler_jax["conv"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{upsampler_pt}.conv.bias")) + + dec_jax["conv_norm_out"]["scale"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.weight")) + dec_jax["conv_norm_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.bias")) + dec_jax["conv_out"]["kernel"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.weight").transpose(2, 3, 1, 0)) + dec_jax["conv_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.bias")) jax_params = jax.tree_util.tree_map( lambda leaf: jnp.zeros(leaf.shape, dtype=leaf.dtype) if isinstance(leaf, jax.ShapeDtypeStruct) else leaf, jax_params diff --git a/src/maxdiffusion/models/flux/vae/__init__.py b/src/maxdiffusion/models/flux/vae/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/maxdiffusion/models/flux/vae/autoencoder_kl_flux2_nnx.py b/src/maxdiffusion/models/flux/vae/autoencoder_kl_flux2_nnx.py new file mode 100644 index 000000000..01adabd1f --- /dev/null +++ b/src/maxdiffusion/models/flux/vae/autoencoder_kl_flux2_nnx.py @@ -0,0 +1,799 @@ +""" +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 math +from typing import Optional, Tuple + +import jax +import jax.numpy as jnp +from flax import nnx + + +class NNXUpsample2D(nnx.Module): + """2D Nearest-neighbor Upsample + Conv layer in NNX.""" + + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + out_channels = out_channels or in_channels + self.conv = nnx.Conv( + in_features=in_channels, + out_features=out_channels, + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + batch, height, width, channels = x.shape + x = jnp.broadcast_to(x[:, :, None, :, None, :], (batch, height, 2, width, 2, channels)) + x = jnp.reshape(x, (batch, height * 2, width * 2, channels)) + return self.conv(x) + + +class NNXDownsample2D(nnx.Module): + """2D Downsample layer with asymmetric padding in NNX.""" + + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + out_channels = out_channels or in_channels + self.conv = nnx.Conv( + in_features=in_channels, + out_features=out_channels, + kernel_size=(3, 3), + strides=(2, 2), + padding="VALID", + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + pad_width = ((0, 0), (0, 1), (0, 1), (0, 0)) + x = jnp.pad(x, pad_width) + return self.conv(x) + + +class NNXResnetBlock2D(nnx.Module): + """2D ResNet Block with GroupNorm and SiLU activations in NNX.""" + + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + groups: int = 32, + use_conv_shortcut: Optional[bool] = None, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + out_channels = out_channels or in_channels + self.in_channels = in_channels + self.out_channels = out_channels + + self.norm1 = nnx.GroupNorm( + num_groups=groups, + num_features=in_channels, + epsilon=1e-6, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.conv1 = nnx.Conv( + in_features=in_channels, + out_features=out_channels, + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.norm2 = nnx.GroupNorm( + num_groups=groups, + num_features=out_channels, + epsilon=1e-6, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.conv2 = nnx.Conv( + in_features=out_channels, + out_features=out_channels, + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + use_shortcut = (in_channels != out_channels) if use_conv_shortcut is None else use_conv_shortcut + if use_shortcut: + self.conv_shortcut = nnx.Conv( + in_features=in_channels, + out_features=out_channels, + kernel_size=(1, 1), + strides=(1, 1), + padding="VALID", + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + else: + self.conv_shortcut = None + + def __call__(self, x: jax.Array) -> jax.Array: + residual = self.conv_shortcut(x) if self.conv_shortcut is not None else x + h = self.norm1(x) + h = nnx.silu(h) + h = self.conv1(h) + h = self.norm2(h) + h = nnx.silu(h) + h = self.conv2(h) + return h + residual + + +class NNXAttentionBlock(nnx.Module): + """Self-Attention block with GroupNorm in NNX.""" + + def __init__( + self, + channels: int, + groups: int = 32, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + self.channels = channels + self.group_norm = nnx.GroupNorm( + num_groups=groups, + num_features=channels, + epsilon=1e-6, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.to_q = nnx.Linear( + in_features=channels, + out_features=channels, + use_bias=True, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.to_k = nnx.Linear( + in_features=channels, + out_features=channels, + use_bias=True, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.to_v = nnx.Linear( + in_features=channels, + out_features=channels, + use_bias=True, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.to_out = nnx.Linear( + in_features=channels, + out_features=channels, + use_bias=True, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + residual = x + b, h, w, c = x.shape + h_states = self.group_norm(x) + h_flat = h_states.reshape((b, h * w, c)) + + q = self.to_q(h_flat) + k = self.to_k(h_flat) + v = self.to_v(h_flat) + + scale = 1.0 / math.sqrt(c) + attn_weights = jnp.einsum("bqc,bkc->bqk", q * scale, k) + attn_weights = jax.nn.softmax(attn_weights, axis=-1) + + out = jnp.einsum("bqk,bkc->bqc", attn_weights, v) + out = self.to_out(out) + out = out.reshape((b, h, w, c)) + return out + residual + + +class NNXUNetMidBlock2D(nnx.Module): + """Mid-Block module in NNX with resnets and attention.""" + + def __init__( + self, + in_channels: int, + groups: int = 32, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + self.resnets_0 = NNXResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + groups=groups, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.attentions_0 = NNXAttentionBlock( + channels=in_channels, + groups=groups, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.resnets_1 = NNXResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + groups=groups, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.resnets_0(x) + x = self.attentions_0(x) + x = self.resnets_1(x) + return x + + +class NNXDownEncoderBlock2D(nnx.Module): + """Down-Encoder block containing ResNet layers and an optional Downsampler in NNX.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 2, + groups: int = 32, + add_downsample: bool = True, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + resnets = [] + for i in range(num_layers): + in_ch = in_channels if i == 0 else out_channels + resnets.append( + NNXResnetBlock2D( + in_channels=in_ch, + out_channels=out_channels, + groups=groups, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + ) + self.resnets = nnx.List(resnets) + + if add_downsample: + self.downsamplers_0 = NNXDownsample2D( + in_channels=out_channels, + out_channels=out_channels, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + else: + self.downsamplers_0 = None + + def __call__(self, x: jax.Array) -> jax.Array: + for resnet in self.resnets: + x = resnet(x) + if self.downsamplers_0 is not None: + x = self.downsamplers_0(x) + return x + + +class NNXUpDecoderBlock2D(nnx.Module): + """Up-Decoder block containing ResNet layers and an optional Upsampler in NNX.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 3, + groups: int = 32, + add_upsample: bool = True, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + resnets = [] + for i in range(num_layers): + in_ch = in_channels if i == 0 else out_channels + resnets.append( + NNXResnetBlock2D( + in_channels=in_ch, + out_channels=out_channels, + groups=groups, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + ) + self.resnets = nnx.List(resnets) + + if add_upsample: + self.upsamplers_0 = NNXUpsample2D( + in_channels=out_channels, + out_channels=out_channels, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + else: + self.upsamplers_0 = None + + def __call__(self, x: jax.Array) -> jax.Array: + for resnet in self.resnets: + x = resnet(x) + if self.upsamplers_0 is not None: + x = self.upsamplers_0(x) + return x + + +class NNXEncoder(nnx.Module): + """Complete VAE Encoder in NNX.""" + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 32, + block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + self.conv_in = nnx.Conv( + in_features=in_channels, + out_features=block_out_channels[0], + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + down_blocks = [] + output_ch = block_out_channels[0] + for i, ch in enumerate(block_out_channels): + input_ch = output_ch + output_ch = ch + is_final = i == len(block_out_channels) - 1 + down_blocks.append( + NNXDownEncoderBlock2D( + in_channels=input_ch, + out_channels=output_ch, + num_layers=layers_per_block, + groups=norm_num_groups, + add_downsample=not is_final, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + ) + self.down_blocks = nnx.List(down_blocks) + + self.mid_block = NNXUNetMidBlock2D( + in_channels=block_out_channels[-1], + groups=norm_num_groups, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + self.conv_norm_out = nnx.GroupNorm( + num_groups=norm_num_groups, + num_features=block_out_channels[-1], + epsilon=1e-6, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.conv_out = nnx.Conv( + in_features=block_out_channels[-1], + out_features=2 * out_channels, # double_z for Gaussian distribution moments + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.conv_in(x) + for block in self.down_blocks: + x = block(x) + x = self.mid_block(x) + x = self.conv_norm_out(x) + x = nnx.silu(x) + x = self.conv_out(x) + return x + + +class NNXDecoder(nnx.Module): + """Complete VAE Decoder in NNX.""" + + def __init__( + self, + in_channels: int = 32, + out_channels: int = 3, + block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 3, + norm_num_groups: int = 32, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + reversed_channels = list(reversed(block_out_channels)) + self.conv_in = nnx.Conv( + in_features=in_channels, + out_features=reversed_channels[0], + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + self.mid_block = NNXUNetMidBlock2D( + in_channels=reversed_channels[0], + groups=norm_num_groups, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + up_blocks = [] + output_ch = reversed_channels[0] + for i, ch in enumerate(reversed_channels): + input_ch = output_ch + output_ch = ch + is_final = i == len(reversed_channels) - 1 + up_blocks.append( + NNXUpDecoderBlock2D( + in_channels=input_ch, + out_channels=output_ch, + num_layers=layers_per_block, + groups=norm_num_groups, + add_upsample=not is_final, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + ) + self.up_blocks = nnx.List(up_blocks) + + self.conv_norm_out = nnx.GroupNorm( + num_groups=norm_num_groups, + num_features=reversed_channels[-1], + epsilon=1e-6, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.conv_out = nnx.Conv( + in_features=reversed_channels[-1], + out_features=out_channels, + kernel_size=(3, 3), + strides=(1, 1), + padding=((1, 1), (1, 1)), + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.conv_in(x) + x = self.mid_block(x) + for block in self.up_blocks: + x = block(x) + x = self.conv_norm_out(x) + x = nnx.silu(x) + x = self.conv_out(x) + return x + + +class NNXAutoencoderKLFlux2(nnx.Module): + """Full FLUX.2-Klein Variational Autoencoder (VAE) in Flax NNX.""" + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + latent_channels: int = 32, + block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + rngs: Optional[nnx.Rngs] = None, + dtype: jnp.dtype = jnp.float32, + param_dtype: jnp.dtype = jnp.float32, + ): + rngs = rngs or nnx.Rngs(0) + self.latent_channels = latent_channels + self.dtype = dtype + + self.encoder = NNXEncoder( + in_channels=in_channels, + out_channels=latent_channels, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.quant_conv = nnx.Conv( + in_features=2 * latent_channels, + out_features=2 * latent_channels, + kernel_size=(1, 1), + strides=(1, 1), + padding="VALID", + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.post_quant_conv = nnx.Conv( + in_features=latent_channels, + out_features=latent_channels, + kernel_size=(1, 1), + strides=(1, 1), + padding="VALID", + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + self.decoder = NNXDecoder( + in_channels=latent_channels, + out_channels=out_channels, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block + 1, # 3 resnet blocks in decoder + norm_num_groups=norm_num_groups, + rngs=rngs, + dtype=dtype, + param_dtype=param_dtype, + ) + + def encode(self, sample: jax.Array) -> jax.Array: + """Encodes image tensor of shape (B, 3, H, W) to mode latents of shape (B, 32, H/8, W/8).""" + # Transpose to channels last (B, H, W, 3) + x = jnp.transpose(sample, (0, 2, 3, 1)) + h = self.encoder(x) + moments = self.quant_conv(h) + # Extract mean / mode (first latent_channels) + mean, _ = jnp.split(moments, 2, axis=-1) # (B, H/8, W/8, 32) + # Transpose back to (B, 32, H/8, W/8) + return jnp.transpose(mean, (0, 3, 1, 2)) + + def decode(self, latents: jax.Array) -> jax.Array: + """Decodes latent tensor of shape (B, 32, H/8, W/8) to image tensor of shape (B, 3, H, W).""" + # Transpose to channels last (B, H/8, W/8, 32) + z = jnp.transpose(latents, (0, 2, 3, 1)) + h = self.post_quant_conv(z) + img = self.decoder(h) + # Transpose back to (B, 3, H, W) + return jnp.transpose(img, (0, 3, 1, 2)) + + +def load_and_convert_flux2klein_nnx_vae_weights( + safetensors_path: str, + nnx_vae: NNXAutoencoderKLFlux2, + dtype: Optional[jnp.dtype] = None, + pt_state_dict: Optional[dict] = None, +): + """Directly loads and maps PyTorch safetensors into NNXAutoencoderKLFlux2 State.""" + from safetensors.numpy import load_file + + if pt_state_dict is None: + pt_state_dict = load_file(safetensors_path) + + target_dtype = dtype if dtype is not None else jnp.float32 + + def get_pt_tensor(key, is_norm=False): + tensor = pt_state_dict[key] + leaf_dtype = jnp.float32 if is_norm else target_dtype + return jnp.array(tensor, dtype=leaf_dtype) + + def get_conv_kernel(key): + return jnp.array(pt_state_dict[key].transpose(2, 3, 1, 0), dtype=target_dtype) + + def get_linear_kernel(key): + return jnp.array(pt_state_dict[key].T, dtype=target_dtype) + + flat_state = dict(nnx.to_flat_state(nnx.state(nnx_vae, nnx.Param))) + + def set_val(var, val): + var[...] = val + + # ========================================================================= + # 1. ENCODER + # ========================================================================= + set_val(flat_state[("encoder", "conv_in", "kernel")], get_conv_kernel("encoder.conv_in.weight")) + set_val(flat_state[("encoder", "conv_in", "bias")], get_pt_tensor("encoder.conv_in.bias")) + + for b_idx in range(4): + down_block_pt = f"encoder.down_blocks.{b_idx}" + for r_idx in range(2): + res_pt = f"{down_block_pt}.resnets.{r_idx}" + res_path = ("encoder", "down_blocks", b_idx, "resnets", r_idx) + + set_val(flat_state[res_path + ("norm1", "scale")], get_pt_tensor(f"{res_pt}.norm1.weight", is_norm=True)) + set_val(flat_state[res_path + ("norm1", "bias")], get_pt_tensor(f"{res_pt}.norm1.bias", is_norm=True)) + set_val(flat_state[res_path + ("conv1", "kernel")], get_conv_kernel(f"{res_pt}.conv1.weight")) + set_val(flat_state[res_path + ("conv1", "bias")], get_pt_tensor(f"{res_pt}.conv1.bias")) + + set_val(flat_state[res_path + ("norm2", "scale")], get_pt_tensor(f"{res_pt}.norm2.weight", is_norm=True)) + set_val(flat_state[res_path + ("norm2", "bias")], get_pt_tensor(f"{res_pt}.norm2.bias", is_norm=True)) + set_val(flat_state[res_path + ("conv2", "kernel")], get_conv_kernel(f"{res_pt}.conv2.weight")) + set_val(flat_state[res_path + ("conv2", "bias")], get_pt_tensor(f"{res_pt}.conv2.bias")) + + shortcut_key = f"{res_pt}.conv_shortcut.weight" + if shortcut_key in pt_state_dict: + set_val(flat_state[res_path + ("conv_shortcut", "kernel")], get_conv_kernel(shortcut_key)) + set_val(flat_state[res_path + ("conv_shortcut", "bias")], get_pt_tensor(f"{res_pt}.conv_shortcut.bias")) + + if b_idx < 3: + ds_pt = f"{down_block_pt}.downsamplers.0.conv" + ds_path = ("encoder", "down_blocks", b_idx, "downsamplers_0", "conv") + set_val(flat_state[ds_path + ("kernel",)], get_conv_kernel(f"{ds_pt}.weight")) + set_val(flat_state[ds_path + ("bias",)], get_pt_tensor(f"{ds_pt}.bias")) + + # Encoder Mid Block + for r_idx in [0, 1]: + res_pt = f"encoder.mid_block.resnets.{r_idx}" + res_path = ("encoder", "mid_block", f"resnets_{r_idx}") + set_val(flat_state[res_path + ("norm1", "scale")], get_pt_tensor(f"{res_pt}.norm1.weight", is_norm=True)) + set_val(flat_state[res_path + ("norm1", "bias")], get_pt_tensor(f"{res_pt}.norm1.bias", is_norm=True)) + set_val(flat_state[res_path + ("conv1", "kernel")], get_conv_kernel(f"{res_pt}.conv1.weight")) + set_val(flat_state[res_path + ("conv1", "bias")], get_pt_tensor(f"{res_pt}.conv1.bias")) + set_val(flat_state[res_path + ("norm2", "scale")], get_pt_tensor(f"{res_pt}.norm2.weight", is_norm=True)) + set_val(flat_state[res_path + ("norm2", "bias")], get_pt_tensor(f"{res_pt}.norm2.bias", is_norm=True)) + set_val(flat_state[res_path + ("conv2", "kernel")], get_conv_kernel(f"{res_pt}.conv2.weight")) + set_val(flat_state[res_path + ("conv2", "bias")], get_pt_tensor(f"{res_pt}.conv2.bias")) + + attn_pt = "encoder.mid_block.attentions.0" + attn_path = ("encoder", "mid_block", "attentions_0") + set_val(flat_state[attn_path + ("group_norm", "scale")], get_pt_tensor(f"{attn_pt}.group_norm.weight", is_norm=True)) + set_val(flat_state[attn_path + ("group_norm", "bias")], get_pt_tensor(f"{attn_pt}.group_norm.bias", is_norm=True)) + set_val(flat_state[attn_path + ("to_q", "kernel")], get_linear_kernel(f"{attn_pt}.to_q.weight")) + set_val(flat_state[attn_path + ("to_q", "bias")], get_pt_tensor(f"{attn_pt}.to_q.bias")) + set_val(flat_state[attn_path + ("to_k", "kernel")], get_linear_kernel(f"{attn_pt}.to_k.weight")) + set_val(flat_state[attn_path + ("to_k", "bias")], get_pt_tensor(f"{attn_pt}.to_k.bias")) + set_val(flat_state[attn_path + ("to_v", "kernel")], get_linear_kernel(f"{attn_pt}.to_v.weight")) + set_val(flat_state[attn_path + ("to_v", "bias")], get_pt_tensor(f"{attn_pt}.to_v.bias")) + set_val(flat_state[attn_path + ("to_out", "kernel")], get_linear_kernel(f"{attn_pt}.to_out.0.weight")) + set_val(flat_state[attn_path + ("to_out", "bias")], get_pt_tensor(f"{attn_pt}.to_out.0.bias")) + + set_val(flat_state[("encoder", "conv_norm_out", "scale")], get_pt_tensor("encoder.conv_norm_out.weight", is_norm=True)) + set_val(flat_state[("encoder", "conv_norm_out", "bias")], get_pt_tensor("encoder.conv_norm_out.bias", is_norm=True)) + set_val(flat_state[("encoder", "conv_out", "kernel")], get_conv_kernel("encoder.conv_out.weight")) + set_val(flat_state[("encoder", "conv_out", "bias")], get_pt_tensor("encoder.conv_out.bias")) + + # ========================================================================= + # 2. QUANT CONV & POST QUANT CONV + # ========================================================================= + set_val(flat_state[("quant_conv", "kernel")], get_conv_kernel("quant_conv.weight")) + set_val(flat_state[("quant_conv", "bias")], get_pt_tensor("quant_conv.bias")) + set_val(flat_state[("post_quant_conv", "kernel")], get_conv_kernel("post_quant_conv.weight")) + set_val(flat_state[("post_quant_conv", "bias")], get_pt_tensor("post_quant_conv.bias")) + + # ========================================================================= + # 3. DECODER + # ========================================================================= + set_val(flat_state[("decoder", "conv_in", "kernel")], get_conv_kernel("decoder.conv_in.weight")) + set_val(flat_state[("decoder", "conv_in", "bias")], get_pt_tensor("decoder.conv_in.bias")) + + for r_idx in [0, 1]: + res_pt = f"decoder.mid_block.resnets.{r_idx}" + res_path = ("decoder", "mid_block", f"resnets_{r_idx}") + set_val(flat_state[res_path + ("norm1", "scale")], get_pt_tensor(f"{res_pt}.norm1.weight", is_norm=True)) + set_val(flat_state[res_path + ("norm1", "bias")], get_pt_tensor(f"{res_pt}.norm1.bias", is_norm=True)) + set_val(flat_state[res_path + ("conv1", "kernel")], get_conv_kernel(f"{res_pt}.conv1.weight")) + set_val(flat_state[res_path + ("conv1", "bias")], get_pt_tensor(f"{res_pt}.conv1.bias")) + set_val(flat_state[res_path + ("norm2", "scale")], get_pt_tensor(f"{res_pt}.norm2.weight", is_norm=True)) + set_val(flat_state[res_path + ("norm2", "bias")], get_pt_tensor(f"{res_pt}.norm2.bias", is_norm=True)) + set_val(flat_state[res_path + ("conv2", "kernel")], get_conv_kernel(f"{res_pt}.conv2.weight")) + set_val(flat_state[res_path + ("conv2", "bias")], get_pt_tensor(f"{res_pt}.conv2.bias")) + + dec_attn_pt = "decoder.mid_block.attentions.0" + dec_attn_path = ("decoder", "mid_block", "attentions_0") + set_val( + flat_state[dec_attn_path + ("group_norm", "scale")], get_pt_tensor(f"{dec_attn_pt}.group_norm.weight", is_norm=True) + ) + set_val(flat_state[dec_attn_path + ("group_norm", "bias")], get_pt_tensor(f"{dec_attn_pt}.group_norm.bias", is_norm=True)) + set_val(flat_state[dec_attn_path + ("to_q", "kernel")], get_linear_kernel(f"{dec_attn_pt}.to_q.weight")) + set_val(flat_state[dec_attn_path + ("to_q", "bias")], get_pt_tensor(f"{dec_attn_pt}.to_q.bias")) + set_val(flat_state[dec_attn_path + ("to_k", "kernel")], get_linear_kernel(f"{dec_attn_pt}.to_k.weight")) + set_val(flat_state[dec_attn_path + ("to_k", "bias")], get_pt_tensor(f"{dec_attn_pt}.to_k.bias")) + set_val(flat_state[dec_attn_path + ("to_v", "kernel")], get_linear_kernel(f"{dec_attn_pt}.to_v.weight")) + set_val(flat_state[dec_attn_path + ("to_v", "bias")], get_pt_tensor(f"{dec_attn_pt}.to_v.bias")) + set_val(flat_state[dec_attn_path + ("to_out", "kernel")], get_linear_kernel(f"{dec_attn_pt}.to_out.0.weight")) + set_val(flat_state[dec_attn_path + ("to_out", "bias")], get_pt_tensor(f"{dec_attn_pt}.to_out.0.bias")) + + for b_idx in range(4): + up_block_pt = f"decoder.up_blocks.{b_idx}" + for r_idx in range(3): + res_pt = f"{up_block_pt}.resnets.{r_idx}" + res_path = ("decoder", "up_blocks", b_idx, "resnets", r_idx) + + set_val(flat_state[res_path + ("norm1", "scale")], get_pt_tensor(f"{res_pt}.norm1.weight", is_norm=True)) + set_val(flat_state[res_path + ("norm1", "bias")], get_pt_tensor(f"{res_pt}.norm1.bias", is_norm=True)) + set_val(flat_state[res_path + ("conv1", "kernel")], get_conv_kernel(f"{res_pt}.conv1.weight")) + set_val(flat_state[res_path + ("conv1", "bias")], get_pt_tensor(f"{res_pt}.conv1.bias")) + + set_val(flat_state[res_path + ("norm2", "scale")], get_pt_tensor(f"{res_pt}.norm2.weight", is_norm=True)) + set_val(flat_state[res_path + ("norm2", "bias")], get_pt_tensor(f"{res_pt}.norm2.bias", is_norm=True)) + set_val(flat_state[res_path + ("conv2", "kernel")], get_conv_kernel(f"{res_pt}.conv2.weight")) + set_val(flat_state[res_path + ("conv2", "bias")], get_pt_tensor(f"{res_pt}.conv2.bias")) + + shortcut_key = f"{res_pt}.conv_shortcut.weight" + if shortcut_key in pt_state_dict: + set_val(flat_state[res_path + ("conv_shortcut", "kernel")], get_conv_kernel(shortcut_key)) + set_val(flat_state[res_path + ("conv_shortcut", "bias")], get_pt_tensor(f"{res_pt}.conv_shortcut.bias")) + + if b_idx < 3: + ups_pt = f"{up_block_pt}.upsamplers.0.conv" + ups_path = ("decoder", "up_blocks", b_idx, "upsamplers_0", "conv") + set_val(flat_state[ups_path + ("kernel",)], get_conv_kernel(f"{ups_pt}.weight")) + set_val(flat_state[ups_path + ("bias",)], get_pt_tensor(f"{ups_pt}.bias")) + + set_val(flat_state[("decoder", "conv_norm_out", "scale")], get_pt_tensor("decoder.conv_norm_out.weight", is_norm=True)) + set_val(flat_state[("decoder", "conv_norm_out", "bias")], get_pt_tensor("decoder.conv_norm_out.bias", is_norm=True)) + set_val(flat_state[("decoder", "conv_out", "kernel")], get_conv_kernel("decoder.conv_out.weight")) + set_val(flat_state[("decoder", "conv_out", "bias")], get_pt_tensor("decoder.conv_out.bias")) + + # Update nnx_vae state + nnx.update(nnx_vae, nnx.from_flat_state(flat_state)) + + # Extract Batch Normalization running stats + bn_mean = jnp.array(get_pt_tensor("bn.running_mean")).reshape(1, -1, 1, 1) + bn_var = jnp.array(get_pt_tensor("bn.running_var")).reshape(1, -1, 1, 1) + batch_norm_eps = 0.0001 + bn_std = jnp.sqrt(bn_var + batch_norm_eps) + + return bn_mean, bn_std diff --git a/src/maxdiffusion/pipelines/controlnet/pipeline_flax_controlnet.py b/src/maxdiffusion/pipelines/controlnet/pipeline_flax_controlnet.py index ecc6814d9..7c2162613 100644 --- a/src/maxdiffusion/pipelines/controlnet/pipeline_flax_controlnet.py +++ b/src/maxdiffusion/pipelines/controlnet/pipeline_flax_controlnet.py @@ -18,9 +18,11 @@ import jax import jax.numpy as jnp import numpy as np -from flax.core.frozen_dict import FrozenDict -from PIL import Image -from transformers import CLIPTokenizer, FlaxCLIPTextModel +from transformers import CLIPTokenizer +try: + from transformers import FlaxCLIPTextModel +except ImportError: + FlaxCLIPTextModel = None from ...models import FlaxAutoencoderKL, FlaxControlNetModel, FlaxUNet2DConditionModel from ...schedulers import ( diff --git a/src/maxdiffusion/pipelines/controlnet/pipeline_flax_controlnet_sdxl.py b/src/maxdiffusion/pipelines/controlnet/pipeline_flax_controlnet_sdxl.py index b8b1cc18b..8cdff06b9 100644 --- a/src/maxdiffusion/pipelines/controlnet/pipeline_flax_controlnet_sdxl.py +++ b/src/maxdiffusion/pipelines/controlnet/pipeline_flax_controlnet_sdxl.py @@ -17,8 +17,11 @@ import jax import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict -from transformers import CLIPTokenizer, FlaxCLIPTextModel +from transformers import CLIPTokenizer +try: + from transformers import FlaxCLIPTextModel +except ImportError: + FlaxCLIPTextModel = None from PIL import Image from maxdiffusion.utils import logging from ...utils import PIL_INTERPOLATION diff --git a/src/maxdiffusion/pipelines/flux/__init__.py b/src/maxdiffusion/pipelines/flux/__init__.py index 39ea05b57..c94ebd8a3 100644 --- a/src/maxdiffusion/pipelines/flux/__init__.py +++ b/src/maxdiffusion/pipelines/flux/__init__.py @@ -19,3 +19,6 @@ from .flux_pipeline import ( FluxPipeline, ) +from .flux2klein_pipeline import ( + FlaxFlux2KleinPipeline, +) diff --git a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py index 4fa257ca2..8ff868a0b 100644 --- a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py +++ b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py @@ -41,7 +41,9 @@ from ...models.flux.util import ( pack_latents, + patchify_latents, prepare_latent_image_ids, + prepare_multi_image_ids, prepare_text_ids, ) @@ -105,6 +107,7 @@ def __init__( # JIT compilation cache self._jitted_qwen3_forward = None self._jitted_transformer_step = None + self._jitted_vae_encode = None self._jitted_vae_decode = None def _setup_jit_functions(self): @@ -123,22 +126,57 @@ def qwen3_forward(q_params, ids, mask): prompt_embeds = jax.lax.with_sharding_constraint(prompt_embeds, jax.sharding.NamedSharding(self.mesh, context_spec)) return prompt_embeds - @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 + if isinstance(self.vae, nnx.Module): + v_graph, _, v_rest = nnx.split(self.vae, nnx.Param, ...) - vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128) - vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128) + @jax.jit + def vae_encode(v_params, img): + merged = nnx.merge(v_graph, v_params, v_rest) + return merged.encode(img) + + @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)) + + merged = nnx.merge(v_graph, v_params, v_rest) + res = merged.decode(latents_unpacked) + return FlaxDecoderOutput(sample=res) + + else: + + @jax.jit + def vae_encode(v_params, img): + # FlaxAutoencoderKL expects (B, 3, H, W) + res = self.vae.apply({"params": v_params}, sample=img, method=self.vae.encode) + moments = res.latent_dist.mode() + return jnp.transpose(moments, (0, 3, 1, 2)) + + @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 - 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)) + vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128) + vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128) - res = self.vae.apply({"params": v_params}, latents=latents_unpacked, method=self.vae.decode) - return FlaxDecoderOutput(sample=res.sample) + 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) if isinstance(self.transformer, nnx.Module): g, nnx_state, r = nnx.split(self.transformer, nnx.Param, ...) @@ -157,8 +195,10 @@ def transformer_step(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, ti return_dict=True, ) - @jax.jit - def fused_denoise_loop(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance): + @jax.jit(static_argnums=(9,)) + def fused_denoise_loop( + t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance, target_len=None + ): sigmas_padded = jnp.concatenate([sigmas, jnp.array([0.0], dtype=sigmas.dtype)]) nnx_merged = nnx.merge(g, t_params, r) @@ -177,7 +217,15 @@ def scan_body(cur_latents, step_idx): ) sigma = sigmas_padded[step_idx] sigma_next = sigmas_padded[step_idx + 1] - prev_sample = cur_latents + model_output.sample * (sigma_next - sigma) + dt = sigma_next - sigma + v = model_output.sample + if target_len is not None and cur_latents.shape[1] > target_len: + target_latents = cur_latents[:, :target_len, :] + v_target = v[:, :target_len, :] + next_target = target_latents + v_target * dt + prev_sample = jnp.concatenate([next_target, cur_latents[:, target_len:, :]], axis=1) + else: + prev_sample = cur_latents + v * dt return prev_sample, None steps = jnp.arange(timesteps.shape[0]) @@ -199,8 +247,10 @@ def transformer_step(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, ti guidance=guidance, ) - @jax.jit - def fused_denoise_loop(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance): + @jax.jit(static_argnums=(9,)) + def fused_denoise_loop( + t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance, target_len=None + ): sigmas_padded = jnp.concatenate([sigmas, jnp.array([0.0], dtype=sigmas.dtype)]) def scan_body(cur_latents, step_idx): @@ -218,7 +268,15 @@ def scan_body(cur_latents, step_idx): ) sigma = sigmas_padded[step_idx] sigma_next = sigmas_padded[step_idx + 1] - prev_sample = cur_latents + model_output.sample * (sigma_next - sigma) + dt = sigma_next - sigma + v = model_output.sample + if target_len is not None and cur_latents.shape[1] > target_len: + target_latents = cur_latents[:, :target_len, :] + v_target = v[:, :target_len, :] + next_target = target_latents + v_target * dt + prev_sample = jnp.concatenate([next_target, cur_latents[:, target_len:, :]], axis=1) + else: + prev_sample = cur_latents + v * dt return prev_sample, None steps = jnp.arange(timesteps.shape[0]) @@ -228,6 +286,7 @@ def scan_body(cur_latents, step_idx): self._jitted_qwen3_forward = qwen3_forward self._jitted_transformer_step = transformer_step self._jitted_fused_denoise_loop = fused_denoise_loop + self._jitted_vae_encode = vae_encode self._jitted_vae_decode = vae_decode def _get_dynamic_batch_sharding(self): @@ -237,25 +296,44 @@ def _get_dynamic_batch_sharding(self): 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 + self, + params, + vae_params, + qwen3_params, + vae_bn_mean, + vae_bn_std, + batch_size=1, + height=1024, + width=1024, + images=None, + image=None, + num_conditioning_images=0, ): """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 + if images is None and image is not None: + images = image if isinstance(image, list) else [image] + + if images is not None and len(images) > 0: + num_conditioning_images = len(images) + seq_len_img = (height // 16) * (width // 16) + total_img_len = (1 + num_conditioning_images) * seq_len_img 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_latents = jnp.zeros((batch_size, total_img_len, 128), dtype=jnp.float32) + dummy_img_ids = jnp.zeros((batch_size, total_img_len, 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_target_latents = jnp.zeros((batch_size, seq_len_img, 128), 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) @@ -277,6 +355,7 @@ def put_data_on_devices(x, sharding): dummy_prompt_embeds = put_data_on_devices(dummy_prompt_embeds, context_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_target_latents = put_data_on_devices(dummy_target_latents, 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) @@ -303,22 +382,32 @@ def compile_transformer(): dummy_timesteps, dummy_sigmas, None, + seq_len_img, ).compile() max_logging.log(f" -> [AOT COMPILED] Fused Flux Transformer Denoise Scan 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() + self._jitted_vae_decode.lower(vae_params, dummy_target_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") + def compile_vae_encode(): + t0 = time.perf_counter() + dummy_rgb = jnp.zeros((1, 3, height, width), dtype=jnp.float32) + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + self._jitted_vae_encode.lower(vae_params, dummy_rgb).compile() + max_logging.log(f" -> [AOT COMPILED] VAE Encoder in {time.perf_counter() - t0:.2f}s") + t_start = time.perf_counter() - with ThreadPoolExecutor(max_workers=3) as executor: + with ThreadPoolExecutor(max_workers=4) as executor: futures = [ executor.submit(compile_qwen3), executor.submit(compile_transformer), executor.submit(compile_vae), ] + if num_conditioning_images > 0 or (images is not None and len(images) > 0): + futures.append(executor.submit(compile_vae_encode)) for future in futures: future.result() aot_duration = time.perf_counter() - t_start @@ -364,6 +453,8 @@ def __call__( width: int = 1024, num_inference_steps: int = 4, batch_size: int = 1, + images: Optional[List[Any]] = None, + image: Optional[Union[Any, List[Any]]] = None, use_latents: bool = False, latents: Optional[Any] = None, measure_time: bool = False, @@ -375,6 +466,9 @@ def __call__( # 1. Setup JIT functions self._setup_jit_functions() + if images is None and image is not None: + images = image if isinstance(image, list) else [image] + # 2. Setup prompts and inputs if isinstance(prompt, str): prompts = [prompt] * batch_size @@ -392,6 +486,8 @@ def __call__( if C == 32: max_logging.log(" [PIPELINE] Unpacked 32-channel latents detected. Packing using pack_latents...") latents_jax = pack_latents(latents_jax) + elif C == 128: + latents_jax = jnp.transpose(jnp.reshape(latents_jax, (B, C, H * W)), (0, 2, 1)) else: latents_jax = jnp.transpose(jnp.reshape(latents_jax, (B, C, H * W)), (0, 2, 1)) else: @@ -401,7 +497,9 @@ def __call__( # 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) + target_img_ids_val = prepare_latent_image_ids(batch_size, height // 16, width // 16) + t_pipeline_start = time.perf_counter() + trace = {} # Scheduler mu = compute_empirical_mu(seq_len_img, num_inference_steps) @@ -414,9 +512,6 @@ def __call__( 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() @@ -432,9 +527,89 @@ def put_data_on_devices(x, sharding): return jax.device_put(x, sharding) return device_put_replicated(x, sharding) + # --------------------------------------------------------------------- + # PHASE 0: Encode Reference Images (VAE) + # --------------------------------------------------------------------- + if images is not None and len(images) > 0: + t0_vae_enc_start = time.perf_counter() + trace["start_to_vae_encode"] = t0_vae_enc_start - t_pipeline_start + max_logging.log(f"{host_prefix} [PHASE 0] Encoding {len(images)} reference image(s) using JAX VAE encoder on TPU...") + norm_ref_latents = [] + packed_ref_latents = [] + bn_mean_arr = jnp.array(vae_bn_mean, dtype=jnp.float32) + bn_std_arr = jnp.array(vae_bn_std, dtype=jnp.float32) + + for img in images: + if isinstance(img, Image.Image): + img = img.convert("RGB").resize((width, height), Image.Resampling.BICUBIC) + arr = np.array(img, dtype=np.float32) / 127.5 - 1.0 + arr = np.transpose(arr, (2, 0, 1)) + img_tensor = jnp.expand_dims(jnp.array(arr), axis=0) + elif isinstance(img, np.ndarray): + if img.ndim == 3: + img = np.expand_dims(img, axis=0) + if img.shape[-1] == 3: + img = np.transpose(img, (0, 3, 1, 2)) + if np.issubdtype(img.dtype, np.integer): + img = img.astype(np.float32) / 127.5 - 1.0 + elif np.issubdtype(img.dtype, np.floating): + if img.max() > 1.0: + img = img / 127.5 - 1.0 + elif img.min() >= 0.0: + img = img * 2.0 - 1.0 + img_tensor = jnp.array(img, dtype=np.float32) + elif isinstance(img, jnp.ndarray): + if img.ndim == 3: + img = jnp.expand_dims(img, axis=0) + if img.shape[-1] == 3: + img = jnp.transpose(img, (0, 3, 1, 2)) + if jnp.issubdtype(img.dtype, jnp.integer): + img = img.astype(jnp.float32) / 127.5 - 1.0 + elif jnp.issubdtype(img.dtype, jnp.floating): + if img.max() > 1.0: + img = img / 127.5 - 1.0 + elif img.min() >= 0.0: + img = img * 2.0 - 1.0 + img_tensor = img + else: + raise ValueError(f"Unsupported image type: {type(img)}") + + raw_ref_latents = self._jitted_vae_encode(vae_params, img_tensor) + raw_ref_latents.block_until_ready() + patchified_ref = patchify_latents(raw_ref_latents) + normalized_ref = (patchified_ref - bn_mean_arr) / bn_std_arr + norm_ref_latents.append(normalized_ref) + + packed = jnp.transpose( + jnp.reshape(normalized_ref, (normalized_ref.shape[0], normalized_ref.shape[1], -1)), (0, 2, 1) + ) + if packed.shape[0] == 1 and batch_size > 1: + packed = jnp.repeat(packed, batch_size, axis=0) + packed_ref_latents.append(packed) + + ref_img_ids_val = prepare_multi_image_ids(norm_ref_latents, scale=10) + if ref_img_ids_val.shape[0] == 1 and batch_size > 1: + ref_img_ids_val = jnp.repeat(ref_img_ids_val, batch_size, axis=0) + img_ids_val = jnp.concatenate([target_img_ids_val, ref_img_ids_val], axis=1) + latents_jax = jnp.concatenate([latents_jax] + packed_ref_latents, axis=1) + max_logging.log(f" [PIPELINE] Joint latents shape: {latents_jax.shape}, Joint img_ids shape: {img_ids_val.shape}") + + t0_vae_enc_end = time.perf_counter() + trace["vae_encode"] = t0_vae_enc_end - t0_vae_enc_start + trace["image_encoding"] = trace["vae_encode"] + max_logging.log(f" -> [TIMING] Reference Image Encoding (VAE): {trace['vae_encode']:.4f} seconds ⏱ī¸") + else: + img_ids_val = target_img_ids_val + trace["vae_encode"] = 0.0 + trace["image_encoding"] = 0.0 + 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 ⏱ī¸") + if trace.get("vae_encode", 0.0) > 0: + trace["vae_encode_to_qwen3"] = t0_qwen3_start - t0_vae_enc_end + max_logging.log(f" -> [TIMING] VAE Encode to Qwen3 Overhead: {trace['vae_encode_to_qwen3']:.4f} seconds ⏱ī¸") + else: + 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) @@ -543,6 +718,7 @@ def put_data_on_devices(x, sharding): timesteps_device, sigmas_device, guidance_vec_val, + seq_len_img, ) latents_jax.block_until_ready() if do_prof_denoise: @@ -569,6 +745,10 @@ def put_data_on_devices(x, sharding): # --------------------------------------------------------------------- max_logging.log("[PHASE C] Decoding final latents to RGB image using JAX VAE decoder on TPU...") + # Slice target latents from joint latents if reference images were present + if latents_jax.shape[1] > seq_len_img: + latents_jax = latents_jax[:, :seq_len_img, :] + # 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()) diff --git a/src/maxdiffusion/pipelines/flux/flux2klein_profile_pipeline.py b/src/maxdiffusion/pipelines/flux/flux2klein_profile_pipeline.py new file mode 100644 index 000000000..eb3871cc7 --- /dev/null +++ b/src/maxdiffusion/pipelines/flux/flux2klein_profile_pipeline.py @@ -0,0 +1,982 @@ +""" +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 flax import nnx +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, + NNXFlux2KleinTransformer2DModel, +) +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, + patchify_latents, + prepare_latent_image_ids, + prepare_multi_image_ids, + prepare_text_ids, +) + + +class FlaxFlux2KleinProfilePipeline(FlaxDiffusionPipeline): + """ + Dedicated profiling pipeline for Flux.2-klein-4B and 9B models on JAX+TPU. + Wraps individual components in isolated jax.profiler.start_trace/stop_trace sessions. + """ + + def __init__( + self, + transformer: Union[Flux2KleinTransformer2DModel, NNXFlux2KleinTransformer2DModel], + 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 + max_layer = getattr(config, "text_encoder_max_layer", 27) + if max_layer is not None and max_layer < 27: + raise ValueError( + f"Invalid configuration `text_encoder_max_layer={max_layer}`. " + f"FLUX.2-Klein requires extracting intermediate prompt embeddings from Qwen3 layers 9, 18, and 27, " + f"so `text_encoder_max_layer` must be >= 27." + ) + self.mesh = mesh + self.tokenizer = tokenizer + if self.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: + self.tokenizer = Qwen2TokenizerFast.from_pretrained(tokenizer_path, local_files_only=True) + except Exception: + self.tokenizer = Qwen2TokenizerFast.from_pretrained(tokenizer_path, subfolder="tokenizer", local_files_only=True) + + # JIT compilation cache + self._jitted_qwen3_forward = None + self._jitted_transformer_step = None + self._jitted_vae_encode = 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): + _, all_hidden_states = self.text_encoder.apply({"params": q_params}, input_ids=ids, attention_mask=mask) + 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) + prompt_embeds = jnp.transpose(out, (0, 2, 1, 3)).reshape((ids.shape[0], ids.shape[1], -1)) + context_spec = P(None, "context") if "context" in self.mesh.axis_names and self.mesh.shape["context"] > 1 else P() + prompt_embeds = jax.lax.with_sharding_constraint(prompt_embeds, jax.sharding.NamedSharding(self.mesh, context_spec)) + return prompt_embeds + + if isinstance(self.vae, nnx.Module): + v_graph, _, v_rest = nnx.split(self.vae, nnx.Param, ...) + + @jax.jit + def vae_encode(v_params, img): + merged = nnx.merge(v_graph, v_params, v_rest) + return merged.encode(img) + + @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)) + + merged = nnx.merge(v_graph, v_params, v_rest) + res = merged.decode(latents_unpacked) + return FlaxDecoderOutput(sample=res) + + else: + + @jax.jit + def vae_encode(v_params, img): + # FlaxAutoencoderKL expects (B, 3, H, W) + res = self.vae.apply({"params": v_params}, sample=img, method=self.vae.encode) + moments = res.latent_dist.mode() + return jnp.transpose(moments, (0, 3, 1, 2)) + + @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) + + if isinstance(self.transformer, nnx.Module): + g, nnx_state, r = nnx.split(self.transformer, nnx.Param, ...) + + @jax.jit + def transformer_step(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timestep, guidance): + nnx_merged = nnx.merge(g, t_params, r) + return nnx_merged( + hidden_states=latents, + encoder_hidden_states=prompt_embeds, + pooled_projections=vec, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + guidance=guidance, + return_dict=True, + ) + + @jax.jit(static_argnums=(9,)) + def fused_denoise_loop( + t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance, target_len=None + ): + sigmas_padded = jnp.concatenate([sigmas, jnp.array([0.0], dtype=sigmas.dtype)]) + nnx_merged = nnx.merge(g, t_params, r) + + 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 = nnx_merged( + 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, + return_dict=True, + ) + sigma = sigmas_padded[step_idx] + sigma_next = sigmas_padded[step_idx + 1] + dt = sigma_next - sigma + v = model_output.sample + if target_len is not None and cur_latents.shape[1] > target_len: + target_latents = cur_latents[:, :target_len, :] + v_target = v[:, :target_len, :] + next_target = target_latents + v_target * dt + prev_sample = jnp.concatenate([next_target, cur_latents[:, target_len:, :]], axis=1) + else: + prev_sample = cur_latents + v * dt + return prev_sample, None + + steps = jnp.arange(timesteps.shape[0]) + final_latents, _ = jax.lax.scan(scan_body, latents, steps) + return final_latents + + @jax.jit(static_argnums=(10, 11)) + def fused_kv_denoise_loop( + t_params, + target_latents, + ref_latents, + target_img_ids, + ref_img_ids, + prompt_embeds, + txt_ids, + vec, + timesteps, + sigmas, + guidance=None, + num_ref_tokens=0, + ): + sigmas_padded = jnp.concatenate([sigmas, jnp.array([0.0], dtype=sigmas.dtype)]) + nnx_merged = nnx.merge(g, t_params, r) + + step0_latents = jnp.concatenate([ref_latents, target_latents], axis=1) + step0_img_ids = jnp.concatenate([ref_img_ids, target_img_ids], axis=1) + t0_val = timesteps[0] + t0_vec = jnp.broadcast_to(t0_val / 1000.0, (target_latents.shape[0],)) + + out0, kv_cache = nnx_merged( + hidden_states=step0_latents, + img_ids=step0_img_ids, + encoder_hidden_states=prompt_embeds, + txt_ids=txt_ids, + pooled_projections=vec, + timestep=t0_vec, + guidance=guidance, + return_dict=True, + kv_cache_mode="extract", + num_ref_tokens=num_ref_tokens, + ) + dt0 = sigmas_padded[1] - sigmas_padded[0] + v0 = out0.sample + latents_step1 = target_latents + v0 * dt0 + + 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 = nnx_merged( + hidden_states=cur_latents, + img_ids=target_img_ids, + encoder_hidden_states=prompt_embeds, + txt_ids=txt_ids, + pooled_projections=vec, + timestep=t_vec, + guidance=guidance, + return_dict=True, + kv_cache=kv_cache, + kv_cache_mode="cached", + num_ref_tokens=num_ref_tokens, + ) + sigma = sigmas_padded[step_idx] + sigma_next = sigmas_padded[step_idx + 1] + dt = sigma_next - sigma + v = model_output.sample + next_latents = cur_latents + v * dt + return next_latents, None + + steps = jnp.arange(1, timesteps.shape[0]) + final_latents, _ = jax.lax.scan(scan_body, latents_step1, steps) + return final_latents + + else: + + @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=(9,)) + def fused_denoise_loop( + t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance, target_len=None + ): + sigmas_padded = jnp.concatenate([sigmas, jnp.array([0.0], dtype=sigmas.dtype)]) + + 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_padded[step_idx] + sigma_next = sigmas_padded[step_idx + 1] + dt = sigma_next - sigma + v = model_output.sample + if target_len is not None and cur_latents.shape[1] > target_len: + target_latents = cur_latents[:, :target_len, :] + v_target = v[:, :target_len, :] + next_target = target_latents + v_target * dt + prev_sample = jnp.concatenate([next_target, cur_latents[:, target_len:, :]], axis=1) + else: + prev_sample = cur_latents + v * dt + return prev_sample, None + + steps = jnp.arange(timesteps.shape[0]) + final_latents, _ = jax.lax.scan(scan_body, latents, steps) + return final_latents + + fused_kv_denoise_loop = fused_denoise_loop + + self._jitted_qwen3_forward = qwen3_forward + self._jitted_transformer_step = transformer_step + self._jitted_fused_denoise_loop = fused_denoise_loop + self._jitted_fused_kv_denoise_loop = fused_kv_denoise_loop + self._jitted_vae_encode = vae_encode + 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, + images=None, + image=None, + num_conditioning_images=0, + use_kv=None, + ): + """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 + + if images is None and image is not None: + images = image if isinstance(image, list) else [image] + + if images is not None and len(images) > 0: + num_conditioning_images = len(images) + + seq_len_img = (height // 16) * (width // 16) + total_img_len = (1 + num_conditioning_images) * seq_len_img + 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, total_img_len, 128), dtype=jnp.float32) + dummy_img_ids = jnp.zeros((batch_size, total_img_len, 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_target_latents = jnp.zeros((batch_size, seq_len_img, 128), 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()) + context_sharding = jax.sharding.NamedSharding(self.mesh, P(None, "context")) + + 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, context_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_target_latents = put_data_on_devices(dummy_target_latents, 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,), dtype=jnp.float32), replicated_sharding) + + use_kv = getattr(self._config, "use_kv", False) if use_kv is None else use_kv + if images is not None and len(images) > 0: + total_ref_tokens = sum( + (img.size[1] // 16) * (img.size[0] // 16) if hasattr(img, "size") else seq_len_img for img in images + ) + else: + total_ref_tokens = num_conditioning_images * seq_len_img + + dummy_ref_latents = ( + put_data_on_devices(jnp.zeros((batch_size, total_ref_tokens, 128), dtype=jnp.float32), data_sharding) + if total_ref_tokens > 0 + else None + ) + dummy_ref_img_ids = ( + put_data_on_devices(jnp.zeros((batch_size, total_ref_tokens, 4), dtype=jnp.int32), data_sharding) + if total_ref_tokens > 0 + else None + ) + dummy_target_img_ids = put_data_on_devices(jnp.zeros((batch_size, seq_len_img, 4), dtype=jnp.int32), data_sharding) + + def compile_transformer(): + t0 = time.perf_counter() + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + if use_kv and total_ref_tokens > 0 and self._jitted_fused_kv_denoise_loop is not None: + self._jitted_fused_kv_denoise_loop.lower( + params, + dummy_target_latents, + dummy_ref_latents, + dummy_target_img_ids, + dummy_ref_img_ids, + dummy_prompt_embeds, + dummy_txt_ids, + None, + dummy_timesteps, + dummy_sigmas, + None, + num_ref_tokens=total_ref_tokens, + ).compile() + max_logging.log(f" -> [AOT COMPILED] Fused Flux Transformer KV Denoise Scan in {time.perf_counter() - t0:.2f}s") + else: + self._jitted_fused_denoise_loop.lower( + params, + dummy_latents, + dummy_img_ids, + dummy_prompt_embeds, + dummy_txt_ids, + None, + dummy_timesteps, + dummy_sigmas, + None, + seq_len_img, + ).compile() + max_logging.log(f" -> [AOT COMPILED] Fused Flux Transformer Denoise Scan 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_target_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") + + def compile_vae_encode(): + t0 = time.perf_counter() + dummy_rgb = jnp.zeros((1, 3, height, width), dtype=jnp.float32) + with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules): + self._jitted_vae_encode.lower(vae_params, dummy_rgb).compile() + max_logging.log(f" -> [AOT COMPILED] VAE Encoder in {time.perf_counter() - t0:.2f}s") + + t_start = time.perf_counter() + with ThreadPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(compile_qwen3), + executor.submit(compile_transformer), + executor.submit(compile_vae), + ] + if num_conditioning_images > 0 or (images is not None and len(images) > 0): + futures.append(executor.submit(compile_vae_encode)) + 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, + images: Optional[List[Any]] = None, + image: Optional[Union[Any, List[Any]]] = None, + 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", + profile_target: Optional[str] = "all", + tensorboard_dir: Optional[str] = None, + use_kv: Optional[bool] = None, + ): + # 1. Setup JIT functions + self._setup_jit_functions() + + if images is None and image is not None: + images = image if isinstance(image, list) else [image] + + # 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) + elif C == 128: + latents_jax = jnp.transpose(jnp.reshape(latents_jax, (B, C, H * W)), (0, 2, 1)) + 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) + target_img_ids_val = prepare_latent_image_ids(batch_size, height // 16, width // 16) + t_pipeline_start = time.perf_counter() + trace = {} + + # 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, + ) + + 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}] " + tb_dir = tensorboard_dir or getattr(self._config, "tensorboard_dir", "/mnt/workspace/maxdiffusion/tensorboard_traces") + + # 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 0: Encode Reference Images (VAE) + # --------------------------------------------------------------------- + if images is not None and len(images) > 0: + do_prof_vae_enc = profile_target in ("all", "vae_encode") + if do_prof_vae_enc and not warmup: + trace_path = os.path.join(tb_dir, "profile_vae_encode") + os.makedirs(trace_path, exist_ok=True) + max_logging.log(f"🚀 [PROFILER] Starting XProf trace for VAE Reference Encode -> {trace_path}") + jax.profiler.start_trace(trace_path) + + t0_vae_enc_start = time.perf_counter() + trace["start_to_vae_encode"] = t0_vae_enc_start - t_pipeline_start + max_logging.log(f"{host_prefix} [PHASE 0] Encoding {len(images)} reference image(s) using JAX VAE encoder on TPU...") + norm_ref_latents = [] + packed_ref_latents = [] + bn_mean_arr = jnp.array(vae_bn_mean, dtype=jnp.float32) + bn_std_arr = jnp.array(vae_bn_std, dtype=jnp.float32) + + for img in images: + if isinstance(img, Image.Image): + img = img.convert("RGB").resize((width, height), Image.Resampling.BICUBIC) + arr = np.array(img, dtype=np.float32) / 127.5 - 1.0 + arr = np.transpose(arr, (2, 0, 1)) + img_tensor = jnp.expand_dims(jnp.array(arr), axis=0) + elif isinstance(img, np.ndarray): + if img.ndim == 3: + img = np.expand_dims(img, axis=0) + if img.shape[-1] == 3: + img = np.transpose(img, (0, 3, 1, 2)) + if np.issubdtype(img.dtype, np.integer): + img = img.astype(np.float32) / 127.5 - 1.0 + elif np.issubdtype(img.dtype, np.floating): + if img.max() > 1.0: + img = img / 127.5 - 1.0 + elif img.min() >= 0.0: + img = img * 2.0 - 1.0 + img_tensor = jnp.array(img, dtype=np.float32) + elif isinstance(img, jnp.ndarray): + if img.ndim == 3: + img = jnp.expand_dims(img, axis=0) + if img.shape[-1] == 3: + img = jnp.transpose(img, (0, 3, 1, 2)) + if jnp.issubdtype(img.dtype, jnp.integer): + img = img.astype(jnp.float32) / 127.5 - 1.0 + elif jnp.issubdtype(img.dtype, jnp.floating): + if img.max() > 1.0: + img = img / 127.5 - 1.0 + elif img.min() >= 0.0: + img = img * 2.0 - 1.0 + img_tensor = img + else: + raise ValueError(f"Unsupported image type: {type(img)}") + + raw_ref_latents = self._jitted_vae_encode(vae_params, img_tensor) + raw_ref_latents.block_until_ready() + patchified_ref = patchify_latents(raw_ref_latents) + normalized_ref = (patchified_ref - bn_mean_arr) / bn_std_arr + norm_ref_latents.append(normalized_ref) + + packed = jnp.transpose( + jnp.reshape(normalized_ref, (normalized_ref.shape[0], normalized_ref.shape[1], -1)), (0, 2, 1) + ) + if packed.shape[0] == 1 and batch_size > 1: + packed = jnp.repeat(packed, batch_size, axis=0) + packed_ref_latents.append(packed) + + ref_img_ids_val = prepare_multi_image_ids(norm_ref_latents, scale=10) + if ref_img_ids_val.shape[0] == 1 and batch_size > 1: + ref_img_ids_val = jnp.repeat(ref_img_ids_val, batch_size, axis=0) + ref_latents_jax = jnp.concatenate(packed_ref_latents, axis=1) + num_ref_tokens = ref_latents_jax.shape[1] + img_ids_val = jnp.concatenate([target_img_ids_val, ref_img_ids_val], axis=1) + latents_jax = jnp.concatenate([latents_jax] + packed_ref_latents, axis=1) + max_logging.log(f" [PIPELINE] Joint latents shape: {latents_jax.shape}, Joint img_ids shape: {img_ids_val.shape}") + + if do_prof_vae_enc and not warmup: + jax.profiler.stop_trace() + max_logging.log("✅ [PROFILER] Completed XProf trace for VAE Reference Encode") + + t0_vae_enc_end = time.perf_counter() + trace["vae_encode"] = t0_vae_enc_end - t0_vae_enc_start + trace["image_encoding"] = trace["vae_encode"] + max_logging.log(f" -> [TIMING] Reference Image Encoding (VAE): {trace['vae_encode']:.4f} seconds ⏱ī¸") + else: + img_ids_val = target_img_ids_val + ref_latents_jax = None + num_ref_tokens = 0 + trace["vae_encode"] = 0.0 + trace["image_encoding"] = 0.0 + + t0_qwen3_start = time.perf_counter() + if trace.get("vae_encode", 0.0) > 0: + trace["vae_encode_to_qwen3"] = t0_qwen3_start - t0_vae_enc_end + max_logging.log(f" -> [TIMING] VAE Encode to Qwen3 Overhead: {trace['vae_encode_to_qwen3']:.4f} seconds ⏱ī¸") + else: + 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 not prompts: + raise ValueError("Prompt must be provided to FlaxFlux2KleinProfilePipeline") + if 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) + + do_prof_qwen3 = profile_target in ("all", "qwen3") + if do_prof_qwen3 and not warmup: + trace_path = os.path.join(tb_dir, "profile_qwen3") + os.makedirs(trace_path, exist_ok=True) + max_logging.log(f"🚀 [PROFILER] Starting XProf trace for Qwen3 Text Encoder -> {trace_path}") + jax.profiler.start_trace(trace_path) + + with jax.named_scope("qwen3_text_encoder"): + prompt_embeds_jax = self._jitted_qwen3_forward(qwen3_params, prompt_ids, prompt_mask) + prompt_embeds_jax.block_until_ready() + + if do_prof_qwen3 and not warmup: + jax.profiler.stop_trace() + max_logging.log("✅ [PROFILER] Completed XProf trace for Qwen3 Text Encoder") + + 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 ⏱ī¸") + + proc_id = jax.process_index() + proc_cnt = jax.process_count() + host_prefix = f"[HOST {proc_id}/{proc_cnt}] " + + # 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) + context_spec = P(None, "context") if "context" in self.mesh.axis_names and self.mesh.shape["context"] > 1 else P() + context_sharding = jax.sharding.NamedSharding(self.mesh, context_spec) + prompt_embeds_jax = put_data_on_devices(prompt_embeds_jax, context_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 (Flux Transformer - Standalone Step JIT) + # --------------------------------------------------------------------- + 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) + + do_prof_denoise = profile_target in ("all", "denoise") + if do_prof_denoise and not warmup: + trace_path = os.path.join(tb_dir, "profile_denoise") + os.makedirs(trace_path, exist_ok=True) + max_logging.log(f"🚀 [PROFILER] Starting XProf trace for Flux Denoising Loop -> {trace_path}") + jax.profiler.start_trace(trace_path) + + use_kv = getattr(self._config, "use_kv", False) if use_kv is None else use_kv + if use_kv and len(packed_ref_latents) > 0: + ref_latents_device = put_data_on_devices(ref_latents_jax, data_sharding) + ref_img_ids_device = put_data_on_devices(ref_img_ids_val, data_sharding) + target_img_ids_device = put_data_on_devices(target_img_ids_val, data_sharding) + target_latents_device = put_data_on_devices(latents_jax[:, :seq_len_img, :], data_sharding) + + with jax.named_scope("fused_flux_kv_denoise_loop"): + latents_jax = self._jitted_fused_kv_denoise_loop( + params, + target_latents_device, + ref_latents_device, + target_img_ids_device, + ref_img_ids_device, + prompt_embeds_jax, + txt_ids_val, + vec_val, + timesteps_device, + sigmas_device, + guidance_vec_val, + num_ref_tokens, + ) + latents_jax.block_until_ready() + else: + 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, + seq_len_img, + ) + latents_jax.block_until_ready() + + if do_prof_denoise and not warmup: + jax.profiler.stop_trace() + max_logging.log("✅ [PROFILER] Completed XProf trace for Flux Denoising Loop") + + 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...") + + # Slice target latents from joint latents if reference images were present + if latents_jax.shape[1] > seq_len_img: + latents_jax = latents_jax[:, :seq_len_img, :] + + # 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 ⏱ī¸") + + do_prof_vae = profile_target in ("all", "vae") + if do_prof_vae and not warmup: + trace_path = os.path.join(tb_dir, "profile_vae") + os.makedirs(trace_path, exist_ok=True) + max_logging.log(f"🚀 [PROFILER] Starting XProf trace for VAE Decoder -> {trace_path}") + jax.profiler.start_trace(trace_path) + + 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() + + if do_prof_vae and not warmup: + jax.profiler.stop_trace() + max_logging.log("✅ [PROFILER] Completed XProf trace for VAE Decoder") + + 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/pipelines/flux/flux_pipeline.py b/src/maxdiffusion/pipelines/flux/flux_pipeline.py index 15b2c4f5e..69b20882f 100644 --- a/src/maxdiffusion/pipelines/flux/flux_pipeline.py +++ b/src/maxdiffusion/pipelines/flux/flux_pipeline.py @@ -18,7 +18,12 @@ import jax import jax.numpy as jnp import math -from transformers import (CLIPTokenizer, FlaxCLIPTextModel, FlaxT5EncoderModel, AutoTokenizer) +from transformers import CLIPTokenizer, AutoTokenizer +try: + from transformers import FlaxCLIPTextModel, FlaxT5EncoderModel +except ImportError: + FlaxCLIPTextModel = None + FlaxT5EncoderModel = None from einops import rearrange from jax.typing import DTypeLike from chex import Array diff --git a/src/maxdiffusion/pipelines/ltx_video/ltx_video_pipeline.py b/src/maxdiffusion/pipelines/ltx_video/ltx_video_pipeline.py index 4aa3baf10..64c3de6bc 100644 --- a/src/maxdiffusion/pipelines/ltx_video/ltx_video_pipeline.py +++ b/src/maxdiffusion/pipelines/ltx_video/ltx_video_pipeline.py @@ -17,8 +17,11 @@ from typing import Optional, List, Union, Tuple from einops import rearrange import torch.nn.functional as F -from maxdiffusion.models.ltx_video.autoencoders.vae_torchax import TorchaxCausalVideoAutoencoder -from transformers import (FlaxT5EncoderModel, AutoTokenizer) +from transformers import AutoTokenizer +try: + from transformers import FlaxT5EncoderModel +except ImportError: + FlaxT5EncoderModel = None from torchax import interop from torchax import default_env from dataclasses import dataclass diff --git a/src/maxdiffusion/pipelines/pipeline_flax_utils.py b/src/maxdiffusion/pipelines/pipeline_flax_utils.py index a2d4cc8db..d28920a1d 100644 --- a/src/maxdiffusion/pipelines/pipeline_flax_utils.py +++ b/src/maxdiffusion/pipelines/pipeline_flax_utils.py @@ -41,7 +41,10 @@ ) -from transformers import FlaxPreTrainedModel +try: + from transformers import FlaxPreTrainedModel +except ImportError: + FlaxPreTrainedModel = None INDEX_FILE = "diffusion_flax_model.bin" @@ -514,7 +517,7 @@ def load_module(name, value): quant=quant, ) params[name] = loaded_params - elif issubclass(class_obj, FlaxPreTrainedModel): + elif FlaxPreTrainedModel is not None and issubclass(class_obj, FlaxPreTrainedModel): if from_pt: # TODO(Suraj): Fix this in Transformers. We should be able to use `_do_init=False` here loaded_sub_model = load_method(loadable_folder, from_pt=from_pt) diff --git a/src/maxdiffusion/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/src/maxdiffusion/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py index 96e2f6d5a..60cee2489 100644 --- a/src/maxdiffusion/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +++ b/src/maxdiffusion/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py @@ -20,8 +20,11 @@ import jax import jax.numpy as jnp from flax.core.frozen_dict import FrozenDict -from packaging import version -from transformers import CLIPTokenizer, FlaxCLIPTextModel +from transformers import CLIPTokenizer +try: + from transformers import FlaxCLIPTextModel +except ImportError: + FlaxCLIPTextModel = None from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel from ...schedulers import ( diff --git a/src/maxdiffusion/pipelines/stable_diffusion/safety_checker_flax.py b/src/maxdiffusion/pipelines/stable_diffusion/safety_checker_flax.py index ccf0ba1ca..3716a3343 100644 --- a/src/maxdiffusion/pipelines/stable_diffusion/safety_checker_flax.py +++ b/src/maxdiffusion/pipelines/stable_diffusion/safety_checker_flax.py @@ -17,8 +17,11 @@ import jax import jax.numpy as jnp from flax import linen as nn -from flax.core.frozen_dict import FrozenDict -from transformers import CLIPConfig, FlaxPreTrainedModel +try: + from transformers import CLIPConfig, FlaxPreTrainedModel +except ImportError: + CLIPConfig = None + FlaxPreTrainedModel = object try: from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule diff --git a/src/maxdiffusion/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py b/src/maxdiffusion/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py index e3d96d6e9..3d0575dae 100644 --- a/src/maxdiffusion/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py +++ b/src/maxdiffusion/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py @@ -17,8 +17,11 @@ import jax import jax.numpy as jnp -from flax.core.frozen_dict import FrozenDict -from transformers import CLIPTokenizer, FlaxCLIPTextModel +from transformers import CLIPTokenizer +try: + from transformers import FlaxCLIPTextModel +except ImportError: + FlaxCLIPTextModel = None from maxdiffusion.utils import logging diff --git a/src/maxdiffusion/test_output/edited_ballroom_9b_b0.png b/src/maxdiffusion/test_output/edited_ballroom_9b_b0.png new file mode 100644 index 000000000..0a205d548 Binary files /dev/null and b/src/maxdiffusion/test_output/edited_ballroom_9b_b0.png differ diff --git a/src/maxdiffusion/test_output/edited_ballroom_9b_b1.png b/src/maxdiffusion/test_output/edited_ballroom_9b_b1.png new file mode 100644 index 000000000..fc55b8462 Binary files /dev/null and b/src/maxdiffusion/test_output/edited_ballroom_9b_b1.png differ diff --git a/src/maxdiffusion/test_output/edited_ballroom_9b_b2.png b/src/maxdiffusion/test_output/edited_ballroom_9b_b2.png new file mode 100644 index 000000000..49e50fd30 Binary files /dev/null and b/src/maxdiffusion/test_output/edited_ballroom_9b_b2.png differ diff --git a/src/maxdiffusion/test_output/edited_ballroom_9b_b3.png b/src/maxdiffusion/test_output/edited_ballroom_9b_b3.png new file mode 100644 index 000000000..7864d9ea4 Binary files /dev/null and b/src/maxdiffusion/test_output/edited_ballroom_9b_b3.png differ diff --git a/src/maxdiffusion/test_output/flux2klein_warmup_b0.png b/src/maxdiffusion/test_output/flux2klein_warmup_b0.png new file mode 100644 index 000000000..0a205d548 Binary files /dev/null and b/src/maxdiffusion/test_output/flux2klein_warmup_b0.png differ diff --git a/src/maxdiffusion/test_output/flux2klein_warmup_b1.png b/src/maxdiffusion/test_output/flux2klein_warmup_b1.png new file mode 100644 index 000000000..fc55b8462 Binary files /dev/null and b/src/maxdiffusion/test_output/flux2klein_warmup_b1.png differ diff --git a/src/maxdiffusion/test_output/flux2klein_warmup_b2.png b/src/maxdiffusion/test_output/flux2klein_warmup_b2.png new file mode 100644 index 000000000..49e50fd30 Binary files /dev/null and b/src/maxdiffusion/test_output/flux2klein_warmup_b2.png differ diff --git a/src/maxdiffusion/test_output/flux2klein_warmup_b3.png b/src/maxdiffusion/test_output/flux2klein_warmup_b3.png new file mode 100644 index 000000000..7864d9ea4 Binary files /dev/null and b/src/maxdiffusion/test_output/flux2klein_warmup_b3.png differ diff --git a/src/maxdiffusion/test_output/test_sdxl.png b/src/maxdiffusion/test_output/test_sdxl.png new file mode 100644 index 000000000..311beb2be Binary files /dev/null and b/src/maxdiffusion/test_output/test_sdxl.png differ diff --git a/src/maxdiffusion/tests/edit_flux2klein_e2e_test.py b/src/maxdiffusion/tests/edit_flux2klein_e2e_test.py new file mode 100644 index 000000000..0987fb54d --- /dev/null +++ b/src/maxdiffusion/tests/edit_flux2klein_e2e_test.py @@ -0,0 +1,349 @@ +""" +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 gc +import unittest +import pytest +import numpy as np +from PIL import Image +from skimage.metrics import structural_similarity as ssim +import torch + +import jax +import jax.numpy as jnp +from flax import nnx +from jax.sharding import Mesh +from transformers import AutoConfig, Qwen2TokenizerFast + +from maxdiffusion import pyconfig +from maxdiffusion.max_utils import create_device_mesh +from maxdiffusion.models.flux.transformers.transformer_flux_flax import NNXFlux2KleinTransformer2DModel +from maxdiffusion.models.flux.vae.autoencoder_kl_flux2_nnx import ( + NNXAutoencoderKLFlux2, + load_and_convert_flux2klein_nnx_vae_weights, +) +from maxdiffusion.models.flux.util import load_and_convert_flux_klein_nnx_weights +from maxdiffusion.models.qwen3_flax import FlaxQwen3Model, FlaxQwen3Config +from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights +from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler +from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline + +IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +PROMPT = "a vibrant artistic painting combining the dog, car, mountain, and fruit bowl in surreal neon lighting" + + +class TestFlux2KleinImageEditE2EParity(unittest.TestCase): + """End-to-End Parity Test between PyTorch Diffusers CPU and MaxDiffusion TPU.""" + + def setUp(self): + jax.config.update("jax_default_matmul_precision", "highest") + jax.config.update("jax_use_shardy_partitioner", True) + + if "FLUX2_KLEIN_4B_MODEL_PATH" in os.environ: + self.model_dir = os.environ["FLUX2_KLEIN_4B_MODEL_PATH"] + else: + hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")) + candidates = [ + os.path.join(hf_home, "hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots"), + os.path.join(hf_home, "hub/models--black-forest-labs--FLUX.2-klein-4b/snapshots"), + "/mnt/hyperdisk_weights/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots", + "/mnt/data/models/flux2klein-4b", + ] + self.model_dir = None + for c in candidates: + if os.path.exists(c): + if "snapshots" in c: + snaps = os.listdir(c) + if snaps: + self.model_dir = os.path.join(c, snaps[0]) + else: + self.model_dir = c + if self.model_dir: + self.transformer_path = os.path.join(self.model_dir, "transformer") + self.vae_path = os.path.join(self.model_dir, "vae", "diffusion_pytorch_model.safetensors") + self.text_encoder_path = os.path.join(self.model_dir, "text_encoder") + self.tokenizer_path = os.path.join(self.model_dir, "tokenizer") + if os.path.exists(self.transformer_path) and os.path.exists(self.vae_path): + break + if self.model_dir is None: + self.model_dir = "black-forest-labs/FLUX.2-klein-4B" + + if hasattr(self, "model_dir") and self.model_dir and not hasattr(self, "transformer_path"): + self.transformer_path = os.path.join(self.model_dir, "transformer") + self.vae_path = os.path.join(self.model_dir, "vae", "diffusion_pytorch_model.safetensors") + self.text_encoder_path = os.path.join(self.model_dir, "text_encoder") + self.tokenizer_path = os.path.join(self.model_dir, "tokenizer") + + self.output_dir = "/tmp/e2e_parity" + os.makedirs(self.output_dir, exist_ok=True) + + # Resolve reference images + ref_dir = os.path.join(THIS_DIR, "images", "flux2klein") + self.ref_images = [] + if os.path.exists(ref_dir): + for i in range(4): + p = os.path.join(ref_dir, f"ref_image_{i}.png") + if os.path.exists(p): + self.ref_images.append(Image.open(p).convert("RGB")) + + if len(self.ref_images) < 4: + # Generate synthetic test reference images if not present + colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)] + for i, c in enumerate(colors): + arr = np.full((512, 512, 3), c, dtype=np.uint8) + self.ref_images.append(Image.fromarray(arr)) + + @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run on Github Actions (requires TPU and full weights)") + def test_e2e_image_edit_parity_vs_diffusers(self): + """Generates an image edit on PyTorch Diffusers CPU and MaxDiffusion TPU and asserts SSIM >= 0.75.""" + from diffusers import Flux2KleinPipeline as DiffusersFlux2KleinPipeline + + print("\n" + "=" * 80) + print("🚀 [STEP 1/3] Running Reference PyTorch Diffusers CPU Pipeline...") + print("=" * 80) + + diffusers_pipe = DiffusersFlux2KleinPipeline.from_pretrained(self.model_dir, torch_dtype=torch.bfloat16) + diffusers_pipe.to("cpu") + + # Generate initial noise latents deterministically on CPU (4D tensor for Diffusers prepare_latents) + gen = torch.Generator(device="cpu").manual_seed(42) + raw_latents_pt = torch.randn( + (1, 128, 512 // 16, 512 // 16), + generator=gen, + dtype=torch.bfloat16, + device="cpu", + ) + + with torch.no_grad(): + diffusers_out = diffusers_pipe( + prompt=PROMPT, + image=self.ref_images, + height=512, + width=512, + num_inference_steps=4, + latents=raw_latents_pt, + guidance_scale=1.0, + ) + + diffusers_image = diffusers_out.images[0] + diffusers_img_path = os.path.join(self.output_dir, "diffusers_cpu_output.png") + diffusers_image.save(diffusers_img_path) + print(f" -> Saved PyTorch Diffusers output to: {diffusers_img_path}") + + # Free PyTorch pipeline memory before TPU run + del diffusers_pipe + gc.collect() + + print("\n" + "=" * 80) + print("🚀 [STEP 2/3] Running MaxDiffusion Unified FlaxFlux2KleinPipeline on TPU...") + print("=" * 80) + + # 1. Device mesh setup + active_devices = jax.devices() + active_device_count = len(active_devices) + + pyconfig._config = None + pyconfig.config = None + config_path = os.path.join(THIS_DIR, "..", "configs", "base_flux2klein.yml") + args = [ + None, + config_path, + "run_name=e2e_parity_test", + f"output_dir={self.output_dir}", + f"per_device_batch_size={1.0 / active_device_count}", + "height=512", + "width=512", + "seed=42", + "weights_dtype=bfloat16", + "activations_dtype=bfloat16", + "precision=DEFAULT", + "text_encoder_attention=dot_product", + ] + pyconfig.initialize(args) + config = pyconfig.config + + if active_device_count > 1: + pyconfig._config.keys["ici_tensor_parallelism"] = active_device_count + pyconfig._config.keys["ici_data_parallelism"] = 1 + pyconfig._config.keys["ici_fsdp_parallelism"] = 1 + pyconfig._config.keys["ici_context_parallelism"] = 1 + + devices_array = create_device_mesh(config, devices=active_devices) + mesh = Mesh(devices_array, config.mesh_axes) + + # 2. Load NNX Transformer + print(" -> Loading NNX Transformer weights...") + rngs = nnx.Rngs(0) + transformer = NNXFlux2KleinTransformer2DModel( + rngs=rngs, + patch_size=1, + in_channels=128, + num_layers=5, + num_single_layers=20, + attention_head_dim=128, + num_attention_heads=24, + joint_attention_dim=7680, + pooled_projection_dim=None, + guidance_embeds=False, + axes_dim=(32, 32, 32, 32), + scale_shift_order="scale_shift", + dtype=jnp.bfloat16, + weights_dtype=jnp.bfloat16, + ) + t_state = load_and_convert_flux_klein_nnx_weights( + self.transformer_path, + nnx.state(transformer, nnx.Param), + num_double_layers=5, + num_single_layers=20, + dtype=jnp.bfloat16, + ) + nnx.update(transformer, t_state) + + # 3. Load NNX VAE + print(" -> Loading NNX VAE weights...") + nnx_vae = NNXAutoencoderKLFlux2(dtype=jnp.bfloat16, param_dtype=jnp.bfloat16) + bn_mean, bn_std = load_and_convert_flux2klein_nnx_vae_weights(self.vae_path, nnx_vae, dtype=jnp.bfloat16) + + # 4. Load Qwen3 + print(" -> Loading Qwen3 weights...") + pt_config = AutoConfig.from_pretrained(self.text_encoder_path) + 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=getattr(pt_config, "num_key_value_heads", pt_config.num_attention_heads), + max_position_embeddings=getattr(pt_config, "max_position_embeddings", 32768), + rms_norm_eps=getattr(pt_config, "rms_norm_eps", 1e-6), + rope_theta=getattr(pt_config, "rope_theta", getattr(pt_config, "rope_base", 1000000.0)), + dtype=jnp.bfloat16, + max_layer_to_run=27, + ) + text_encoder = FlaxQwen3Model(config=qwen3_config) + abstract_q_vars = text_encoder.init( + jax.random.PRNGKey(0), jnp.zeros((1, 512), dtype=jnp.int32), jnp.zeros((1, 512), dtype=jnp.int32) + ) + q_params = load_and_convert_qwen3_weights(self.text_encoder_path, abstract_q_vars["params"], qwen3_config) + + tokenizer = Qwen2TokenizerFast.from_pretrained(self.tokenizer_path) + 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", + ) + + # 5. Place parameters on TPU HBM + t_params = nnx.state(transformer, nnx.Param) + v_params = nnx.state(nnx_vae, nnx.Param) + + t_params = jax.device_put(t_params) + v_params = jax.device_put(v_params) + q_params = jax.device_put(q_params) + + # 6. Instantiate Unified FlaxFlux2KleinPipeline + pipeline = FlaxFlux2KleinPipeline( + transformer=transformer, + vae=nnx_vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + config=config, + mesh=mesh, + ) + + # 7. AOT Compile async + pipeline.compile_aot_async( + params=t_params, + vae_params=v_params, + qwen3_params=q_params, + vae_bn_mean=bn_mean, + vae_bn_std=bn_std, + batch_size=1, + height=512, + width=512, + images=self.ref_images, + ) + + # Convert PyTorch initial noise latents to JAX array (shape: 1, 32, 64, 64) + initial_latents_jax = jnp.array(raw_latents_pt.detach().float().cpu().numpy()) + + # 8. Run pipeline + print(f" -> Running FlaxFlux2KleinPipeline with {len(self.ref_images)} reference images on TPU...") + pipeline( + prompt=PROMPT, + params=t_params, + vae_params=v_params, + qwen3_params=q_params, + vae_bn_mean=bn_mean, + vae_bn_std=bn_std, + transformer_shardings=None, + vae_shardings=None, + qwen3_shardings=None, + height=512, + width=512, + num_inference_steps=4, + batch_size=1, + images=self.ref_images, + use_latents=True, + latents=initial_latents_jax, + output_dir=self.output_dir, + output_name="maxdiffusion_tpu_output.png", + ) + + maxdiff_img_path = os.path.join(self.output_dir, "maxdiffusion_tpu_output.png") + self.assertTrue(os.path.exists(maxdiff_img_path), "MaxDiffusion output image was not saved!") + maxdiff_image = Image.open(maxdiff_img_path).convert("RGB") + + print("\n" + "=" * 80) + print("📊 [STEP 3/3] Evaluating End-to-End Parity (SSIM & PSNR)...") + print("=" * 80) + + diffusers_arr = np.array(diffusers_image).astype(np.uint8) + maxdiff_arr = np.array(maxdiff_image).astype(np.uint8) + + self.assertEqual(diffusers_arr.shape, maxdiff_arr.shape) + + ssim_val = ssim(diffusers_arr, maxdiff_arr, channel_axis=-1, data_range=255) + mse = np.mean((diffusers_arr.astype(np.float64) - maxdiff_arr.astype(np.float64)) ** 2) + psnr_val = 10.0 * np.log10(255.0**2 / (mse + 1e-10)) + + print(f" -> SSIM (Diffusers CPU vs MaxDiffusion TPU): {ssim_val:.6f}") + print(f" -> PSNR (Diffusers CPU vs MaxDiffusion TPU): {psnr_val:.2f} dB") + print(f" -> MSE: {mse:.4f}") + + # Create side-by-side comparison image + side_by_side = Image.new("RGB", (1024, 512)) + side_by_side.paste(diffusers_image, (0, 0)) + side_by_side.paste(maxdiff_image, (512, 0)) + comparison_path = os.path.join(self.output_dir, "e2e_parity_diffusers_vs_maxdiffusion.png") + side_by_side.save(comparison_path) + print(f" -> Saved side-by-side comparison to: {comparison_path}") + + self.assertGreaterEqual(ssim_val, 0.75, f"SSIM score {ssim_val:.4f} is below target threshold 0.75!") + print("🎉 END-TO-END PARITY TEST PASSED! MaxDiffusion matches Diffusers reference!") + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py b/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py index c61db29f7..b0c4f3f2f 100644 --- a/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py +++ b/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py @@ -37,11 +37,11 @@ class GenerateFlux2KleinSmokeTest(unittest.TestCase): @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)") def test_flux2klein_4b_smoke(self): """End-to-end smoke test for Flux.2-klein-4B image generation at 1024x1024.""" - ref_path = os.path.join(THIS_DIR, "images", "ref_flux2klein_4b.png") + ref_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b.png") self.assertTrue(os.path.exists(ref_path), f"Reference image not found: {ref_path}") base_image = np.array(Image.open(ref_path)).astype(np.uint8) - output_dir = "/mnt/data/smoke_test_4b" if os.path.exists("/mnt/data") else "/tmp/smoke_test_4b" + output_dir = "/tmp/smoke_test_4b" os.makedirs(output_dir, exist_ok=True) out_path = os.path.join(output_dir, "flux2klein_generated_image.png") if os.path.exists(out_path): @@ -77,16 +77,16 @@ def test_flux2klein_4b_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 4B] SSIM Score: {ssim_compare:.6f}") - self.assertGreaterEqual(ssim_compare, 0.75) + self.assertGreaterEqual(ssim_compare, 0.8) @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)") def test_flux2klein_9b_smoke(self): """End-to-end smoke test for Flux.2-klein-9B image generation at 1024x1024.""" - ref_path = os.path.join(THIS_DIR, "images", "ref_flux2klein_9b.png") + ref_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_9b.png") self.assertTrue(os.path.exists(ref_path), f"Reference image not found: {ref_path}") base_image = np.array(Image.open(ref_path)).astype(np.uint8) - output_dir = "/mnt/data/smoke_test_9b" if os.path.exists("/mnt/data") else "/tmp/smoke_test_9b" + output_dir = "/tmp/smoke_test_9b" os.makedirs(output_dir, exist_ok=True) out_path = os.path.join(output_dir, "flux2klein_generated_image.png") if os.path.exists(out_path): @@ -124,6 +124,104 @@ def test_flux2klein_9b_smoke(self): print(f"\n[SMOKE TEST 9B] SSIM Score: {ssim_compare:.6f}") self.assertGreaterEqual(ssim_compare, 0.8) + @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)") + def test_flux2klein_4b_image_edit_smoke(self): + """End-to-end smoke test for Flux.2-klein-4B image editing at 512x512.""" + ref_gold_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b_image_edit.png") + self.assertTrue(os.path.exists(ref_gold_path), f"Golden reference image not found: {ref_gold_path}") + base_image = np.array(Image.open(ref_gold_path)).astype(np.uint8) + + input_img_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b.png") + self.assertTrue(os.path.exists(input_img_path), f"Input reference image not found: {input_img_path}") + + output_dir = "/tmp/smoke_test_image_edit_4b" + os.makedirs(output_dir, exist_ok=True) + out_path = os.path.join(output_dir, "flux2klein_generated_image.png") + if os.path.exists(out_path): + os.remove(out_path) + + pyconfig._config = None + pyconfig.config = None + args = [ + None, + os.path.join(THIS_DIR, "..", "configs", "base_flux2klein.yml"), + "run_name=smoke_test_image_edit_4b", + f"output_dir={output_dir}", + "jax_cache_dir=/tmp/cache_dir", + f"image_paths=['{input_img_path}']", + "prompt=change the lighting to evening", + "height=512", + "width=512", + f"per_device_batch_size={1.0 / jax.device_count()}", + "seed=42", + "weights_dtype=bfloat16", + "activations_dtype=bfloat16", + "precision=DEFAULT", + "num_reps=5", + "text_encoder_attention=dot_product", + ] + + generate_flux2klein.main(args) + + rep_out_path = os.path.join(output_dir, "rep_1_flux2klein_generated_image.png") + final_out_path = rep_out_path if os.path.exists(rep_out_path) else out_path + self.assertTrue(os.path.exists(final_out_path), "Smoke test 4B image edit failed to produce output image!") + test_image = np.array(Image.open(final_out_path)).astype(np.uint8) + + 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 4B IMAGE EDIT] SSIM Score: {ssim_compare:.6f}") + self.assertGreaterEqual(ssim_compare, 0.8) + + @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)") + def test_flux2klein_9b_image_edit_smoke(self): + """End-to-end smoke test for Flux.2-klein-9B image editing at 512x512.""" + ref_gold_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_9b_image_edit.png") + self.assertTrue(os.path.exists(ref_gold_path), f"Golden reference image not found: {ref_gold_path}") + base_image = np.array(Image.open(ref_gold_path)).astype(np.uint8) + + input_img_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b.png") + self.assertTrue(os.path.exists(input_img_path), f"Input reference image not found: {input_img_path}") + + output_dir = "/tmp/smoke_test_image_edit_9b" + os.makedirs(output_dir, exist_ok=True) + out_path = os.path.join(output_dir, "flux2klein_generated_image.png") + if os.path.exists(out_path): + os.remove(out_path) + + pyconfig._config = None + pyconfig.config = None + args = [ + None, + os.path.join(THIS_DIR, "..", "configs", "base_flux2klein_9B.yml"), + "run_name=smoke_test_image_edit_9b", + f"output_dir={output_dir}", + "jax_cache_dir=/tmp/cache_dir", + f"image_paths=['{input_img_path}']", + "prompt=change the lighting to evening", + "height=512", + "width=512", + f"per_device_batch_size={1.0 / jax.device_count()}", + "seed=42", + "weights_dtype=bfloat16", + "activations_dtype=bfloat16", + "precision=DEFAULT", + "num_reps=5", + "text_encoder_attention=dot_product", + ] + + generate_flux2klein.main(args) + + rep_out_path = os.path.join(output_dir, "rep_1_flux2klein_generated_image.png") + final_out_path = rep_out_path if os.path.exists(rep_out_path) else out_path + self.assertTrue(os.path.exists(final_out_path), "Smoke test 9B image edit failed to produce output image!") + test_image = np.array(Image.open(final_out_path)).astype(np.uint8) + + 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 IMAGE EDIT] SSIM Score: {ssim_compare:.6f}") + self.assertGreaterEqual(ssim_compare, 0.8) + if __name__ == "__main__": unittest.main() diff --git a/src/maxdiffusion/tests/images/ref_flux2klein_4b.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b.png similarity index 100% rename from src/maxdiffusion/tests/images/ref_flux2klein_4b.png rename to src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b.png diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b_image_edit.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b_image_edit.png new file mode 100644 index 000000000..f24599b58 Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b_image_edit.png differ diff --git a/src/maxdiffusion/tests/images/ref_flux2klein_9b.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b.png similarity index 100% rename from src/maxdiffusion/tests/images/ref_flux2klein_9b.png rename to src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b.png diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_image_edit.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_image_edit.png new file mode 100644 index 000000000..7d938fb17 Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_image_edit.png differ diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_image_0.png b/src/maxdiffusion/tests/images/flux2klein/ref_image_0.png new file mode 100644 index 000000000..476ba5984 Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_image_0.png differ diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_image_1.png b/src/maxdiffusion/tests/images/flux2klein/ref_image_1.png new file mode 100644 index 000000000..d14acbe2a Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_image_1.png differ diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_image_2.png b/src/maxdiffusion/tests/images/flux2klein/ref_image_2.png new file mode 100644 index 000000000..cb379a6e8 Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_image_2.png differ diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_image_3.png b/src/maxdiffusion/tests/images/flux2klein/ref_image_3.png new file mode 100644 index 000000000..bffdbe23c Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_image_3.png differ diff --git a/warmup_prxpixel_cfg_1.png b/warmup_prxpixel_cfg_1.png new file mode 100644 index 000000000..8edfbd673 Binary files /dev/null and b/warmup_prxpixel_cfg_1.png differ diff --git a/warmup_prxpixel_cfg_2.png b/warmup_prxpixel_cfg_2.png new file mode 100644 index 000000000..110fcdf37 Binary files /dev/null and b/warmup_prxpixel_cfg_2.png differ diff --git a/warmup_prxpixel_cfg_2_no_np.png b/warmup_prxpixel_cfg_2_no_np.png new file mode 100644 index 000000000..d4905a3c1 Binary files /dev/null and b/warmup_prxpixel_cfg_2_no_np.png differ diff --git a/warmup_prxpixel_cfg_4.5.png b/warmup_prxpixel_cfg_4.5.png new file mode 100644 index 000000000..7aa75dfa8 Binary files /dev/null and b/warmup_prxpixel_cfg_4.5.png differ