Guia Rápido do Krea2¶

Este guia cobre treinamento LoRA do Krea2 no SimpleTuner. Krea2 é um transformer grande de imagem com flow matching, condicionamento de texto estilo Qwen e VAE do Qwen Image. Ele é mais confortável em GPUs NVIDIA com muita memória.

O exemplo inicial está em:

simpletuner/examples/krea2/raw/default/peft-lora/config.json

Ponto de Partida Recomendado¶

Para a primeira execução, use a configuração de exemplo e mantenha o modelo conservador:

{
  "model_family": "krea2",
  "model_flavour": "raw",
  "model_type": "lora",
  "pretrained_model_name_or_path": "krea/Krea-2-Raw",
  "mixed_precision": "bf16",
  "attention_mechanism": "cudnn",
  "gradient_checkpointing": true,
  "fuse_qkv_projections": true,
  "train_batch_size": 1,
  "base_model_precision": "no_change"
}

Krea2 é nativo em 1024px, mas 512px e 768px são úteis para iteração rápida e checagem de dataset. Use um dataloader de 1024px depois que a execução estiver estável.

Notas de Hardware¶

Krea2 pode treinar em bf16 em uma H100 de 80GB com batch 1 em 1024px. Batches maiores couberam sem compile nos nossos testes, mas compile adiciona memória de grafo/cudagraph suficiente para causar OOM em muitas configurações maiores.

TorchAO int8 weight-only reduz bastante a VRAM, mas não foi mais rápido que bf16 no caminho de treinamento testado. Use quando memória for mais importante que tempo por passo.

Recomendações:

• Use bf16 quando couber.
• Use int8-torchao quando precisar de folga de memória.
• Use attention_mechanism=cudnn em GPUs NVIDIA; o caminho SDPA nativo do Krea2 pode usar memória suficiente para causar OOM em runs de classe L40S que normalmente caberiam.
• Mantenha gradient_checkpointing=true.
• Mantenha fuse_qkv_projections=true.
• Use dynamo_backend=inductor, dynamo_mode=reduce-overhead e dynamo_use_regional_compilation=true apenas depois de confirmar que batch/resolução cabem.

Valores Principais de Configuração¶

{
  "model_family": "krea2",
  "model_flavour": "raw",
  "model_type": "lora",
  "pretrained_model_name_or_path": "krea/Krea-2-Raw",
  "base_model_precision": "no_change",
  "mixed_precision": "bf16",
  "attention_mechanism": "cudnn",
  "gradient_checkpointing": true,
  "fuse_qkv_projections": true,
  "optimizer": "optimi-lion",
  "learning_rate": 1e-4,
  "lora_rank": 64,
  "train_batch_size": 1,
  "resolution": 1024,
  "validation_resolution": "1024x1024"
}

Para TorchAO int8:

{
  "base_model_precision": "int8-torchao",
  "quantize_via": "cpu"
}

Para compile reduce-overhead:

{
  "dynamo_backend": "inductor",
  "dynamo_mode": "reduce-overhead",
  "dynamo_use_regional_compilation": true
}

Versão do LyCORIS¶

Use o LyCORIS 4.1.3 ou posterior para treinar LoKR. Essa versão lida com os fatores grandes que causavam o erro kron apply factors too large for the kernel nas versões 4.x anteriores.

Para atualizar uma instalação existente, execute no seu ambiente do SimpleTuner:

pip install --upgrade "lycoris-lora>=4.1.3"

Se você definiu LYCORIS_KERNEL_BACKEND=torch para contornar esse erro, remova essa configuração do ambiente de inicialização para usar a seleção automática de backend.

Versão 4.1.3 do LyCORIS.

Treinamento com Imagem de Referência¶

Krea2 suporta condicionamento opcional por latentes de referência para datasets de edição. Ative quando o dataloader fornecer imagens de referência pareadas ou latentes de referência em cache:

{
  "krea2_reference_latents": true
}

Os latentes de referência devem ter a mesma forma dos latentes alvo.

Configuração do Dataloader¶

Krea2 usa a estrutura geral de dataloader de imagem dos outros modelos transformer. A resolução real de treino vem do JSON do dataloader, não apenas de resolution no config principal. Para treinar em 1024px, confirme que resolution, maximum_image_size e target_downsample_size também são 1024 no dataloader.

Datasets em 512px são úteis para testes rápidos, checar captions e encontrar crops ruins. Para sinal de qualidade final, 1024px costuma ser mais representativo.

Para datasets locais, use type: local, defina instance_data_dir e escolha uma estratégia de caption. Para subject LoRA pequeno, caption_strategy=instanceprompt é um bom começo. Para estilos, filenames ou captions completas tendem a funcionar melhor.

Validação¶

Validação do Krea2 é cara, então comece com poucos prompts. Um único prompt pode esconder overfit ou memorização. Depois que o run estiver estável, adicione uma pequena prompt library.

{
  "validation_prompt": "a studio portrait of <token>, soft directional light, detailed fabric texture",
  "validation_negative_prompt": "ugly, cropped, blurry, low-quality, mediocre average",
  "validation_num_inference_steps": 28,
  "validation_guidance": 4.5,
  "validation_resolution": "1024x1024"
}

Notas de Quantização¶

int8-torchao armazena os pesos base do transformer em int8 e treina pesos LoRA bf16 por cima. Na H100 reduziu bastante a VRAM, mas foi mais lento que bf16 neste caminho de treinamento. Pense nele como uma opção de capacidade, não como garantia de throughput.

Resultados de Benchmark¶

As medições abaixo foram feitas em uma NVIDIA H100 80GB usando o trainer real do SimpleTuner, Krea2 LoRA, QKV fusionado, gradient checkpointing e um dataset pequeno Domokun. A VRAM foi amostrada externamente com nvidia-smi. Use estes valores apenas como comparação; versões diferentes de PyTorch, CUDA, driver, dataset, rank LoRA, otimizador, backend de atenção e GPU podem mudar os resultados.

QKV Fusionado + Checkpointing, Compile Desligado¶

QKV Fusionado + Checkpointing + Compile Reduce-Overhead¶

Orientação Prática¶

• Para iterar mais rápido em uma H100, use bf16, QKV fusionado, checkpointing, compile ligado e batch 1.
• Para batches efetivos maiores, prefira bf16 sem compile e aumente train_batch_size até a VRAM virar o limite.
• Para execuções com pouca memória, use int8-torchao; espere menos VRAM, mas passos mais lentos.
• Compile ajuda em batch 1, mas pode consumir VRAM suficiente para fazer batches maiores falharem.

Problemas Comuns¶

• Se você esperava 1024px mas o log mostra 512px, revise o JSON do dataloader.
• Se compile causa OOM mas o run sem compile cabe, reduza batch size ou desligue compile.
• Se int8 usa menos VRAM mas é mais lento, isso corresponde aos nossos testes H100.
• Se a imagem de referência não afeta a validação, confirme krea2_reference_latents=true e dataset pareado.
• Se overfitar rápido, reduza learning rate, reduza steps ou aumente a variedade do dataset.

Exportação para ComfyUI de LoKR de matriz completa com QKV fundido¶

Para um LoKR LyCORIS de matriz completa treinado com projeções QKV fundidas, exporte projeções separadas para o ComfyUI com o comando abaixo. Use o config.json do transformer do modelo base. Para DoRA normalizado pela saída, informe também o raw.safetensors original do Krea2 com --base-model. O conversor materializa as diferenças mescladas nativas do DoRA para projeções fundidas e separadas, evitando a normalização diferente do ComfyUI. Entradas Q/K/V separadas não utilizadas com delta zero ao lado de um adaptador fundido, incluindo suas escalas DoRA, são removidas; adaptadores conflitantes com delta diferente de zero são rejeitados. Os demais tensores são preservados. Omita --base-model para adaptadores sem DoRA. Ele materializa os deltas fundidos como matrizes FP32 para manter os cortes Q/K/V exatos mesmo quando cruzam os limites dos fatores de Kronecker; a exportação pode ser consideravelmente maior. Fatores LoKR decompostos e DoRA normalizado pela entrada não são suportados. Mantenha o checkpoint nativo para treinamento e retomada.

.venv/bin/python scripts/convert_krea2_lokr_to_comfyui.py \
  --input adapter.safetensors \
  --output adapter-comfyui.safetensors \
  --transformer-config transformer/config.json \
  --base-model raw.safetensors