Upload folder using huggingface_hub

0c17c2b verified about 2 months ago

7.02 kB

# 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`.

	# 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`.