more clean ups

Author: amepas

README.md

@@ -701,36 +701,10 @@ We added ring attention support for Wan models. Below are the stats for one `720
   > * For 9B: `flux_name="flux2-klein-9b" pretrained_model_name_or_path="black-forest-labs/FLUX.2-klein-9B"`
 
   **Interactive Mode (Instant Subsequent Generation):**
-  To avoid JIT re-compilation latency when testing multiple prompts, enable `interactive=True`. The model weights remain loaded in TPU memory, and after the initial warmup pass, each new prompt generates in **~0.9s**:
+  To avoid JIT re-compilation latency when testing multiple prompts, enable `interactive=True`. The model weights remain loaded in TPU memory, and after the initial warmup pass, each new prompt generates in **<1s**:
   ```bash
   python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein.yml interactive=True
   ```
-
-  ### 2. Running Tests & Verification
-
-  MaxDiffusion includes a full suite of unit, parity, and end-to-end smoke tests for Flux.2-klein:
-
-  * **Modular Component Unit Tests (Self-Contained Dummy Weights)**:
-    Verifies SwiGLU MLP, Attention, FlowMatch Scheduler, Double/Single Transformer Blocks, and VAE Decoder:
-    ```bash
-    pytest src/maxdiffusion/tests/flux2klein/generate_flux2klein_test.py -s -vv
-    ```
-
-  * **End-to-End Smoke Test Suite (1024x1024 Reference Image Parity)**:
-    Runs full 4B and 9B text-to-image pipelines and asserts `SSIM >= 0.90` against TPU-generated reference images:
-    ```bash
-    pytest src/maxdiffusion/tests/generate_flux2klein_smoke_test.py -s -vv
-    ```
-
-  * **PyTorch vs JAX 4B End-to-End Parity Test**:
-    ```bash
-    pytest src/maxdiffusion/tests/flux2klein/test_4b_e2e_parity.py -s -vv
-    ```
-
-  * **PyTorch vs JAX 9B End-to-End Parity Test**:
-    ```bash
-    pytest src/maxdiffusion/tests/flux2klein/test_9b_e2e_parity.py -s -vv
-    ```
   ## Fused Attention for GPU:
   Fused Attention for GPU is supported via TransformerEngine. Installation instructions: