AI Toolkit Adds Krea 2 Reference Image Training: Train Edit Concepts as Fast as Dedicated Edit Models
Ostris pushed support in AI Toolkit for training Krea 2 models with reference images, enabling edit-style LoRA training that masters concepts like 'make this a cyclops' in just 1,750 steps.
On July 4, 2026, Ostris pushed support in AI Toolkit for training Krea 2 with reference images. This lets you train edit-style LoRAs directly on Krea 2 Turbo using the training adapter — and in Ostris's testing, it learns edit concepts just as fast as a dedicated edit model would.
TL;DR — AI Toolkit now supports reference-image training for Krea 2 Turbo via a training adapter on HuggingFace. Edit-style LoRAs train in as few as 1,750 steps. A companion custom node for ComfyUI inference handles the reference image encoding and model patching. The same architecture works on the base Krea 2 model (Raw) too.
How It Works
Under the hood, the training pipeline uses the Qwen3-VL image encoder to encode reference images together with the prompt. The encoded reference is then fed into the Krea 2 transformer at timestep 0 alongside the clean training images — the same reference conditioning approach used by dedicated edit models.
| Component | Description |
|---|---|
| Training adapter | ostris/krea2_turbo_training_adapter on HuggingFace |
| Training engine | AI Toolkit — set arch: krea2 with model_kwargs.edit: true in config |
| Reference encoding | Qwen3-VL image encoder, downscaled to 384×384 total pixels |
| Steps to master | ~1,750 steps (tested: "make this a cyclops" concept) |
ComfyUI Inference: Custom Node Required
For inference in ComfyUI, you will need the ComfyUI-Krea2-Ostris-Edit custom node pack. It provides two nodes:
Text Encode Krea 2 Ostris Edit — Encodes the prompt together with up to 3 reference images through the Krea 2 Qwen3-VL text encoder, using the same Picture N: vision placeholder template used during training. When a VAE is connected, it also VAE-encodes the reference images and attaches them as reference latents.
Krea 2 Ostris Edit Model Patch — Patches the Krea 2 model so it consumes the reference latents from the conditioning. Each reference is appended to the image token sequence and conditioned at timestep 0 (the index_timestep_zero method). If no reference latents are present, the model behaves identically to stock Krea 2.
Installation
cd ComfyUI/custom_nodes
git clone https://github.com/ostris/ComfyUI-Krea2-Ostris-Edit.gitNo extra dependencies required. Nodes appear under the ostris/krea2 category.
Example Workflow
Load Diffusion Model (krea2) → Load LoRA → Krea 2 Ostris Edit Model Patch → KSampler
CLIPLoader (krea2) → Text Encode Krea 2 Ostris Edit (prompt + images + VAE) → positive
CLIPLoader (krea2) → Text Encode Krea 2 Ostris Edit (negative prompt) → negativeTraining Setup
To train an edit-style Krea 2 LoRA with AI Toolkit:
- Set
arch: krea2andmodel_kwargs.edit: truein your training config - Use paired training data: input images + edited target images + text descriptions
- The adapter is automatically downloaded from HuggingFace
- Training runs on consumer GPUs with standard VRAM requirements
Ostris demonstrated this with a "make this person a cyclops" test LoRA trained on Krea 2 Turbo — the model mastered the edit concept in 1,750 steps.
Base Model Compatibility
The same reference-image weights work on the base Krea 2 model (Raw) as well. As Ostris explained in response to community questions:
"Yes, the weights are the same. I am just running the reference images through the Qwen-VL image encoder and encoding them with the prompt, then feeding in the clean images with time-0 with the transformer."