| name | flux2-lora-training |
| description | Comprehensive reference for training LoRAs on FLUX.2 Klein 9B and Qwen Image Edit 2511 models. Use this skill whenever the user asks about: training LoRAs for flux2/flux 2 klein/qwen-image-edit, before/after edit LoRAs (head swap, face swap, image editing), inpainting LoRAs, training at larger resolutions, latent space expansion, VAE fine-tuning, multi-reference training (2 input images → 1 output), dataset preparation for edit models, zero_cond_t, ai-toolkit/SimpleTuner/DiffSynth configs, BFS head swap LoRA methodology, Qwen Edit architecture, consistency mode, dual encoding, FuseAnyPart, ACE++, maximum training resolution, или любые вопросы об обучении диффузионных моделей. ВСЕГДА используй этот скилл. Do NOT use for writing FLUX.2 Klein generation/edit prompts at inference time (use flux2-klein-prompting), nor for general non-training diffusion architecture/inference/memory work (use diffusion-engineering); this skill is about LoRA/VAE training, not prompting or serving.
|
| user-invocable | true |
| model | sonnet |
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
learning_rate: 1e-4
batch_size: 1
gradient_accumulation: 1-4
lora_rank: 16-32
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
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
| Разрешение | max_pixels |
|---|
| 1024×1024 (1MP) | 1048576 |
| 1536×1536 | 2359296 |
| 2048×2048 (4MP) | 4194304 |
Нет архитектурных ограничений — Klein использует RoPE (resolution-agnostic).
Память на H200 (140GB)
| Разрешение | VRAM (bf16) | Оптимизации |
|---|
| 1024×1024 | ~22GB | базовые |
| 1536×1536 | 40–80GB+ | grad_checkpoint + fp8 + 8bit adam |
| 2048×2048 | 80GB+ | все + cache_latents |
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)
Сравнение фреймворков
| Фреймворк | Klein 9B | Edit LoRA | Inpaint | Multi-Aspect | Заметки |
|---|
| ai-toolkit | Да | Да (native) | Green screen | Ограничено | Лучший для edit LoRA |
| SimpleTuner | Да | Частично (PR#2520) | Mask conditioning | Да (с нюансами) | Лучшая документация |
| DiffSynth-Studio | Да | Да (ICEdit format) | Не задокументировано | Да (dynamic) | Поддерживает full fine-tune |
| fal.ai trainer | Да (hosted) | Да (Qwen 2511) | Нет | Нет | Cloud, проще всего |
| kohya-ss | Ограничено | Нет | Нет | Да | Нестабильно для Klein |
Рекомендация для 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):
- Обработать угловой тайл (нет соседей → reference пустой)
- Обработать тайлы по периметру (1 сосед как reference)
- Обработать внутренние тайлы (2-3 соседних тайла как references)
Это работает за счёт того, что Klein использует те же RoPE time offsets что и для head swap: ref1=t:1, ref2=t:2, output=t:0. Тайлы по-сути становятся "видеокадрами" в пространстве последовательности.
Существующие подходы для tiled diffusion
| Метод | Где | Принцип | Для FLUX/DiT |
|---|
| MultiDiffusion | ComfyUI-TiledDiffusion | Усреднение latents в overlapping зонах | Работает (базовый) |
| SyncDiffusion | отдельно | LPIPS loss для согласованности цветов | Нет реализации для Klein |
| SpotDiffusion | ComfyUI-TiledDiffusion | Shifting windows, снижает артефакты | Лучший из TiledDiffusion |
| DemoFusion | отдельно | Progressive upscaling + skip residuals | Нет для Klein |
| DyPE | отдельно | Training-free 16MP для FLUX | Только для FLUX.1 |
| Scale-DiT | arxiv:2501.12900 | LoRA для масштабирования DiT | Обучается |
| Kontext-concat | нативно в Klein | Соседние тайлы как references | Лучший |
Практический ComfyUI пайплайн
Большое изображение (напр. 3072×2048)
↓
VAEEncodeTiled (1024px тайлы, 25% overlap)
↓
Для каждого тайла:
[уже_обработанные_соседи] → LoadImage → reference_1, reference_2
[текущий_тайл] → noise latent
Klein KSampler с 2 references
↓
VAEDecodeTiled
↓
Stitch с Gaussian blending в зонах overlap
Ключевые ноды: ComfyUI-TiledDiffusion (SpotDiffusion режим), VAEEncodeTiled, VAEDecodeTiled.
Почему circular convolution НЕ работает для Klein
FLUX.1/Klein трансформер не содержит conv слоёв в основной части — только attention и MLP. Circular padding применимо только к VAE (там есть conv), но не к самому DiT. Поэтому бесшовных панорам через circular padding не получить — используй StitchDiffusion LoRA.
Когда нужен тайлинг vs когда нет
| Сценарий | Решение |
|---|
| Изображение ≤ 4MP | Klein напрямую (поддерживает до 2048px нативно) |
| Изображение > 4MP | VAEEncodeTiled + SpotDiffusion |
| Панорама (seamless) | StitchDiffusion LoRA + wrap padding в VAE |
| Batch крупных продуктовых кадров | Тайлинг с Kontext-соседями |
Мульти-референс тренировка (2 входа → 1 выход)
Нужна ли модификация архитектуры?
Нет. Qwen-Image-Edit-2509/2511 уже нативно поддерживает 2 control image слота. Для head swap (body + face → result) архитектурных изменений не требуется:
dataset/
control_1/ ← BODY (поза, одежда, окружение)
control_2/ ← FACE (лицо донора)
targets/ ← RESULT (голова пересажена)
captions/ ← "head_swap: ..."
Порядок слотов критичен и должен быть консистентным во всём датасете.
Расширение до 3 слотов (body + face + pose map)
Только добавление control_3/ папки + новый time offset в position embeddings (t=3). Веса трансформера не меняются. Это уже поддерживается в Qwen-Image-Edit нативно.
Строгий attribute masking (без attribute leakage)
Проблема 2-slot LoRA: модель иногда берёт волосы с body вместо face. Решения:
A) Текстовые якоря (текущий подход): явно в промпте "hair MUST match Picture 2", повторить 2-3 раза в разных формулировках.
B) FuseAnyPart-style adapter (arxiv:2410.22771): отдельный lightweight adapter принимает masked_face + masked_body → combined conditioning. Жёсткое маскирование на уровне латентов. Требует нового компонента поверх модели.
C) ACE++ (arxiv:2501.02487): Long-context Condition Unit — все conditioning images (reference + mask + output) конкатенируются в единый extended sequence. Two-stage training: сначала 0-ref, потом multi-ref задачи.
Что нужно изменить в датасете vs 1-input → 1-output
| Аспект | 1-input | 2-input (head swap) |
|---|
| Датасет | (source, target) пары | (body, face, result) триплеты |
| Синтетика | Можно генерировать | Нужны реальные результаты head swap как GT |
| Caption | Общая инструкция | Явные ссылки на "image 1" и "image 2" |
| Порядок | Не критичен | Строго: body=control_1, face=control_2 |
| zero_cond_t | Желательно | Обязательно (иначе references зашумляются) |
Известные проблемы
| Проблема | Причина | Решение |
|---|
| Модель учит трансформацию в обратную сторону | Баг в ai-toolkit PR#536 — перепутаны control/target | Обновить ai-toolkit |
| LyCORIS создаёт 0 модулей | Key mismatch с Klein architecture | Использовать стандартный LoRA format |
| SimpleTuner зависает на RTX 5090 | Issue #2477 | Fallback на H200 |
| Деградация likeness при >1024px | Прямой jump на высокое разрешение | Curriculum: 256→512→1024→1536 |
| identity drift при editing | zero_cond_t не включён | Включить и при тренировке, и при inference |
| Attribute leakage (волосы с body вместо face) | 2-slot LoRA не имеет жёсткого mask control | Текстовые якоря ИЛИ FuseAnyPart adapter |
Источники
