Standardize logging to max_logging, update reference images, and verify 4B/9B smoke tests

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -58,7 +58,9 @@ def partition_prompts(prompt_str: str, batch_size: int) -> List[str]:
       active.extend([raw_prompts[-1]] * (batch_size - len(active)))
     return active
   else:
-    print(f"⚠️ Warning: Found {num_prompts} prompts, but batch_size is {batch_size}. Truncating to the first {batch_size}.")
+    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]
 
 
@@ -405,7 +407,7 @@ def unbox_fn(x):
   )
 
   # 10. Instantiate and invoke FlaxFlux2KleinPipeline
-  print("Instantiating JAX FlaxFlux2KleinPipeline...")
+  max_logging.log("Instantiating JAX FlaxFlux2KleinPipeline...")
   pipeline = FlaxFlux2KleinPipeline(
       transformer=transformer,
       vae=vae,

src/maxdiffusion/models/attention_flax.py

@@ -1697,11 +1697,21 @@ def apply_rope(xq: Array, xk: Array, freqs_cis: Any) -> tuple[Array, Array]:
   if isinstance(freqs_cis, (tuple, list)):
     cos, sin = freqs_cis
     if cos.ndim == 2:
-      cos = cos[None, :, None, :]
-      sin = sin[None, :, None, :]
+      seq_len = cos.shape[0]
+      if xq.ndim == 4 and xq.shape[2] == seq_len:
+        cos = cos[None, None, :, :]
+        sin = sin[None, None, :, :]
+      else:
+        cos = cos[None, :, None, :]
+        sin = sin[None, :, None, :]
     elif cos.ndim == 3 and cos.shape[0] == 1:
-      cos = cos[:, :, None, :]
-      sin = sin[:, :, None, :]
+      seq_len = cos.shape[1]
+      if xq.ndim == 4 and xq.shape[2] == seq_len:
+        cos = cos[:, None, :, :]
+        sin = sin[:, None, :, :]
+      else:
+        cos = cos[:, :, None, :]
+        sin = sin[:, :, None, :]
 
     def _rotate(x):
       x_reshaped = x.reshape(*x.shape[:-1], -1, 2)

src/maxdiffusion/models/embeddings_flax.py

@@ -20,6 +20,7 @@
 from .modeling_flax_utils import get_activation
 from ..models.attention_flax import NNXSimpleFeedForward
 from ..models.normalization_flax import FP32LayerNorm
+from maxdiffusion import max_logging
 from maxdiffusion.tpu_utils import get_tpu_type, TpuType
 from maxdiffusion.max_utils import safe_getattr
 
@@ -296,7 +297,7 @@ def __call__(self, encoder_hidden_states_image: jax.Array) -> tuple[jax.Array, j
       # Apply pos_embed to the original sequence length
       hidden_states = hidden_states.at[:, :add_len, :].add(self.pos_embed.value[:, :add_len, :])
       if current_seq_len > pe_len:
-        print(f"[WARN] Input seq_len {current_seq_len} > pos_embed len {pe_len}")
+        max_logging.log(f"[WARN] Input seq_len {current_seq_len} > pos_embed len {pe_len}")
 
     hidden_states = self.norm1(hidden_states)
     hidden_states = self.ff(hidden_states)

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -110,7 +110,7 @@ def _prepare_latents(self, config, batch_size, height, width):
     seed_val = getattr(config, "seed", None)
     if seed_val is None:
       seed_val = int(time.time()) & 0x7FFFFFFF
-    print(
+    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)
@@ -215,7 +215,10 @@ def __call__(
         tokenizer = Qwen2TokenizerFast.from_pretrained(tokenizer_path, subfolder="tokenizer", local_files_only=True)
 
       # 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<think>\n\n</think>\n\n" for p in prompts]
+      templated_texts = [
+          f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
+          for p in prompts
+      ]
       max_logging.log(f"DEBUG: templated_texts[0] = {repr(templated_texts[0])}")
       inputs = tokenizer(templated_texts, return_tensors="np", padding="max_length", truncation=True, max_length=seq_len_txt)
       prompt_ids = jnp.array(inputs["input_ids"])
@@ -243,9 +246,7 @@ def __call__(
       # ---------------------------------------------------------------------
       # PHASE B: Denoising Loop (Flux Transformer)
       # ---------------------------------------------------------------------
-      max_logging.log(
-          f"[PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images..."
-      )
+      max_logging.log(f"[PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images...")
       t0 = time.perf_counter()
 
       guidance_vec_val = None
@@ -270,9 +271,7 @@ def __call__(
 
         # Print progress
         sigma_val = scheduler_state.sigmas[step_idx]
-        max_logging.log(
-            f" -> Step {step_idx}: Timestep = {scheduler_state.timesteps[step_idx]:.4f}, Sigma = {sigma_val:.4f}"
-        )
+        max_logging.log(f" -> Step {step_idx}: Timestep = {scheduler_state.timesteps[step_idx]:.4f}, Sigma = {sigma_val:.4f}")
 
       latents_jax.block_until_ready()
 

src/maxdiffusion/tests/generate_flux2klein_e2e_test.py

@@ -156,6 +156,7 @@ def main():
       "python3",
       "src/maxdiffusion/generate_flux2klein.py",
       "src/maxdiffusion/configs/base_flux2klein.yml",
+      "skip_jax_distributed_system=True",
       f"prompt={prompt}",
       "output_dir=/tmp/",
       f"seed={seed}",
@@ -175,14 +176,13 @@ def main():
   jax_4b_paths = [f"/tmp/jax_4b_dancing_b{b}.png" for b in range(batch_size)]
   print(f"[JAX 4B] Execution complete! Verifying outputs at: {jax_4b_paths}")
 
+  # =========================================================================
   # =========================================================================
   # PART II: FLUX.2-KLEIN-9B PARITY
   # =========================================================================
   print("\n" + "#" * 80)
   print("🎬 STARTING PARITY EVALUATION FOR FLUX.2-KLEIN-9B")
   print("#" * 80)
-
-  # 1. Run PyTorch 9B (FP32 & BF16)
   pt_9b_fp32_paths, pt_9b_bf16_paths = run_pytorch_pipeline(
       model_id="black-forest-labs/FLUX.2-klein-9B",
       prompt=prompt,
@@ -195,12 +195,12 @@ def main():
       prefix="9b",
   )
 
-  # 2. Run JAX 9B (via generate_flux2klein.py)
   print("\n[JAX 9B] Executing pipeline script generate_flux2klein.py...")
   cmd_jax_9b = [
       "python3",
       "src/maxdiffusion/generate_flux2klein.py",
       "src/maxdiffusion/configs/base_flux2klein_9B.yml",
+      "skip_jax_distributed_system=True",
       f"prompt={prompt}",
       "output_dir=/tmp/",
       f"seed={seed}",

src/maxdiffusion/tests/generate_flux2klein_smoke_test.py

@@ -118,3 +118,7 @@ def test_flux2klein_9b_smoke(self):
     ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
     print(f"\n[SMOKE TEST 9B] SSIM Score: {ssim_compare:.6f}")
     self.assertGreaterEqual(ssim_compare, 0.80)
+
+
+if __name__ == "__main__":
+  unittest.main()