---
library_name: diffusers
base_model: krea/Krea-2-Turbo
tags:
  - text-to-image
  - image-to-image
  - custom-pipeline
  - krea2
  - lora
license: apache-2.0
---

# Krea2OstrisEdit

A self-contained [community pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/custom_pipeline_overview) for [Krea 2](https://huggingface.co/krea/Krea-2-Turbo) 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](https://github.com/ostris/ai-toolkit)'s Krea 2 reference-image trainer and how they run with the [ComfyUI-Krea2-Ostris-Edit](https://github.com/ostris/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 |
| :---: | :---: | :---: |
| ![reference](images/style_ref_clean.png) | ![output](images/out_style_yeti_clean.png) | ![no reference](images/out_yeti_no_ref.png) |

*"a white yeti with horns reading a book" with the [Style Reference LoRA](https://huggingface.co/ostris/krea2_turbo_style_reference) — the reference image drives the style.*

## Usage

```python
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
).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). |
| `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.

## 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](https://github.com/huggingface/diffusers)). The Krea 2 model weights are covered by the [Krea 2 Community License](https://huggingface.co/krea/Krea-2-Turbo/blob/main/LICENSE.pdf).