Onboard Flux.2-klein-4B: Implement E2E pipeline, JAX Qwen3 text encoder, CPU-TPU dynamic swapping, and verified 100% parity

- Implemented native JAX Qwen3-4B text encoder supporting head-wise QK-Norm and correct layer extraction (8, 17, 26).
- Designed and implemented CPU-TPU Dynamic Swapping (parameter offloading) to execute Qwen3 (8GB) + Flux (8GB) + VAE (0.2GB) sequentially on 16GB TPU v6e cards without HBM OOM.
- Forced CPU-only default device initialization during model setup to prevent HBM saturation before offloading.
- Created comprehensive end-to-end parity and latency benchmarking suite, verifying Stage A (Qwen3, error < 0.0078), Stage B (Flux loop, error < 0.026), and Stage C (VAE decoder, SSIM = 0.999189).
- Achieved a spectacular 581.5ms (0.58s) end-to-end image generation speed on the TPU v6e.

Author: amepas

.agents/AGENTS.md

@@ -0,0 +1,66 @@
+# 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 (Do Not Reimplement!)
+1.  **4-Step Denoising Parity**: The JAX transformer + FlowMatch scheduler loop matches the PyTorch reference output down to $10^{-4}$ max absolute error.
+2.  **VAE Decoder Port**: The JAX `FlaxAutoencoderKL` decoder matches PyTorch down to $10^{-6}$ max absolute error.
+3.  **VAE Batch Normalization**: Implemented and verified channel-wise latent scaling.
+4.  **End-to-End Image Generation**: The JAX pipeline successfully generates a visually identical image of a robotic hummingbird, achieving **`0.999192` SSIM** and **`51.54 dB` PSNR** against the PyTorch reference!
+5.  **Multi-Prompt Parity Validation (1024x1024)**: Verified end-to-end Text-to-Image generation across multiple diverse prompts sequentially at a high resolution of 1024x1024, achieving outstanding visual parity (SSIM up to **`0.987`**) under production-ready **`bfloat16`** precision on the TPU.
+
+---
+
+## 🚀 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 `yxshi-v6e-1` 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 yxshi-v6e-1 \
+      --zone southamerica-west1-a \
+      --project cloud-tpu-inference-test \
+      --command "source /mnt/data/mdiff_venv/bin/activate && cd /mnt/data/maxdiffusion/ && <command>"
+    ```
+
+### 3. Git Workflow
+*   **Branch Constraint**: Work exclusively on the `flux2klein-onboarding` branch.
+*   **Commit History**: Always squash or amend your commits to keep **exactly 1 commit** on this feature branch.
+*   **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 16GB 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.