flux2-lora-training

# FLUX.2 Klein 9B — LoRA Training Reference

## Архитектура моделей

### FLUX.2 Klein — Rectified Flow Transformer (DiT)

| Параметр | Klein 9B | Klein 4B | Qwen-Image-Edit | FLUX.1 dev |
|---------|---------|---------|---------|---------|
| Blocks | 32 (8+24) | 25 (5+20) | 60 (MM-DiT) | 56 (8+48) |
| Embedding dim | 12,288 | 7,680 | — | 15,360 |
| VAE latent channels | **128** | 128 | 16 (стандарт) | **16** |
| Text encoder | Qwen3 (bundled) | Qwen3 | Qwen2.5-VL (7B) | Mistral-Small-3.1 |
| Guidance embeddings | **НЕТ** | НЕТ | — | Есть |
| Total params | 9B | 4B | **20B DiT + 7B VL** | 12B |

**FLUX.2 Klein VAE:** `AutoencoderKLFlux2`, **32 latent channels** (FLUX.1: 16) → после 2×2 patch packing в трансформере: 32×4 = **128 dims per token**, 16× spatial compression. **Несовместим с FLUX.1 LoRA** — другой VAE, другой latent space. Tiling для больших разрешений: 1024px тайлы с 25% overlap, обрабатывает произвольное разрешение.

**Klein editing механизм (Kontext-style):** reference image VAE-кодируется и конкатенируется с noise latent вдоль sequence dim. Positional embeddings разделяют reference и output через **3D RoPE time offsets** (ref1=t:1, ref2=t:2, output=t:0). Поддерживает до 10 reference images теоретически, обучен на 2.

**Text encoder Klein:** Qwen3 (встроен в 9B), выходы из слоёв 9, 18, 27.

**Guidance embeddings в Klein отсутствуют** — `flux_guidance_mode`/`flux_guidance_value` — no-ops.

**Для LoRA тренировки: base модель** `FLUX.2-klein-base-9B`, не distilled 4-step.

---

### Qwen Image Edit 2511 — MMDiT с dual encoding

Архитектурно совершенно другая модель. **20B MMDiT** (Multimodal DiT) + **Qwen2.5-VL 7B** как visual-language encoder.

**Dual encoding** — ключевое отличие от Klein:
```
Входное изображение →┬→ Qwen2.5-VL 7B → семантические фичи (кто/что/стиль)
                     └→ VAE encoder   → reconstructive latent (текстуры/пиксели)
                                ↓ оба пути сходятся в MMDiT
```

Это даёт модели одновременно **понимание** содержимого (VL) и **воспроизводимость** деталей (VAE). Поэтому identity drift у Qwen Edit ниже, чем у Klein при тех же rank — модель буквально "видит" что на референсе через VLM.

**Consistency Mode** (`qe2511_consis_alpha`) — результат I2I reconstruction training objective: модель дообучена реконструировать входное изображение без изменений. Это выравнивает VL и DiT latent spaces, делая консервативные правки более точными. Триггер `restore image details` активирует этот bias.

**LoRA target modules для Qwen-Image-Edit-2511** (13 модулей, в отличие от Klein):
```
to_q, to_k, to_v, to_out.0
add_q_proj, add_k_proj, add_v_proj, to_add_out
img_mlp.net.2    ← image-specific FFN (в Klein нет разделения)
img_mod.1        ← image AdaLayerNorm modulation
txt_mlp.net.2    ← text-specific FFN
txt_mod.1        ← text modulation
linear_in, linear_out
```

**Qwen нативно поддерживает 3 control image слота** (control_1, control_2, control_3) через sequence concatenation + MSRoPE.

---

## Edit LoRA (before/after): датасет и тренировка

### Датасет (ai-toolkit формат)

```
dataset/
  control_1/    ← "before" изображение (источник)
  control_2/    ← опциональный 2-й референс (напр., лицо донора)
  targets/      ← "after" изображение (результат)
  captions/     ← .txt файл с инструкцией на сэмпл
```

**Критично:** имена файлов должны совпадать между папками (`0001.png ↔ 0001.png ↔ 0001.txt`).

**Известный баг в ai-toolkit** (issue #536, исправлен в PR #629): `folder_path` и `control_path_1` были перепутаны — модель учила трансформацию в обратную сторону. Обновись до актуальной версии.

**Порядок входных изображений важен.** Для BFS Head Swap:
- V1–V2: `[face, body]`
- **V3+: `[body, face]`** — инвертирование порядка дало значительный прирост качества

**Размер датасета:** для узких edit LoRA (head swap, relighting): **100–300 пар, качество > количество.** BFS эволюционировал 628 → 138 → 76 → 300+ высококачественных пар, подобранных по тону кожи.

### zero_cond_t — ключевой параметр

`zero_cond_t` трактует control images как чистые референсы при t=0, пока основное изображение следует нормальному diffusion schedule. Предотвращает identity drift.

**Должен быть включён и при тренировке, И при инференсе.** Для Qwen 2511 это особенно критично.

### Trigger word

Trigger ОБЯЗАТЕЛЬНО ставить в **user message** (slot `prompt`), НЕ в system/instruction message. LoRA-триггеры обучаются в user message контексте. Это объясняет почему в head swap воркфлоу `head_swap:` идёт в `prompt`, а не в `instruction`.

---

## Гиперпараметры и конфиги

### ai-toolkit (Klein 9B) — рекомендуется для edit LoRA

```yaml
learning_rate: 1e-4          # снизь до 5e-5 при нестабильности
batch_size: 1
gradient_accumulation: 1-4
lora_rank: 16-32             # стандарт; 64-128 для сложных edits
steps: 2000-6000             # зависит от датасета
save_every: 250-500
optimizer: adamw8bit
dtype: bf16
noise_scheduler: flowmatch
```

**BFS Head Swap для Klein 9B конкретно:**
- rank: 128 (начало), 64 (позже)
- steps: 3500–3750
- Curriculum resolution: начинал с 256/512px → прогрессивно до 1024/1536

### DiffSynth-Studio

```bash
python train.py \
  --learning_rate 1e-4 \
  --epochs 5 \
  --lora_rank 32 \
  --max_pixels 1048576  # 1024×1024
```

LoRA target modules (24 transformer blocks):
```
to_q, to_k, to_v, to_out.0,
add_q_proj, add_k_proj, add_v_proj, to_add_out,
linear_in, linear_out, to_qkv_mlp_proj
```

Полный fine-tune скрипт: `examples/flux2/model_training/full/FLUX.2-klein-base-9B.sh`

### SimpleTuner (Klein 9B)

```json
{
  "model_family": "flux2",
  "model_flavour": "klein-9b",
  "base_model_precision": "int8-quanto",
  "lora_rank": 16,
  "learning_rate": 1e-4,
  "lr_scheduler": "constant",
  "flux_guidance_mode": "constant",
  "flux_guidance_value": 1.0,
  "gradient_checkpointing": true
}
```

VRAM: ~14GB (int8-quanto), ~22GB (bf16).

### 50+ training runs (Calvin Herbst, Feb 2026)

Что реально работает:
- `linear: 128, linear_alpha: 64, conv: 64, conv_alpha: 32` → улучшает