Standardize logging to max_logging.log and update AGENTS.md workflow

Author: amepas

.agents/AGENTS.md

@@ -0,0 +1,69 @@
+# Developer Guidelines & Workflow Rules — Flux.2-klein-4B
+
+## Project Goal & Context
+* **Objective**: Onboard the `Flux.2-klein-4B` model (currently implemented in PyTorch's `diffusers` library) into this `maxdiffusion` repository so it can run with JAX+TPU support.
+* **Reference Implementation**: We are porting and improving upon the PyTorch reference implementation: [pipeline_flux2_klein.py](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py).
+
+### Key Development Files
+* **Pipeline Entry Point**: `src/maxdiffusion/generate_flux2klein.py`. This is where we load configurations, load weights, stitch together all model blocks, run the image generation pipeline, apply VAE Batch Normalization, and run the JAX VAE Decoder.
+* **Component Testing**: `src/maxdiffusion/tests/generate_flux2klein_test.py`. This is where we write granular unit tests verifying the scheduler, RoPE position grids, packing/unpacking, individual attention blocks, E2E transformer, and the VAE decoder.
+
+### Completed Milestones (⚠️ LOST CHECKPOINT - ACTIVE RECOVERY IN PROGRESS)
+> [!IMPORTANT]
+> The parity milestones listed below were achieved on a previous repository checkpoint that is no longer available. We are currently working to recover the codebase to this state from scratch.
+1.  **4-Step Denoising Parity**: [In Progress] Recovering parity between the JAX transformer + FlowMatch scheduler loop and the PyTorch reference.
+2.  **VAE Decoder Port**: [In Progress] Verification of `FlaxAutoencoderKL` decoder mapping.
+3.  **VAE Batch Normalization**: [In Progress] Restoring channel-wise latent scaling.
+4.  **End-to-End Image Generation**: [In Progress] Achieving high SSIM/PSNR visual match against the PyTorch reference.
+5.  **Multi-Prompt Parity Validation (1024x1024)**: [In Progress] Sequential multi-prompt verification under bfloat16.
+
+---
+
+## 🚀 Concrete Next Action Items (The Roadmap)
+
+If you are picking up this project, your immediate priority tasks are:
+
+1.  **Task 1: Batched Generation ($B > 1$)**
+    *   Verify the pipeline under `batch_size > 1` (parallel prompt and image generation).
+    *   Ensure coordinate ID replication (`jnp.repeat`) and parallel latent unpacking loop work flawlessly without collapsing dimensions.
+    *   *Note*: Flux was **not** trained with attention masks (verified by Hugging Face `diffusers` codebase search), so do **not** attempt to implement an attention mask; let the model attend to the padded sequence naturally, as this is how it maintains its learned representations!
+2.  **Task 2: Dynamic Resolutions & Static Bucketing**
+    *   Verify JAX generation for non-square aspect ratios (e.g. 1024x768).
+    *   To prevent slow JAX/XLA re-compilation runs in production when the resolution changes, implement a **static shape bucketing strategy** (e.g. compiling for a few standard sizes like 512x512, 768x512, 1024x768, and padding inputs).
+3.  **Task 3: VAE Encoder Porting (Img2Img support)**
+    *   Port the VAE Encoder block (`FlaxEncoder` inside `vae_flax.py`) to JAX and write its weights mapping function.
+    *   Implement the inverse Batch Normalization step on encoded latents and verify the entire roundtrip (image $\rightarrow$ latents $\rightarrow$ image) is mathematically lossless.
+4.  **Task 4: ControlNet & LoRA Adapter Merging**
+    *   Verify how to dynamically load LoRA weights and merge them into the main parameter PyTrees (`vae_params` and `transformer_params`) efficiently without triggering XLA re-compilations.
+
+---
+
+## ⚠️ CRITICAL WORKFLOW RULES
+
+### 1. The Continuous Rsync Race Condition Warning
+*   **The Problem**: The user has a continuous `rsync` script running from their local workspace to the remote TPU VM (`local -> VM`).
+*   **The Risk**: If you generate any large file on the VM (like a new 1.8GB diagnostic `.npz` bundle, or a generated image), the local workspace won't have it. During the next `rsync` pass, the user's script will see the mismatch and **overwrite or delete** your newly generated file on the VM with the old/missing local one!
+*   **The Rule**: As soon as you generate a new diagnostic bundle or output on the VM, you **must immediately run a `gcloud compute tpus tpu-vm scp` command to download it to your local workspace**. Once it is local, the `rsync` script will see both environments are identical and won't touch it.
+
+### 2. Testing & Execution Environment
+*   **Target VM**: Run all code and tests exclusively on the remote TPU VM `ameypasarkar-tpu` in zone `southamerica-west1-a` (Project: `cloud-tpu-inference-test`).
+*   **No Local Runs**: Do not attempt to run the model or tests in your local workspace environment.
+*   **Permissions**: Always ask the user for explicit permission before running any tests or pipeline scripts on the TPU VM.
+*   **gcloud SSH Command Format**:
+    ```bash
+    gcloud compute tpus tpu-vm ssh ameypasarkar-tpu \
+      --zone southamerica-west1-a \
+      --project cloud-tpu-inference-test \
+      --command "source ~/.bashrc && source /mnt/data/mdiff_venv/bin/activate && cd /mnt/data/maxdiffusion/ && <command>"
+    ```
+    *(Note: Always prepend `source ~/.bashrc` to ensure user-specific environment variables like `HF_HOME` and Hugging Face authentication tokens are correctly loaded).*
+
+### 3. Git Workflow
+*   **Branch Constraint**: Work exclusively on the `flux2klein-onboarding` branch.
+*   **Commit History**: Create individual, focused commits per task/fix. Do NOT squash commits into 1 unless explicitly requested by the user.
+*   **Code Push**: Never push code to the remote repository without explicit user permission.
+
+### 4. Precision Policy: Debugging vs Production
+*   **For Parity Verification**: Set `jax_default_matmul_precision = "highest"` and load weights in `float32`. This removes TPU hardware-level `bfloat16` rounding noise, allowing you to isolate structural bugs by achieving clean $10^{-6}$ differences. (Note: Only feasible for resolutions up to 512x512 due to memory limits).
+*   **For Production/Inference & High Resolutions (e.g. 1024x1024)**: Set `weights_dtype: 'bfloat16'` and `activations_dtype: 'bfloat16'` in the config, and let JAX use default matmul precision. This activates the TPU's hardware matrix units, giving maximum speed and efficiency.
+*   **TPU HBM Memory Optimization**: Running 1024x1024 generation in float32 completely saturates the 32GB HBM of a single TPU v6e. You **must** run in `bfloat16` and use a recursive, leaf-by-leaf in-place parameter casting strategy (`cast_dict_to_bfloat16_inplace` with active Python `gc.collect()`) during weight conversion to avoid `RESOURCE_EXHAUSTED` OOM crashes.

src/maxdiffusion/generate_flux2klein.py

@@ -444,14 +444,14 @@ def unbox_fn(x):
     latents_to_use = None
     use_latents_flag = False
     if getattr(config, "latents_path", ""):
-      print(f"Loading custom starting noise latents from: {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
-      print(f" -> Custom latents shape: {latents_to_use.shape} | sum: {latents_to_use.sum():.6f}")
+      max_logging.log(f" -> Custom latents shape: {latents_to_use.shape} | sum: {latents_to_use.sum():.6f}")
 
-    print("\n" + "=" * 80)
-    print("🚀 Running initial dry run (Warmup Pass) to compile XLA graphs...")
-    print("=" * 80)
+    max_logging.log("\n" + "=" * 80)
+    max_logging.log("🚀 Running initial dry run (Warmup Pass) to compile XLA graphs...")
+    max_logging.log("=" * 80)
     _, warmup_trace = pipeline(
         prompt=active_prompts,
         params=params,
@@ -477,9 +477,9 @@ def unbox_fn(x):
         + warmup_trace.get("vae_decode", 0.0)
     )
 
-    print("\n" + "=" * 80)
-    print("⏱️ Running timed pass at full TPU speed...")
-    print("=" * 80)
+    max_logging.log("\n" + "=" * 80)
+    max_logging.log("⏱️ Running timed pass at full TPU speed...")
+    max_logging.log("=" * 80)
     _, main_trace = pipeline(
         prompt=active_prompts,
         params=params,
@@ -503,23 +503,23 @@ 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)
     )
 
-    print("\n" + "=" * 80)
-    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"2) Cold-Start / Warmup Pass (XLA Compilation): {warmup_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=======================================================")
-    print(f"SUCCESS! Batched generation complete for {config.batch_size} images! 🎨🎉")
-    print("=======================================================\n")
+    max_logging.log("\n" + "=" * 80)
+    max_logging.log("📊 FLUX.2-KLEIN LATENCY & TIMING BREAKDOWN (PURE MODEL INFERENCE)")
+    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"   - 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"   - 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")
+    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__":

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -24,6 +24,7 @@
 import numpy as np
 from flax.linen import partitioning as nn_partitioning
 
+from maxdiffusion import max_logging
 from ..pipeline_flax_utils import FlaxDiffusionPipeline
 from ...models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel
 from ...models.vae_flax import FlaxAutoencoderKL
@@ -163,13 +164,13 @@ def __call__(
       if latents_jax.ndim == 4:
         B, C, H, W = latents_jax.shape
         if C == 32:
-          print("  [PIPELINE] Unpacked 32-channel latents detected. Packing using pack_latents...")
+          max_logging.log("  [PIPELINE] Unpacked 32-channel latents detected. Packing using pack_latents...")
           latents_jax = pack_latents(latents_jax)
         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)
-      print(f"DEBUG PIPELINE: latents_numpy sum = {latents_numpy.sum():.6f} | mean = {latents_numpy.mean():.6f}")
+      max_logging.log(f"DEBUG PIPELINE: latents_numpy sum = {latents_numpy.sum():.6f} | mean = {latents_numpy.mean():.6f}")
       B, C, H, W = latents_numpy.shape
       latents_jax = jnp.transpose(jnp.reshape(latents_numpy, (B, C, H * W)), (0, 2, 1))
 
@@ -194,7 +195,7 @@ def __call__(
       # ---------------------------------------------------------------------
       # PHASE A: Encode Prompt (Qwen3)
       # ---------------------------------------------------------------------
-      print(f"[PHASE A] Encoding {len(prompts)} prompt(s) using JAX Qwen3 on TPU...")
+      max_logging.log(f"[PHASE A] Encoding {len(prompts)} prompt(s) using JAX Qwen3 on TPU...")
       t0 = time.perf_counter()
 
       # Resolve tokenizer path from config
@@ -215,11 +216,11 @@ def __call__(
 
       # 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]
-      print(f"DEBUG: templated_texts[0] = {repr(templated_texts[0])}")
+      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"])
       prompt_mask = jnp.array(inputs["attention_mask"])
-      print(f"DEBUG: prompt_ids[0][:15] = {prompt_ids[0][:15]}")
+      max_logging.log(f"DEBUG: prompt_ids[0][:15] = {prompt_ids[0][:15]}")
 
       # Run Text Encoding
       hidden_states, all_hidden_states = self._jitted_qwen3_forward(qwen3_params, prompt_ids, prompt_mask)
@@ -232,17 +233,19 @@ def __call__(
       # Transpose shape to [B, seq_len, 3*hidden_size]
       prompt_embeds_jax = jnp.transpose(out, (0, 2, 1, 3)).reshape((batch_size, seq_len_txt, -1))
       prompt_embeds_jax.block_until_ready()
-      print(
+      max_logging.log(
           f"DEBUG: prompt_embeds_jax min={float(prompt_embeds_jax.min()):.4f}, max={float(prompt_embeds_jax.max()):.4f}, mean={float(prompt_embeds_jax.mean()):.4f}, sum={float(prompt_embeds_jax.sum()):.4f}"
       )
 
       trace["prompt_encoding"] = time.perf_counter() - t0
-      print(f" -> [TIMING] Prompt Encoding (Qwen3): {trace['prompt_encoding']:.4f} seconds ⏱️")
+      max_logging.log(f" -> [TIMING] Prompt Encoding (Qwen3): {trace['prompt_encoding']:.4f} seconds ⏱️")
 
       # ---------------------------------------------------------------------
       # PHASE B: Denoising Loop (Flux Transformer)
       # ---------------------------------------------------------------------
-      print(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
@@ -267,17 +270,19 @@ def __call__(
 
         # Print progress
         sigma_val = scheduler_state.sigmas[step_idx]
-        print(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()
 
     trace["denoise_loop"] = time.perf_counter() - t0
-    print(f" -> [TIMING] Denoising Loop (Flux): {trace['denoise_loop']:.4f} seconds ⏱️")
+    max_logging.log(f" -> [TIMING] Denoising Loop (Flux): {trace['denoise_loop']:.4f} seconds ⏱️")
 
     # ---------------------------------------------------------------------
     # PHASE C: Decode Latents (VAE Decoder)
     # ---------------------------------------------------------------------
-    print("[PHASE C] Decoding final latents to RGB image using JAX VAE decoder on TPU...")
+    max_logging.log("[PHASE C] Decoding final latents to RGB image using JAX VAE decoder on TPU...")
     t0 = time.perf_counter()
 
     # Apply Channel-wise Batch Normalization Scaling in packed sequence format (denormalize)
@@ -295,12 +300,12 @@ def __call__(
     images_rgb.block_until_ready()
 
     trace["vae_decode"] = time.perf_counter() - t0
-    print(f" -> [TIMING] VAE Decoding: {trace['vae_decode']:.4f} seconds ⏱️")
+    max_logging.log(f" -> [TIMING] VAE Decoding: {trace['vae_decode']:.4f} seconds ⏱️")
 
     # ---------------------------------------------------------------------
     # POST-PROCESS: Format and Save Outputs
     # ---------------------------------------------------------------------
-    print("Postprocessing and saving generated images...")
+    max_logging.log("Postprocessing and saving generated images...")
     saved_paths = []
     # Clamp pixels and scale to [0, 255]
     images_rgb = jnp.clip((images_rgb + 1.0) / 2.0, 0.0, 1.0)
@@ -322,7 +327,7 @@ def __call__(
 
       output_png_path = os.path.join(output_dir, batch_output_name)
       img.save(output_png_path)
-      print(f" -> Saved image: {output_png_path} | Prompt: '{prompts[b_idx]}'")
+      max_logging.log(f" -> Saved image: {output_png_path} | Prompt: '{prompts[b_idx]}'")
       saved_paths.append(output_png_path)
 
     return saved_paths, trace