Krea 2 Identity-Edit (Diffusers)

A self-contained 🤗 Diffusers pipeline that brings the conradlocke/krea2-identity-edit v1.2 identity-preserving instruction-edit recipe to Diffusers, running on top of the Krea-2-Turbo flow-matching model.

Given a source image and a natural-language edit instruction (e.g. "change the suit jacket to a red leather jacket"), the pipeline edits the image while holding the subject's identity / likeness. It also works as a plain Krea 2 text-to-image pipeline when no source image is supplied.

---

✨ Features

• Identity-preserving instruction editing — text-guided edits that keep the subject's likeness, mirroring the ComfyUI-Krea2Edit / conradlocke/krea2-identity-edit v1.2 recipe.
• Vendored Krea 2 transformer — a from-scratch Krea2Transformer2DModel (single-stream MMDiT) whose module tree and state-dict keys match the Krea 2 checkpoints, so the sharded weights load directly. The forward pass is extended with clean reference / source tokens that the stock upstream class does not support.
• Drop-in LoRA loading — loads Ostris AI-Toolkit / ComfyUI-format Krea 2 LoRAs (diffusion_model.blocks... with lora_down/lora_up) as well as already-converted diffusers-format state dicts. Alpha folding and key remapping are handled automatically.
• No custom diffusers build required — the pipeline registers the vendored transformer into the diffusers namespace at import time, so DiffusionPipeline.from_pretrained with trust_remote_code=True resolves it on stock diffusers releases.
• Turbo-fast — 10 denoising steps, guidance disabled (guidance_scale=0.0), thanks to the Krea-2-Turbo distilled checkpoint.

---

📋 Requirements

Install the pinned dependencies:

pip install -r requirements.txt

---

🚀 Quick start

python inference.py

---

🎛️ Tuning the edit

The two dials that matter most for identity edits:

Other useful arguments to pipe(...):

image = pipe(
    "change the suit jacket to a red leather jacket",
    image=Image.open("example.png"),
    num_inference_steps=10,   # Turbo recipe; more rarely helps
    guidance_scale=0.0,       # Krea convention: 0.0 disables CFG
    ref_boost=4.0,            # bump up for tougher identity holds
    grounding_px=768,         # 384–768
    height=1024,              # defaults to source AR, capped at ~1 MP,
    width=768,                # snapped to a multiple of 16
    generator=torch.Generator("cuda").manual_seed(0),  # for reproducibility
).images[0]

Output resolution for edits defaults to the source aspect ratio, capped at ~1 megapixel (MAX_EDIT_MEGAPIXELS = 1.0) and snapped to a multiple of vae_scale_factor * patch_size (= 16). The edit path doubles the transformer token count (it prepends the source latent), so the area cap keeps memory in check.

---

🔧 How it works

The identity-edit recipe (mirrors conradlocke/krea2-identity-edit v1.2 / ComfyUI-Krea2Edit) has two parts:

1. Image-grounded instruction

The edit instruction is grounded on the source image through the Qwen3-VL chat template. A <|vision_start|><|image_pad|><|vision_end|> block is inserted before the instruction so the VLM grounds the edit semantics on the source. The system prefix is byte-identical to the text-only Krea 2 template:

<|im_start|>system
Describe the image by detailing the color, shape, size, texture, quantity, text,
spatial relationships of the objects and background:<|im_end|>
<|im_start|>user
<|vision_start|><|image_pad|><|vision_end|>{instruction}<|im_end|>
<|im_start|>assistant

The Qwen3-VL text encoder produces a stack of hidden states tapped from several layers, which the transformer's text-fusion stage collapses into one text sequence.

2. Prepended source latent (Kontext-style)

The source image is VAE-encoded into a latent and prepended to the transformer sequence as a clean token block. The combined sequence is [text | source (frame=1) | target (frame=0)]. Source and target are distinguished purely by the 3-axis RoPE frame index (source frame = 1, target = 0) — not by timestep — so the target tokens can attend back to the clean source to preserve identity.

The ref_boost likeness dial then adds log(ref_boost) to the target-on-source attention logits, multiplicatively boosting how strongly the target copies the source before softmax renormalization. The float bias is applied chunk-wise by the memory-efficient SDPA backend, so the full seq × seq score matrix is never materialized.

See the module and method docstrings in krea2_edit_pipeline.py for the full detail (notably _grounded_encode, _encode_source_latent, _edit_position_ids, _ref_boost_bias, and _edit_transformer_forward).

---

🙏 Credits

• Model: krea/Krea-2-Turbo.

• Identity-edit LoRA: conradlocke/krea2-identity-edit(v1.2).

• Edit Pipeline: ostris/Krea2OstrisEdit and conradlocke/krea2-identity-edit

• Framework: 🤗 Diffusers.

---