anastasiyaw/codex-claude-code-config/flux2-lora-training

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)

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 воркфлоу head_swap: идёт в prompt, а не в instruction.

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

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

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

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)

{
  "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 → улучшает в почти каждом тесте
• lr: 0.000095, timestep_type: shift, weight_decay: 0.00015
• Более высокий rank = более чёткие детали, более естественный grain

Инпеинтинг LoRA

Подход 1: ostris "Green Screen" (проще всего)

Обучаешь модель на парах где маскированные области закрашены чистым зелёным (RGB 0, 255, 0). Модель учит: зелёный = "перегенерировать эту область по тексту".

• Использует существующий control_image slot в ai-toolkit — не нужна модификация архитектуры
• Inference: красишь область зелёным, пишешь промпт что должно быть
• Модель: ostris/qwen_image_edit_inpainting на HuggingFace
• При тренировке: normal t2i dropout для generalization, random inpaint blobs по умолчанию

Подход 2: SimpleTuner mask conditioning (PR #2520, merged Jan 29, 2026)

mask_conditioning_type — добавляет proper mask-conditioned training для FLUX.2 Klein. Подходит для более строгого контроля маски.

Чего НЕ делать

FLUX.1-Fill approach (concatenated latents, 16+16+64 channels) требует полной модификации архитектуры модели — не реализуемо как LoRA. Нужен full fine-tune с добавлением input channels.

Тренировка на больших разрешениях (1024 → 1536 → 2048)

Klein официально поддерживает 64×64 до 4MP (2048×2048), dimensions кратны 16.

DiffSynth-Studio: изменить --max_pixels

Нет архитектурных ограничений — Klein использует RoPE (resolution-agnostic).

Память на H200 (140GB)

H200 (140GB) комфортно берёт 1536 с оптимизациями.

Curriculum learning (метод BFS)

Начинай с маленького разрешения (256–512px), постепенно увеличивай. Прыжок сразу на 1536 может вызвать нестабильность. Kohya-ss задокументировали деградацию likeness при >1024px без curriculum подхода.

SimpleTuner bucketing

{
  "resolution_type": "pixel_area",
  "crop_aspect": "square",
  "crop_style": "center"
}

crop_aspect: square — наиболее надёжный вариант для multi-aspect.

Расширение latent space / checkpoint fine-tune

Нет опубликованных методов расширения native resolution Klein через checkpoint fine-tuning (на март 2026).

На практике это не нужно — Klein уже генерирует до 4MP нативно.

Если всё же нужен "latent expansion":

• Просто тренируй на большем --max_pixels через DiffSynth-Studio full fine-tune
• Теоретически применимы RoPE scaling техники (YaRN, NTK-aware) из LLM мира, но требуют данных на целевом разрешении

Full checkpoint fine-tune скрипт: examples/flux2/model_training/full/FLUX.2-klein-base-9B.sh (DiffSynth-Studio)

Сравнение фреймворков

Рекомендация для edit LoRA (head swap, face swap): ai-toolkit + base модель.

Тайловая обработка крупных изображений (tile-as-video)

Проблема

Банальная тайловая обработка через ComfyUI-TiledDiffusion даёт цветовые скачки между тайлами: каждый тайл генерируется независимо без знания соседей.

Решение: Kontext-style sequence concatenation (без обучения)

Klein нативно поддерживает произвольное число reference images через sequence concatenation + 3D RoPE. Это можно использовать для тайлов:

Tile 1 (обработан) → reference_1
Tile 2 (обработан) → reference_2
Tile 3 (текущий)   → output (Klein видит соседей при генерации)

Практический порядок обхода (BFS-style):

1. Обработать угловой тайл (нет соседей → reference пустой)
2. Обработать тайлы по периметру (1 сосед как reference)
3. Обработать внутренние тайлы (2-3 соседних тайла как references)

Это работает за счёт того, что Klein использует те же RoPE time offsets что и для head swap: ref1=t:1, ref2=t:2, output=t:0. Тайлы по-сути становятся "видеокадрами" в пространстве последовательности.

Существующие подходы для tiled diffusion