print statement fixes

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -20,6 +20,7 @@
 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
@@ -458,8 +459,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(
+    _, warmup_trace = pipeline(
         prompt=active_prompts,
         params=params,
         vae_params=vae_params,
@@ -478,13 +478,16 @@ def unbox_fn(x):
         output_dir=config.output_dir,
         output_name="flux2klein_warmup.png",
     )
-    warmup_time = time.time() - t_warmup_start
+    warmup_time = (
+        warmup_trace.get("prompt_encoding", 0.0)
+        + warmup_trace.get("denoise_loop", 0.0)
+        + warmup_trace.get("vae_decode", 0.0)
+    )
 
     print("\n" + "=" * 80)
     print("⏱️ Running timed pass at full TPU speed...")
     print("=" * 80)
-    t_main_start = time.time()
-    pipeline(
+    _, main_trace = pipeline(
         prompt=active_prompts,
         params=params,
         vae_params=vae_params,
@@ -503,14 +506,22 @@ def unbox_fn(x):
         output_dir=config.output_dir,
         output_name="flux2klein_generated_image.png",
     )
-    main_time = time.time() - t_main_start
+    main_time = (
+        main_trace.get("prompt_encoding", 0.0) + main_trace.get("denoise_loop", 0.0) + main_trace.get("vae_decode", 0.0)
+    )
 
     print("\n" + "=" * 80)
-    print("📊 FLUX.2-KLEIN LATENCY & TIMING BREAKDOWN")
+    print("📊 FLUX.2-KLEIN LATENCY & TIMING BREAKDOWN (PURE MODEL INFERENCE)")
     print("=" * 80)
-    print(f"1) Total Model Loading & Placement Time: {load_time:.2f} seconds ⏱️")
+    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(f"   - Qwen3 Encoding:  {warmup_trace.get('prompt_encoding', 0.0):.2f}s")
+    print(f"   - Flux Denoising:  {warmup_trace.get('denoise_loop', 0.0):.2f}s")
+    print(f"   - VAE Decoding:    {warmup_trace.get('vae_decode', 0.0):.2f}s")
+    print(f"3) Main Warmed-Up Pass (Pure Model Inference): {main_time:.2f} seconds ⏱️")
+    print(f"   - Qwen3 Encoding:  {main_trace.get('prompt_encoding', 0.0):.2f}s")
+    print(f"   - Flux Denoising:  {main_trace.get('denoise_loop', 0.0):.2f}s")
+    print(f"   - VAE Decoding:    {main_trace.get('vae_decode', 0.0):.2f}s")
     print("=" * 80)
 
     print("\n=======================================================")

src/maxdiffusion/models/flux/util.py

@@ -454,6 +454,7 @@ def load_and_convert_vae_weights(safetensors_path, jax_params):
   """Loads PyTorch VAE weights from safetensors, maps them to JAX, and extracts BN stats."""
   from safetensors.torch import load_file
   import torch
+  import numpy as np
   import flax
   import jax.numpy as jnp
 

src/maxdiffusion/models/qwen3_flax.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import math
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Union
 import flax.linen as nn
 import jax
 import jax.numpy as jnp

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -15,12 +15,13 @@
 import gc
 import os
 import time
-from typing import List, Union
+from typing import List, Union, Optional
 from PIL import Image
 
 import jax
 import jax.numpy as jnp
 import numpy as np
+import flax
 from flax.linen import partitioning as nn_partitioning
 
 from ..pipeline_flax_utils import FlaxDiffusionPipeline

src/maxdiffusion/tests/flux2klein/generate_flux2klein_test.py

@@ -81,12 +81,16 @@ def prepare_text_ids(batch_size, seq_len):
 class GenerateFlux2KleinTest(unittest.TestCase):
 
   def test_generate_random_latents_shape(self):
+    import torch
+
     latents = torch.randn((2, 32, 1024 // 8, 512 // 8))
     expected_shape = (2, 32, 1024 // 8, 512 // 8)
     self.assertEqual(tuple(latents.shape), expected_shape)
 
   def test_load_golden_latents_shape(self):
     # Deterministically generate initial latents on CPU using seed 0
+    import torch
+
     generator = torch.Generator(device="cpu").manual_seed(0)
     latents_pt = torch.randn((1, 32, 64, 64), generator=generator, dtype=torch.float32)
     expected_shape = (1, 32, 512 // 8, 512 // 8)
@@ -110,6 +114,7 @@ def test_qwen3_prompt_embeddings(self):
       self.fail(f"Failed to generate prompt embeddings: {e}")
 
   def test_context_embedder_projection(self):
+    import os
     import torch
     import jax
     import jax.numpy as jnp
@@ -329,11 +334,14 @@ def test_scheduler_timesteps_parity(self):
 
   def test_attention_blocks_parity(self):
     """Verifies that JAX joint-attention (double) and single-stream blocks match PyTorch golden outputs."""
+    import os
+    import torch
     import jax
     import jax.numpy as jnp
     import flax
     from flax.linen import partitioning as nn_partitioning
     from jax.sharding import Mesh
+    from safetensors.torch import load_file
 
     from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel
     from maxdiffusion import pyconfig
@@ -491,6 +499,7 @@ def test_attention_blocks_parity(self):
 
   def test_full_transformer_and_multistep_parity(self):
     """Verifies full JAX transformer forward pass (all blocks) and 4-step denoising loop parity against PyTorch."""
+    import os
     import torch
     import jax
 
@@ -704,6 +713,7 @@ def callback_fn(pipe, step_idx, timestep, callback_kwargs):
 
   def test_vae_decoder_parity(self):
     """Verifies JAX FlaxAutoencoderKL VAE Decoder parity against PyTorch."""
+    import os
     import jax
     import jax.numpy as jnp
     import numpy as np
@@ -876,7 +886,7 @@ def get_w(key):
 
     # 7. Compare raw decoder output
     diff_raw = jnp.abs(jax_decoder_out.sample - golden_decoder_out_pt)
-    print("\n[VAE DIAG] Raw Decoder Output Comparison:")
+    print(f"\n[VAE DIAG] Raw Decoder Output Comparison:")
     print(f"[VAE DIAG]   Max absolute diff: {jnp.max(diff_raw)}")
     print(f"[VAE DIAG]   Mean absolute diff: {jnp.mean(diff_raw)}")
 
@@ -904,6 +914,7 @@ def test_10_point_isolated_parity_benchmark(self):
 
   def test_swiglu_mlp_math_parity(self):
     """Verifies SwiGLU MLP mathematical parity between JAX and PyTorch."""
+    import torch
     import jax
     import jax.numpy as jnp
     import numpy as np
@@ -956,6 +967,7 @@ def test_attention_math_parity(self):
     pyconfig._config.keys["activations_dtype"] = "float32"
     config = pyconfig.config
 
+    import torch
     import jax
     import jax.numpy as jnp
     import numpy as np
@@ -1052,6 +1064,8 @@ def test_attention_math_parity(self):
 
   def test_flowmatch_scheduler_parity(self):
     """Verifies FlowMatch Euler Discrete Scheduler stepping parity against PyTorch."""
+    import torch
+    import jax
     import jax.numpy as jnp
     import numpy as np
     from diffusers import FlowMatchEulerDiscreteScheduler
@@ -1103,6 +1117,7 @@ def test_double_transformer_block_dummy_parity(self):
     pyconfig._config.keys["activations_dtype"] = "float32"
     config = pyconfig.config
 
+    import torch
     import jax
     import jax.numpy as jnp
     import numpy as np
@@ -1231,6 +1246,7 @@ def test_single_transformer_block_dummy_parity(self):
     pyconfig._config.keys["activations_dtype"] = "float32"
     config = pyconfig.config
 
+    import torch
     import jax
     import jax.numpy as jnp
     import numpy as np

src/maxdiffusion/tests/flux2klein/test_4b_e2e_parity.py

@@ -496,19 +496,19 @@ def jitted_vae_decode(v_params, latents_unpatched):
   l2_direct = np.sqrt(np.sum((jax_bf16_np - pt_bf16_np) ** 2))
   max_err_direct = np.max(np.abs(jax_bf16_np - pt_bf16_np))
 
-  print("\n--- Leg 2 (PyTorch CPU BF16) vs Leg 1 (PyTorch CPU FP32) Baseline ---")
+  print(f"\n--- Leg 2 (PyTorch CPU BF16) vs Leg 1 (PyTorch CPU FP32) Baseline ---")
   print(f"  SSIM:             {ssim_pt_bf16:.6f}")
   print(f"  RMSE:             {rmse_pt_bf16:.4f} / 255")
   print(f"  L2 Distance:      {l2_pt_bf16:.4f}")
   print(f"  Max Absolute Err: {max_err_pt_bf16:.4f} / 255")
 
-  print("\n--- Leg 3 (JAX TPU BF16) vs Leg 1 (PyTorch CPU FP32) Parity ---")
+  print(f"\n--- Leg 3 (JAX TPU BF16) vs Leg 1 (PyTorch CPU FP32) Parity ---")
   print(f"  SSIM:             {ssim_jax_bf16:.6f}  (Target: > 0.88)")
   print(f"  RMSE:             {rmse_jax_bf16:.4f} / 255")
   print(f"  L2 Distance:      {l2_jax_bf16:.4f}")
   print(f"  Max Absolute Err: {max_err_jax_bf16:.4f} / 255")
 
-  print("\n--- Leg 3 (JAX TPU BF16) vs Leg 2 (PyTorch CPU BF16) Direct Alignment ---")
+  print(f"\n--- Leg 3 (JAX TPU BF16) vs Leg 2 (PyTorch CPU BF16) Direct Alignment ---")
   print(f"  SSIM:             {ssim_direct:.6f}")
   print(f"  RMSE:             {rmse_direct:.4f} / 255")
   print(f"  L2 Distance:      {l2_direct:.4f}")

src/maxdiffusion/tests/flux2klein/test_9b_e2e_parity.py

@@ -506,19 +506,19 @@ def jitted_vae_decode(v_params, latents_unpatched):
   l2_direct = np.sqrt(np.sum((jax_bf16_np - pt_bf16_np) ** 2))
   max_err_direct = np.max(np.abs(jax_bf16_np - pt_bf16_np))
 
-  print("\n--- Leg 2 (PyTorch CPU BF16) vs Leg 1 (PyTorch CPU FP32) Baseline ---")
+  print(f"\n--- Leg 2 (PyTorch CPU BF16) vs Leg 1 (PyTorch CPU FP32) Baseline ---")
   print(f"  SSIM:             {ssim_pt_bf16:.6f}")
   print(f"  RMSE:             {rmse_pt_bf16:.4f} / 255")
   print(f"  L2 Distance:      {l2_pt_bf16:.4f}")
   print(f"  Max Absolute Err: {max_err_pt_bf16:.4f} / 255")
 
-  print("\n--- Leg 3 (JAX TPU BF16) vs Leg 1 (PyTorch CPU FP32) Parity ---")
+  print(f"\n--- Leg 3 (JAX TPU BF16) vs Leg 1 (PyTorch CPU FP32) Parity ---")
   print(f"  SSIM:             {ssim_jax_bf16:.6f}  (Target: > 0.88)")
   print(f"  RMSE:             {rmse_jax_bf16:.4f} / 255")
   print(f"  L2 Distance:      {l2_jax_bf16:.4f}")
   print(f"  Max Absolute Err: {max_err_jax_bf16:.4f} / 255")
 
-  print("\n--- Leg 3 (JAX TPU BF16) vs Leg 2 (PyTorch CPU BF16) Direct Alignment ---")
+  print(f"\n--- Leg 3 (JAX TPU BF16) vs Leg 2 (PyTorch CPU BF16) Direct Alignment ---")
   print(f"  SSIM:             {ssim_direct:.6f}")
   print(f"  RMSE:             {rmse_direct:.4f} / 255")
   print(f"  L2 Distance:      {l2_direct:.4f}")