extras/README.md

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 based on the recipe in 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

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

# 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

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)

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)

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:

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)

curl http://localhost:4501/v1/memory_stats | jq .

Expected output when idle: "allocated_gb": 0.01, "reserved_gb": 0.02.