Krea2OstrisEdit

A self-contained community pipeline for Krea 2 that adds:

• Reference-image (edit) conditioning — pass 1–2 reference images and the model generates with them as context (style transfer, editing, subject reference, etc., depending on the LoRA you load). This matches how edit LoRAs are trained with AI Toolkit's Krea 2 reference-image trainer and how they run with the ComfyUI-Krea2-Ostris-Edit custom nodes.
• LoRA loading for AI Toolkit / ComfyUI-format Krea 2 LoRAs (diffusion_model.* keys, lora_A/lora_B or lora_down/lora_up + alpha) as well as diffusers-format state dicts.

Everything lives in a single pipeline.py, so it works on diffusers releases that don't ship Krea 2 yet. Without a reference image it is a plain Krea 2 text-to-image sampler.

Reference	Output	Same seed, no reference

"a white yeti with horns reading a book" with the Style Reference LoRA — the reference image drives the style.

Usage

import torch
from diffusers import DiffusionPipeline
from PIL import Image

pipe = DiffusionPipeline.from_pretrained(
    "krea/Krea-2-Turbo",
    custom_pipeline="ostris/Krea2OstrisEdit",
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()  # or pipe.to("cuda") with ~40+ GB of VRAM

# An AI-Toolkit Krea 2 LoRA, e.g. the style reference LoRA
pipe.load_lora_weights(
    "ostris/krea2_turbo_style_reference", weight_name="krea2_style_reference.safetensors"
)

image = pipe(
    "a white yeti with horns reading a book",
    image=Image.open("style_reference.png"),  # one reference image or a list of them
    # kv_cache=True,  # reference K/V computed once and reused every step; only for
    #                 # LoRAs trained with AI-Toolkit's kv_cache model kwarg
).images[0]
image.save("output.png")

Works the same with krea/Krea-2-Raw (the non-distilled base model); sampling defaults adapt automatically (see below).

Call arguments

Beyond the standard diffusers text-to-image arguments (prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, generator, ...):

Argument	Default	Description
image	None	Reference image(s): a PIL image, numpy array, [0,1] CHW tensor, or a list of them. References keep their own aspect ratio; output size is set by height/width independently.
reference_max_pixels	1024 * 1024	Pixel budget each reference is downscaled to fit (never upscaled) before VAE encoding.
vl_image_max_pixels	384 * 384	Pixel budget for the coarse Qwen3-VL view of each reference.
encode_reference_in_prompt	True	Also embed references into the text conditioning through the Qwen3-VL vision tower (matches AI-Toolkit edit training).
kv_cache	False	Cache the reference tokens' attention K/V: precomputed in a single t=0 pass and reused on every denoising step, so the references never ride along in the per-step sequence (faster, especially with CFG or many steps). The LoRA must be trained with AI Toolkit's kv_cache model kwarg for this to work properly; leave off for normally trained edit LoRAs.
max_sequence_length	512	Maximum prompt token length (truncation only; prompts are encoded at natural length, not padded).

Defaults for num_inference_steps / guidance_scale follow the loaded checkpoint: 8 / 0.0 for the distilled Turbo model, 28 / 4.5 for the base model. Guidance uses the Krea 2 convention cond + scale * (cond - uncond), enabled whenever scale > 0 (this equals standard CFG with scale 1 + scale).

How reference conditioning works

Reference images condition the model in two places:

1. Through the Qwen3-VL text encoder — each image is embedded in the user message ahead of the prompt via Picture N: <|vision_start|><|image_pad|><|vision_end|> placeholders, so the text embeddings "see" the references.
2. As clean VAE latents appended after the noisy image tokens in the transformer sequence. They keep flow time t=0 (they are never noised) and sit on rotary-position frame axis i + 1 — the Kontext-style "index" placement.

With a LoRA trained using AI Toolkit's kv_cache option, the reference tokens attend only to each other, which makes their per-block attention K/V independent of the timestep and of everything else in the sequence. Passing kv_cache=True then computes those K/V once in a reference-only precompute pass and injects them as extra attention keys on every denoising step, instead of recomputing the full reference tokens each step (OminiControl2-style conditioning feature reuse). The LoRA must be trained with kv_cache enabled for this to work properly.

LoRA support

pipe.load_lora_weights(...) accepts a hub repo id (+ weight_name), a local .safetensors file or directory, or a state dict, in any of these formats:

• AI Toolkit / reference-trainer keys: diffusion_model.blocks.N.attn.wq.lora_A.weight, ...
• ComfyUI-style lora_down.weight / lora_up.weight with optional .alpha tensors (folded into the effective scale)
• Already-converted diffusers keys: transformer.transformer_blocks.N.attn.to_q.lora_A.weight, ...

unload_lora_weights(), fuse_lora() / unfuse_lora(), set_adapters(), and per-call scaling via attention_kwargs={"scale": 0.8} are also available.

Hardware notes

• bf16 weights are ~24 GB (transformer) + ~8 GB (Qwen3-VL text encoder) + VAE, so use pipe.enable_model_cpu_offload() on cards with less than ~40 GB of VRAM. On a 32 GB RTX 5090 a 1024×1024 Turbo image takes ~40–50 s with offloading.
• The Qwen3-VL image processor (only needed when passing reference images) is lazily loaded from Qwen/Qwen3-VL-4B-Instruct.

License

The pipeline code is Apache-2.0 (portions of the transformer implementation adapted from huggingface/diffusers). The Krea 2 model weights are covered by the Krea 2 Community License.