/flux2-lora-training

name: flux2-lora-training
description: Plan or review LoRA and edit-training work specifically for FLUX.2 Klein or Qwen-Image-Edit, including paired datasets, trainer-version contracts, and held-out fidelity checks. Do not use for generic Stable Diffusion/DiT training, prompt authoring, or model serving; route those tasks to their specialized skill.

FLUX.2 Klein 9B — LoRA Training Reference

Evidence boundary

This file is a route to evaluate a training plan, not a tested recipe. Model cards establish model identity and license; trainer repositories establish only the behavior of their pinned version. Numeric settings and BFS notes below are historical local experiment candidates unless an exact source is linked. Before training, record the model revision, trainer commit, dataset contract, hardware, and a small held-out acceptance set. Do not promote a historical score, memory estimate, slot order, or private result into a universal default.

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

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 | **32** (128 packed features) | 32 (128 packed features) | verify the exact revision | **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.

Do not assume a trainer’s `flux_guidance_*` fields are no-ops or active from their names. Verify their behavior against the pinned trainer release and the selected model before relying on them.

**Для 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

This is an architectural interpretation, not a measured cross-model guarantee. Assess identity preservation on the task’s held-out images rather than claiming that one model drifts less at an equal LoRA rank.

**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`).

Do not infer `folder_path`/`control_path_1` semantics from issue #536 or PR #629. The primary records do not establish a reversed-path bug: #536 was closed as not planned, while #629 fixes Flux2 training that ignored a *single* control image. Pin the trainer commit and run a small non-production fixture that proves which input becomes which control/target before a training job. Sources: https://github.com/ostris/ai-toolkit/issues/536 and https://github.com/ostris/ai-toolkit/pull/629

**Порядок входных изображений важен.** Для BFS Head Swap:

• V1–V2: `[face, body]`
• **V3+: `[body, face]`** — инвертирование порядка дало значительный прирост качества

**Historical BFS observation:** the recorded dataset counts are not a minimum or maximum for another edit. Start from a representative, consented dataset and use held-out fidelity/attribute checks to decide whether more diversity or cleaner pairs are needed.

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

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

Use it only when the pinned trainer and inference implementation document the field; then keep train/inference conditioning consistent and verify it on a held-out edit. It is not a universal Qwen-Image-Edit switch.

Trigger word

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