# FLUX.2 Klein 9B KV Image Generation & Multi-Reference Edit Server

This directory contains the self-contained image generation and editing server integrating **`FLUX.2-klein-9b-kv-Nunchaku`** with this **`FLUX.2-klein-text-encoder-NVFP4`** text encoder.

It enables fast, 4-step distilled text-to-image generation and multi-reference image editing (with KV caching) on single consumer GPUs (e.g., 16 GB RTX 5060 Ti / RTX 4080 / RTX 5090) without running out of memory.

---

## 🧪 Quantization Details & Strategy

The text encoder was quantized using `llmcompressor` / `compressed-tensors` via [`quantize_flux2_klein_text_encoder_nvfp4.py`](file:///auto/home/amano/olegk/Nikola/src/nvfp4/quantize_flux2_klein_text_encoder_nvfp4.py) based on the recipe in [`recipe.yaml`](../recipe.yaml).

### High-Fidelity Representation Preservation Strategy

FLUX.2 Klein conditions its flow transformer on intermediate representations extracted from specific layers of the Qwen3 text encoder (specifically layer indices **9, 18, and 27**). Blanket quantization across all layers degrades prompt adherence and text grounding. To prevent this, the following selective quantization scheme was applied:

1. **Unquantized Self-Attention Projections (`q_proj`, `k_proj`, `v_proj`, `o_proj`)**:
   - Kept in full **bfloat16** across all 36 layers.
   - Preserves semantic query-key routing and relational representations without quantization noise.

2. **Unquantized Root & Near-Tap Layers**:
   - Layer 0 (root) and layers **8, 9, 17, 18, 26, 27** (the critical tap layers and their immediate predecessors) are preserved in full **bfloat16**.
   - Ensures the hidden state outputs at indices 9, 18, and 27 have near-perfect representation fidelity (cosine similarity > 0.999 vs unquantized reference).

3. **Compressed Intermediate MLP Blocks**:
   - The MLP blocks (`gate_proj`, `up_proj`, `down_proj`) of intermediate layers (1–7, 10–16, 19–25, 28–35) are compressed to **NVFP4 (W4A4)** using tensor-group scaling (`format: nvfp4-pack-quantized`).
   - Yields ~4x compression of MLP weights, reducing memory footprint and load times significantly.

4. **Preserved Head & Embeddings**:
   - `lm_head` and `embed_tokens` are kept unquantized in **bfloat16**.

---

## 📦 Required Packages in Python Environment

The server runs on Python 3.11 or 3.12 with CUDA 12.8 / 13.0 (e.g. Blackwell sm_120 / Thor sm_110):

| Package | Minimum Version | Purpose |
| :--- | :---: | :--- |
| `torch` | `2.11.0+` | PyTorch with Blackwell/sm_120 CUDA compute support |
| `diffusers` | `0.38.0+` | Built with native `Flux2KleinKVPipeline` support |
| `transformers` | `5.3.0+` | Qwen3 causal language model architecture support |
| `vllm` | `0.20.0+` | Provides precompiled CUTLASS NVFP4 GEMM kernels |
| `nunchaku` | `1.2.1+` | SVDQuant FP4/INT4 4-bit transformer execution |
| `fastapi`, `uvicorn` | Latest | OpenAI-compatible REST server infrastructure |
| `pillow` | Latest | Image processing and aspect-ratio padding |

> [!IMPORTANT]
> When installing packages in environments with custom CUDA builds, use `--no-deps` to avoid overwriting PyTorch or flash-attention wheels tuned for your compute capability.

---

## ⚡ Architecture & Sequential GPU Memory Management

On a 16 GB VRAM GPU, loading both the NVFP4 text encoder (~10.1 GB) and the Nunchaku transformer (~5.8 GB) simultaneously would total ~15.9 GB, exceeding available VRAM and causing an out-of-memory crash.

`Flux2KleinBackend.py` solves this using **Sequential GPU Stage Execution**:

```
[Idle: 0.01 GB VRAM]
       │
       ▼
Stage 1: Load Text Encoder to CUDA ──> Encode Prompt & IDs ──> Offload to CPU [0.01 GB VRAM]
       │
       ▼ (if reference images passed)
Stage 2: Load VAE to CUDA ───────────> Encode Image Latents ──> Offload to CPU [0.01 GB VRAM]
       │
       ▼
Stage 3: Load Transformer to CUDA ───> 4-Step Denoise (KV) ──> Offload to CPU [0.01 GB VRAM]
       │                                (Step 0: Extract KV)
       │                                (Steps 1-3: Cached KV)
       ▼
Stage 4: Load VAE to CUDA ───────────> Decode Latents to PIL ─> Offload to CPU [0.01 GB VRAM]
       │
       ▼
[Idle: 0.01 GB VRAM] (GPU completely freed for other tasks!)
```

### Performance & Memory Metrics (RTX 5060 Ti 16 GB)

- **Text-to-Image (1024x1024, 4 steps)**: ~9.5s total (~0.5s/step), **10.43 GB** peak VRAM.
- **1-Image Edit (1024x1024, 4 steps)**: ~12.5s total, **10.46 GB** peak VRAM.
- **2-Image Multi-Reference Edit (1024x1024, 4 steps)**: ~15.0s total, **13.18 GB** peak VRAM.
- **Idle Memory Footprint**: **0.010 GB** (10 MB).
- **Concurrency**: All requests are serialized with an `asyncio.Lock` queue, preventing concurrent memory spikes.

---

## 🚀 Running the Server

### 1. Launch via Helper Script

```bash
# Set paths (or edit defaults inside the script)
export MODEL_PATH="/path/to/FLUX.2-klein-9b-kv-Nunchaku"
export NVFP4_TEXT_ENCODER_PATH="/path/to/FLUX.2-klein-text-encoder-NVFP4"

./imagegen_flux2_klein.sh 4501
```

### 2. Launch Manually via Python

```bash
export VLLM_NVFP4_GEMM_BACKEND=cutlass

python ImageEditServer.py \
    --port 4501 \
    --model /path/to/FLUX.2-klein-9b-kv-Nunchaku \
    --optimized-model /path/to/FLUX.2-klein-9b-kv-Nunchaku/svdq-fp4_r32-FLUX.2-klein-9b-kv-Nunchaku.safetensors \
    --nvfp4-text-encoder /path/to/FLUX.2-klein-text-encoder-NVFP4 \
    --backend flux2-klein \
    --steps 4 \
    --guidance-scale 1.0
```

---

## 📡 API Usage Examples

### 1. Text-to-Image Generation (`POST /v1/images/generations`)

```bash
curl -X POST http://localhost:4501/v1/images/generations \
  -H "Content-Type: application/json" \
  -d '{
    "prompt": "A cute cheerful anime girl with a ponytail in a simple village dress smiling, anime style",
    "size": "1024x1024",
    "num_inference_steps": 4,
    "seed": 42
  }' | jq -r '.data[0].b64_json' | base64 -d > generated.png
```

### 2. Single-Image Reference Edit (`POST /v1/images/edits`)

```bash
curl -X POST http://localhost:4501/v1/images/edits \
  -F "image=@generated.png" \
  -F "prompt=The anime girl wearing stylish round glasses and a bright red ribbon in her hair, anime style" \
  -F "num_inference_steps=4" \
  -F "seed": 42 | jq -r '.data[0].b64_json' | base64 -d > edited_1img.png
```

### 3. Multi-Reference 2-Image Edit (`POST /v1/images/edits`)

Pass multiple `-F "image=@..."` form fields:

```bash
curl -X POST http://localhost:4501/v1/images/edits \
  -F "image=@girl.png" \
  -F "image=@sofa.png" \
  -F "prompt=The anime girl from image 1 is sitting comfortably on the sofa from image 2 in the living room, anime style" \
  -F "num_inference_steps=4" \
  -F "seed": 42 | jq -r '.data[0].b64_json' | base64 -d > edited_2imgs.png
```

### 4. Memory Introspection (`GET /v1/memory_stats`)

```bash
curl http://localhost:4501/v1/memory_stats | jq .
```
Expected output when idle: `"allocated_gb": 0.01`, `"reserved_gb": 0.02`.