AI Toolkit Adds Krea 2 Reference Image Training: Train Edit Concepts as Fast as Dedicated Edit Models

news

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.

ComponentDescription
Training adapterostris/krea2_turbo_training_adapter on HuggingFace
Training engineAI Toolkit — set arch: krea2 with model_kwargs.edit: true in config
Reference encodingQwen3-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.git

No 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) → negative

Training Setup

To train an edit-style Krea 2 LoRA with AI Toolkit:

  1. Set arch: krea2 and model_kwargs.edit: true in your training config
  2. Use paired training data: input images + edited target images + text descriptions
  3. The adapter is automatically downloaded from HuggingFace
  4. 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."

AI Toolkit Adds Krea 2 Reference Image Training: Train Edit Concepts as Fast as Dedicated Edit Models | ComfyUI Wiki