more verbose timing

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -14,6 +14,7 @@
 
 import gc
 import os
+import time
 from typing import List
 
 from absl import app
@@ -300,6 +301,7 @@ def qwen3_init_fn():
 
   # 8. Load weights on Host CPU
   print("Loading parameters on Host CPU...")
+  t_load_start = time.time()
   cpu_device = jax.devices("cpu")[0]
   with jax.default_device(cpu_device):
     with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
@@ -378,6 +380,9 @@ def unbox_fn(x):
         gc.collect()
         jax.effects_barrier()
 
+  load_time = time.time() - t_load_start
+  print(f" -> [TIMING] Total Model Loading & Device Placement: {load_time:.2f} seconds ⏱️\n")
+
   # 9. Setup FlowMatch Scheduler
   scheduler = FlaxFlowMatchScheduler(
       num_train_timesteps=1000,
@@ -453,6 +458,7 @@ def unbox_fn(x):
     print("\n" + "=" * 80)
     print("🚀 Running initial dry run (Warmup Pass) to compile XLA graphs...")
     print("=" * 80)
+    t_warmup_start = time.time()
     pipeline(
         prompt=active_prompts,
         params=params,
@@ -472,10 +478,12 @@ def unbox_fn(x):
         output_dir=config.output_dir,
         output_name="flux2klein_warmup.png",
     )
+    warmup_time = time.time() - t_warmup_start
 
     print("\n" + "=" * 80)
     print("⏱️ Running timed pass at full TPU speed...")
     print("=" * 80)
+    t_main_start = time.time()
     pipeline(
         prompt=active_prompts,
         params=params,
@@ -495,6 +503,15 @@ def unbox_fn(x):
         output_dir=config.output_dir,
         output_name="flux2klein_generated_image.png",
     )
+    main_time = time.time() - t_main_start
+
+    print("\n" + "=" * 80)
+    print("📊 FLUX.2-KLEIN LATENCY & TIMING BREAKDOWN")
+    print("=" * 80)
+    print(f"1) Total Model Loading & Placement Time: {load_time:.2f} seconds ⏱️")
+    print(f"2) Cold-Start / Warmup Pass (XLA Compilation): {warmup_time:.2f} seconds ⏱️")
+    print(f"3) Main Warmed-Up Pass (Full-Speed Inference): {main_time:.2f} seconds ⏱️")
+    print("=" * 80)
 
     print("\n=======================================================")
     print(f"SUCCESS! Batched generation complete for {config.batch_size} images! 🎨🎉")