feat(flux2klein): optimize weight loading speed and fix pipeline static jit dimensions

Author: amepas

src/maxdiffusion/configs/base_flux2klein_9B.yml

@@ -27,7 +27,7 @@ gcs_metrics: False
 save_config_to_gcs: False
 log_period: 100
 
-pretrained_model_name_or_path: 'black-forest-labs/FLUX.2-klein-9B'
+pretrained_model_name_or_path: '/mnt/hyperdisk_weights/hub/flux2klein-9b'
 clip_model_name_or_path: 'ariG23498/clip-vit-large-patch14-text-flax'
 t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax'
 
@@ -43,7 +43,7 @@ max_shift: 1.15
 
 
 unet_checkpoint: ''
-revision: 'refs/pr/95'
+revision: ''
 # This will convert the weights to this dtype.
 # When running inference on TPUv5e, use weights_dtype: 'bfloat16'
 weights_dtype: 'bfloat16'

src/maxdiffusion/configs/base_wan_14b.yml

@@ -274,10 +274,6 @@ load_tfrecord_cached: True
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''
-# Directory for per-shape AOT serialized executables ('' = disabled).
-aot_cache_dir: ''
-# Directory for memoized torch->flax converted weights ('' = disabled).
-converted_weights_dir: ''
 hf_data_dir: ''
 hf_train_files: ''
 hf_access_token: ''

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -227,10 +227,6 @@ load_tfrecord_cached: True
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''
-# Directory for per-shape AOT serialized executables ('' = disabled).
-aot_cache_dir: ''
-# Directory for memoized torch->flax converted weights ('' = disabled).
-converted_weights_dir: ''
 hf_data_dir: ''
 hf_train_files: ''
 hf_access_token: ''

src/maxdiffusion/configs/base_wan_27b.yml

@@ -248,10 +248,6 @@ load_tfrecord_cached: True
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''
-# Directory for per-shape AOT serialized executables ('' = disabled).
-aot_cache_dir: ''
-# Directory for memoized torch->flax converted weights ('' = disabled).
-converted_weights_dir: ''
 hf_data_dir: ''
 hf_train_files: ''
 hf_access_token: ''

src/maxdiffusion/configs/base_wan_animate.yml

@@ -237,10 +237,6 @@ load_tfrecord_cached: True
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: '.jax_cache'
-# Directory for per-shape AOT serialized executables ('' = disabled).
-aot_cache_dir: ''
-# Directory for memoized torch->flax converted weights ('' = disabled).
-converted_weights_dir: ''
 hf_data_dir: ''
 hf_train_files: ''
 hf_access_token: ''

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -240,10 +240,6 @@ load_tfrecord_cached: True
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''
-# Directory for per-shape AOT serialized executables ('' = disabled).
-aot_cache_dir: ''
-# Directory for memoized torch->flax converted weights ('' = disabled).
-converted_weights_dir: ''
 hf_data_dir: ''
 hf_train_files: ''
 hf_access_token: ''

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -241,10 +241,6 @@ load_tfrecord_cached: True
 train_data_dir: ''
 dataset_config_name: ''
 jax_cache_dir: ''
-# Directory for per-shape AOT serialized executables ('' = disabled).
-aot_cache_dir: ''
-# Directory for memoized torch->flax converted weights ('' = disabled).
-converted_weights_dir: ''
 hf_data_dir: ''
 hf_train_files: ''
 hf_access_token: ''

src/maxdiffusion/generate_flux2klein.py

@@ -62,7 +62,7 @@ def partition_prompts(prompt_str: str, batch_size: int) -> List[str]:
     max_logging.log(
         f"⚠️ Warning: Found {num_prompts} prompts, but batch_size is {batch_size}. Truncating to the first {batch_size}."
     )
-    return raw_prompts[:batch_size]
+
 
 
 def encode_prompt(prompt: str, snapshot_dir: str = None, repo_id: str = "black-forest-labs/FLUX.2-klein-4B"):
@@ -78,8 +78,18 @@ def encode_prompt(prompt: str, snapshot_dir: str = None, repo_id: str = "black-f
 
   text_encoder_path = os.path.join(snapshot_dir, "text_encoder")
   tokenizer_path = os.path.join(snapshot_dir, "tokenizer")
-  if not os.path.exists(tokenizer_path):
-    tokenizer_path = text_encoder_path
+
+  if not os.path.exists(os.path.join(text_encoder_path, "config.json")) or not os.path.exists(tokenizer_path):
+    try:
+      fb_dir = snapshot_download(repo_id=repo_id, local_files_only=True)
+      if not os.path.exists(os.path.join(text_encoder_path, "config.json")):
+        text_encoder_path = os.path.join(fb_dir, "text_encoder")
+      if not os.path.exists(tokenizer_path):
+        tokenizer_path = os.path.join(fb_dir, "tokenizer") if os.path.exists(os.path.join(fb_dir, "tokenizer")) else os.path.join(fb_dir, "text_encoder")
+    except Exception as e:
+      if not os.path.exists(tokenizer_path):
+        tokenizer_path = text_encoder_path
+
   tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
   text_encoder = AutoModelForCausalLM.from_pretrained(text_encoder_path, torch_dtype=torch.float32)
   text_encoder.eval()
@@ -183,8 +193,13 @@ def main(argv):
   else:
     from huggingface_hub import snapshot_download
 
-    max_logging.log(f"Resolving snapshot directory for model '{repo_id}' from HF Hub...")
-    snapshot_dir = snapshot_download(repo_id=repo_id)
+    rev = getattr(config, "revision", None)
+    if not rev or rev == "refs/pr/95":
+      rev = "main"
+    try:
+      snapshot_dir = snapshot_download(repo_id=repo_id, revision=rev, local_files_only=True)
+    except Exception:
+      snapshot_dir = snapshot_download(repo_id=repo_id, local_files_only=True)
 
   max_logging.log(f"Host {jax.process_index()} using HF snapshot directory: {snapshot_dir}")
   safetensors_path = os.path.join(snapshot_dir, "transformer")
@@ -194,8 +209,13 @@ def main(argv):
   # 4. Load Qwen3 Config & Setup model layout
   from transformers import AutoConfig
 
-  max_logging.log(f"Loading Qwen3 config from text_encoder path: {text_encoder_path}...")
-  pt_config = AutoConfig.from_pretrained(text_encoder_path, local_files_only=True)
+  try:
+    pt_config = AutoConfig.from_pretrained(text_encoder_path, local_files_only=True)
+  except Exception as e:
+    depth_val = getattr(config, "depth", 24)
+    hf_repo = "black-forest-labs/FLUX.2-klein-9B" if depth_val in (24, -1) else "black-forest-labs/FLUX.2-klein-4B"
+    max_logging.log(f"ℹ️ Config not found in {text_encoder_path}. Resolving from HF cache: {hf_repo}")
+    pt_config = AutoConfig.from_pretrained(hf_repo, subfolder="text_encoder", local_files_only=True)
 
   qwen3_config = FlaxQwen3Config(
       vocab_size=pt_config.vocab_size,
@@ -216,9 +236,25 @@ def main(argv):
 
   transformer_config_json = os.path.join(safetensors_path, "config.json")
   transformer_pt_cfg = {}
+  loaded_cfg = False
   if os.path.exists(transformer_config_json):
-    with open(transformer_config_json, "r") as f:
-      transformer_pt_cfg = json.load(f)
+    try:
+      with open(transformer_config_json, "r") as f:
+        transformer_pt_cfg = json.load(f)
+        loaded_cfg = True
+    except Exception as e:
+      max_logging.log(f"ℹ️ Could not parse {transformer_config_json}: {e}. Falling back to HF cache...")
+
+  if not loaded_cfg:
+    depth_val = getattr(config, "depth", 24)
+    hf_repo = "black-forest-labs/FLUX.2-klein-9B" if depth_val in (24, -1) else "black-forest-labs/FLUX.2-klein-4B"
+    try:
+      from huggingface_hub import hf_hub_download
+      cfg_file = hf_hub_download(repo_id=hf_repo, filename="transformer/config.json", local_files_only=True)
+      with open(cfg_file, "r") as f:
+        transformer_pt_cfg = json.load(f)
+    except Exception as e:
+      max_logging.log(f"⚠️ Warning resolving transformer config fallback: {e}")
 
   num_double_layers = getattr(config, "num_double_layers", -1)
   if num_double_layers is None or num_double_layers <= 0:
@@ -481,9 +517,22 @@ def unbox_fn(x):
       max_logging.log(f" -> Custom latents shape: {latents_to_use.shape} | sum: {latents_to_use.sum():.6f}")
 
     max_logging.log("\n" + "=" * 80)
-    max_logging.log("🚀 Running initial dry run (Warmup Pass) to compile XLA graphs...")
+    max_logging.log("🚀 Pre-compiling XLA graphs concurrently (AOT Compilation)...")
+    max_logging.log("=" * 80)
+    aot_time = pipeline.compile_aot_async(
+        params=params,
+        vae_params=vae_params,
+        qwen3_params=qwen3_params,
+        vae_bn_mean=vae_bn_mean,
+        vae_bn_std=vae_bn_std,
+        batch_size=config.batch_size,
+        height=config.height,
+        width=config.width,
+    )
+
+    max_logging.log("\n" + "=" * 80)
+    max_logging.log("🚀 Running initial dry run (Warmup Pass) to verify compiled graph execution...")
     max_logging.log("=" * 80)
-    pipeline.compile_aot_async(params, vae_params, qwen3_params, batch_size=config.batch_size, height=config.height, width=config.width)
     _, warmup_trace = pipeline(
         prompt=active_prompts,
         params=params,
@@ -500,7 +549,6 @@ def unbox_fn(x):
         batch_size=config.batch_size,
         use_latents=use_latents_flag,
         latents=latents_to_use,
-        warmup=True,
         output_dir=config.output_dir,
         output_name="flux2klein_warmup.png",
     )
@@ -536,15 +584,19 @@ def unbox_fn(x):
         main_trace.get("prompt_encoding", 0.0) + main_trace.get("denoise_loop", 0.0) + main_trace.get("vae_decode", 0.0)
     )
 
+    total_cold_start = load_time + aot_time + warmup_time
+
     max_logging.log("\n" + "=" * 80)
-    max_logging.log("📊 FLUX.2-KLEIN LATENCY & TIMING BREAKDOWN (PURE MODEL INFERENCE)")
+    max_logging.log("📊 FLUX.2-KLEIN COMPLETE LATENCY & TIMING BREAKDOWN")
     max_logging.log("=" * 80)
-    max_logging.log(f"1) Total Model Loading & Placement Time:  {load_time:.2f} seconds ⏱️")
-    max_logging.log(f"2) Cold-Start / Warmup Pass (XLA Compilation): {warmup_time:.2f} seconds ⏱️")
+    max_logging.log(f"1) Model Loading & Placement Time:              {load_time:.2f} seconds ⏱️")
+    max_logging.log(f"2) Concurrent AOT XLA Compilation Time:         {aot_time:.2f} seconds ⚡")
+    max_logging.log(f"3) Warmup Pass Execution Time:                   {warmup_time:.2f} seconds ⏱️")
     max_logging.log(f"   - Qwen3 Encoding:  {warmup_trace.get('prompt_encoding', 0.0):.2f}s")
     max_logging.log(f"   - Flux Denoising:  {warmup_trace.get('denoise_loop', 0.0):.2f}s")
     max_logging.log(f"   - VAE Decoding:    {warmup_trace.get('vae_decode', 0.0):.2f}s")
-    max_logging.log(f"3) Main Warmed-Up Pass (Pure Model Inference): {main_time:.2f} seconds ⏱️")
+    max_logging.log(f"👉 TOTAL COLD-START TIME (Loading + AOT + Warmup): {total_cold_start:.2f} seconds 🎯")
+    max_logging.log(f"4) Main Warmed-Up Pass (Pure Inference Latency): {main_time:.2f} seconds ⏱️")
     max_logging.log(f"   - Qwen3 Encoding:  {main_trace.get('prompt_encoding', 0.0):.2f}s")
     max_logging.log(f"   - Flux Denoising:  {main_trace.get('denoise_loop', 0.0):.2f}s")
     max_logging.log(f"   - VAE Decoding:    {main_trace.get('vae_decode', 0.0):.2f}s")

src/maxdiffusion/generate_wan.py

@@ -21,7 +21,7 @@
 from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2
 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p1 import WanCheckpointerI2V_2_1
 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p2 import WanCheckpointerI2V_2_2
-from maxdiffusion import aot_cache, pyconfig, max_logging, max_utils
+from maxdiffusion import pyconfig, max_logging, max_utils
 from absl import app
 from maxdiffusion.train_utils import transformer_engine_context
 from maxdiffusion.utils import export_to_video
@@ -279,30 +279,6 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
             dtype=config.weights_dtype,
         )
 
-  # Per-shape AOT executable cache: deserialization starts on background
-  # threads now and overlaps the remaining setup; unknown shapes silently
-  # fall back to jit and are serialized by save_pending() after warmup.
-  aot_cache.install(
-      getattr(config, "aot_cache_dir", ""),
-      meta={
-          "model": config.pretrained_model_name_or_path,
-          "attention": config.attention,
-          # Kernel block sizes change the lowered graph, not the input
-          # shapes — they must key the executable or a re-tuned config
-          # would silently hit stale binaries.
-          "flash_block_sizes": str(config.flash_block_sizes),
-          "mesh_shape": str(pipeline.mesh.shape),
-          "weights_dtype": str(config.weights_dtype),
-          "activations_dtype": str(config.activations_dtype),
-          "scan_layers": str(config.scan_layers),
-          "jax": jax.__version__,
-      },
-      mesh=pipeline.mesh,
-  )
-  # Deserialization is seconds and warmup must see the loaded executables
-  # to hit them; without this the first call races the loader threads.
-  aot_cache.wait_for_loads()
-
   s0 = time.perf_counter()
 
   # Disable profiler for the first two runs to avoid duplicate uploads
@@ -323,21 +299,11 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
   # step count only changes the Python loop trip count, not traced shapes.
   warmup_steps = min(2, config.num_inference_steps)
   max_logging.log(f"Compile warmup: {warmup_steps} denoising steps")
-  # Zero-execution warmup: wrapped transformer passes lower+compile (or
-  # reuse the deserialized AOT executable) and return sharded zeros, so
-  # the warmup pays compile time only, never real denoise compute. The
-  # returned videos are garbage by design and are discarded below.
-  with aot_cache.warmup_mode():
-    videos = call_pipeline(config, pipeline, prompt, negative_prompt, num_inference_steps=warmup_steps)
+  videos = call_pipeline(config, pipeline, prompt, negative_prompt, num_inference_steps=warmup_steps)
   if isinstance(videos, tuple):
     videos, warmup_trace = videos
     max_logging.log("Warmup breakdown: " + ", ".join(f"{stage}={seconds:.1f}s" for stage, seconds in warmup_trace.items()))
 
-  # Serialize any newly-compiled shapes synchronously while still inside
-  # warmup-accounted time; a background save would compete with the first
-  # real generation (DiffusionServing PR#39 first-generation-stall lesson).
-  aot_cache.save_pending()
-
   max_logging.log("===================== Model details =======================")
   max_logging.log(f"model name: {config.model_name}")
   max_logging.log(f"model path: {config.pretrained_model_name_or_path}")

src/maxdiffusion/models/flux/util.py

@@ -300,17 +300,17 @@ def unpack_latents(latents, batch_size, num_channels_latents, height, width):
   Unpacks packed sequence of shape (batch_size, (height//16)*(width//16), channels*4)
   back to the unpacked spatial grid shape (batch_size, channels, height//8, width//8).
   """
-  import numpy as np
+  import jax.numpy as jnp
 
   h_latent = height // 8
   w_latent = width // 8
 
   # 1. Reshape to split spatial grid and packed channel blocks
-  latents = np.reshape(latents, (batch_size, h_latent // 2, w_latent // 2, num_channels_latents, 2, 2))
+  latents = jnp.reshape(latents, (batch_size, h_latent // 2, w_latent // 2, num_channels_latents, 2, 2))
   # 2. Permute dimensions back to unpacked order
-  latents = np.transpose(latents, (0, 3, 1, 4, 2, 5))
+  latents = jnp.transpose(latents, (0, 3, 1, 4, 2, 5))
   # 3. Flatten back to 4D unpacked latent shape
-  latents = np.reshape(latents, (batch_size, num_channels_latents, h_latent, w_latent))
+  latents = jnp.reshape(latents, (batch_size, num_channels_latents, h_latent, w_latent))
   return latents
 
 
@@ -411,7 +411,7 @@ def cast_dict_to_bfloat16_inplace(d, device=None, exclude_keywords=None, parent_
 # -----------------------------------------------------------------------------
 
 
-def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, num_single_layers, dtype=None):
+def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, num_single_layers, dtype=None, pt_state_dict=None):
   """
   Loads weights from safetensors via zero-copy safetensors.numpy and converts them to JAX parameter dictionary.
   Supports dynamic layer counts (double and single stream blocks) and sharded safetensors directories.
@@ -423,16 +423,17 @@ def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_lay
   import os
   import gc
 
-  pt_state_dict = {}
-  if os.path.isdir(safetensors_path):
-    shards = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
-    max_logging.log(f"Loading sharded weights from directory: {safetensors_path} (Found {len(shards)} shards)...")
-    for shard in sorted(shards):
-      max_logging.log(f"Loading shard: {shard}...")
-      pt_state_dict.update(load_file(shard))
-  else:
-    max_logging.log(f"Loading weights from: {safetensors_path}")
-    pt_state_dict = load_file(safetensors_path)
+  if pt_state_dict is None:
+    pt_state_dict = {}
+    if os.path.isdir(safetensors_path):
+      shards = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
+      max_logging.log(f"Loading sharded weights from directory: {safetensors_path} (Found {len(shards)} shards)...")
+      for shard in sorted(shards):
+        max_logging.log(f"Loading shard: {shard}...")
+        pt_state_dict.update(load_file(shard))
+    else:
+      max_logging.log(f"Loading weights from: {safetensors_path}")
+      pt_state_dict = load_file(safetensors_path)
 
   max_logging.log("Mapping weights to JAX parameters...")
 
@@ -564,14 +565,15 @@ def convert_and_transpose_tensor(tensor, transpose=False, is_norm=False):
   return params
 
 
-def load_and_convert_vae_weights(safetensors_path, jax_params, dtype=None):
+def load_and_convert_vae_weights(safetensors_path, jax_params, dtype=None, pt_state_dict=None):
   """Loads VAE weights from safetensors via zero-copy safetensors.numpy, maps them to JAX, and extracts BN stats."""
   from safetensors.numpy import load_file
   import flax
   import jax.numpy as jnp
 
-  max_logging.log(f"Loading VAE weights from: {safetensors_path}")
-  pt_state_dict = load_file(safetensors_path)
+  if pt_state_dict is None:
+    max_logging.log(f"Loading VAE weights from: {safetensors_path}")
+    pt_state_dict = load_file(safetensors_path)
 
   # Unfreeze JAX params so we can load the weights
   jax_params = flax.core.unfreeze(jax_params)