removing extra files

Author: amepas

README.md

@@ -677,45 +677,60 @@ We added ring attention support for Wan models. Below are the stats for one `720
   python src/maxdiffusion/generate_flux.py src/maxdiffusion/configs/base_flux_schnell.yml jax_cache_dir=/tmp/cache_dir run_name=flux_test output_dir=/tmp/ prompt="photograph of an electronics chip in the shape of a race car with trillium written on its side" per_device_batch_size=1 ici_data_parallelism=1 ici_fsdp_parallelism=-1 offload_encoders=False
   ```
 
-  ## Flux.2-klein-4B
+  ## Flux.2-klein (4B & 9B)
 
-  MaxDiffusion supports JAX+TPU inference for the **Flux.2-klein-4B** model. We provide both a standard one-shot generation script and an **interactive generation mode** that reuses the JAX/XLA compiled graphs for instant subsequent runs.
+  MaxDiffusion supports JAX+TPU inference for both **Flux.2-klein-4B** and **Flux.2-klein-9B** models. We provide both a standard one-shot generation script and an **interactive generation mode** that keeps parameters in TPU HBM and reuses compiled JAX/XLA execution graphs for instant subsequent image generation.
 
-  ### Running Inference
+  ### 1. Running Inference
 
-  To generate images, run the following command:
+  **One-Shot Mode (4B Model):**
+  ```bash
+  python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein.yml \
+    prompt="an anime corgi eating sushi in the mountains"
+  ```
 
+  **One-Shot Mode (9B Model):**
   ```bash
-  python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein.yml
+  python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein_9B.yml \
+    prompt="an anime corgi eating sushi in the mountains"
   ```
 
-  To enable **interactive mode** (reusing the compiled graphs to avoid JIT compilation latency on subsequent prompts):
+  > **Switching Between 4B & 9B**:
+  > You can switch between 4B and 9B by passing the appropriate YAML config (`base_flux2klein.yml` for 4B, `base_flux2klein_9B.yml` for 9B) or by overriding `flux_name` via CLI arguments:
+  > * For 4B: `flux_name="flux2-klein-4b" pretrained_model_name_or_path="black-forest-labs/FLUX.2-klein-4B"`
+  > * 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**:
   ```bash
   python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein.yml interactive=True
   ```
 
-  ### Running Tests & Verification
-
-  All tests and verification scripts are located in `src/maxdiffusion/tests/flux2klein/`.
+  ### 2. Running Tests & Verification
 
-  To run the 19-point mathematical parity verification suite against PyTorch CPU:
-
-  ```bash
-  python src/maxdiffusion/tests/flux2klein/verify_blockwise_parity.py
-  ```
+  MaxDiffusion includes a full suite of unit, parity, and end-to-end smoke tests for Flux.2-klein:
 
-  To run the pure TPU latency benchmark (permanent HBM residence):
+  * **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
+    ```
 
-  ```bash
-  python src/maxdiffusion/tests/flux2klein/benchmark_tpu_pure.py
-  ```
+  * **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
+    ```
 
-  To run all unit and integration tests:
+  * **PyTorch vs JAX 4B End-to-End Parity Test**:
+    ```bash
+    pytest src/maxdiffusion/tests/flux2klein/test_4b_e2e_parity.py -s -vv
+    ```
 
-  ```bash
-  pytest src/maxdiffusion/tests/flux2klein/
-  ```
+  * **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: