Krea2OstrisEdit

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.

"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, ...):

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.